|
16 | 16 |
|
17 | 17 | package io.opentelemetry.instrumentation.hypertrace.grpc.v1_5.server; |
18 | 18 |
|
| 19 | +import static io.opentelemetry.javaagent.tooling.ClassLoaderMatcher.hasClassesNamed; |
| 20 | +import static io.opentelemetry.javaagent.tooling.bytebuddy.matcher.AgentElementMatchers.safeHasSuperType; |
19 | 21 | import static java.util.Collections.singletonMap; |
20 | 22 | import static net.bytebuddy.matcher.ElementMatchers.isMethod; |
| 23 | +import static net.bytebuddy.matcher.ElementMatchers.isPublic; |
21 | 24 | import static net.bytebuddy.matcher.ElementMatchers.named; |
| 25 | +import static net.bytebuddy.matcher.ElementMatchers.takesArguments; |
22 | 26 |
|
23 | | -import io.grpc.ServerInterceptor; |
| 27 | +import io.grpc.ServerBuilder; |
| 28 | +import io.opentelemetry.javaagent.instrumentation.api.CallDepthThreadLocalMap; |
24 | 29 | import io.opentelemetry.javaagent.tooling.TypeInstrumentation; |
25 | | -import java.util.List; |
26 | 30 | import java.util.Map; |
27 | 31 | import net.bytebuddy.asm.Advice; |
28 | 32 | import net.bytebuddy.description.method.MethodDescription; |
|
31 | 35 |
|
32 | 36 | public class GrpcServerBodyInstrumentation implements TypeInstrumentation { |
33 | 37 |
|
| 38 | + @Override |
| 39 | + public ElementMatcher<ClassLoader> classLoaderOptimization() { |
| 40 | + return hasClassesNamed("io.grpc.ServerBuilder"); |
| 41 | + } |
| 42 | + |
34 | 43 | @Override |
35 | 44 | public ElementMatcher<TypeDescription> typeMatcher() { |
36 | | - return named("io.grpc.internal.AbstractServerImplBuilder"); |
| 45 | + return safeHasSuperType(named("io.grpc.ServerBuilder")); |
37 | 46 | } |
38 | 47 |
|
39 | 48 | @Override |
40 | 49 | public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() { |
41 | 50 | return singletonMap( |
42 | | - isMethod().and(named("build")), |
| 51 | + isMethod().and(isPublic()).and(named("build")).and(takesArguments(0)), |
43 | 52 | GrpcServerBodyInstrumentation.class.getName() + "$AddInterceptorAdvice"); |
44 | 53 | } |
45 | 54 |
|
46 | 55 | public static class AddInterceptorAdvice { |
47 | 56 | @Advice.OnMethodEnter(suppress = Throwable.class) |
48 | | - public static void addInterceptor( |
49 | | - @Advice.FieldValue("interceptors") List<ServerInterceptor> interceptors) { |
50 | | - boolean shouldRegister = true; |
51 | | - for (ServerInterceptor interceptor : interceptors) { |
52 | | - if (interceptor instanceof GrpcServerInterceptor) { |
53 | | - shouldRegister = false; |
54 | | - break; |
55 | | - } |
56 | | - } |
57 | | - if (shouldRegister) { |
58 | | - // Interceptors run in the reverse order in which they are added |
59 | | - // OTEL interceptor is last |
60 | | - interceptors.add(interceptors.size() - 1, new GrpcServerInterceptor()); |
| 57 | + public static void onEnter(@Advice.This ServerBuilder<?> serverBuilder) { |
| 58 | + int callDepth = CallDepthThreadLocalMap.incrementCallDepth(GrpcServerInterceptor.class); |
| 59 | + if (callDepth == 0) { |
| 60 | + serverBuilder.intercept(new GrpcServerInterceptor()); |
61 | 61 | } |
62 | 62 | } |
| 63 | + |
| 64 | + @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) |
| 65 | + public static void onExit() { |
| 66 | + CallDepthThreadLocalMap.decrementCallDepth(GrpcServerInterceptor.class); |
| 67 | + } |
63 | 68 | } |
64 | 69 | } |
0 commit comments