From 71995957307175214e7dd938644cbbf341aacce8 Mon Sep 17 00:00:00 2001
From: Anuraag Agrawal <aanuraag@amazon.co.jp>
Date: Mon, 7 Dec 2020 16:24:10 +0900
Subject: [PATCH] Switch gRPC instrumentation to instrument public
 ServerBuilder class. (#1839)

* Switch gRPC instrumentation to instrument public ServerBuilder class.

* A bit more readability (hopefully)

* hack

* Remove library dedupe for now since we need a better story.
---
 .../javaagent/grpc-1.5-javaagent.gradle       |  8 +---
 .../GrpcServerBuilderInstrumentation.java     | 45 +++++++++++--------
 2 files changed, 27 insertions(+), 26 deletions(-)

diff --git a/instrumentation/grpc-1.5/javaagent/grpc-1.5-javaagent.gradle b/instrumentation/grpc-1.5/javaagent/grpc-1.5-javaagent.gradle
index 8cac1ce4ec56..603cb3f08159 100644
--- a/instrumentation/grpc-1.5/javaagent/grpc-1.5-javaagent.gradle
+++ b/instrumentation/grpc-1.5/javaagent/grpc-1.5-javaagent.gradle
@@ -20,7 +20,7 @@ muzzle {
   pass {
     group = "io.grpc"
     module = "grpc-core"
-    versions = "[1.5.0, 1.33.0)" // see https://github.com/open-telemetry/opentelemetry-java-instrumentation/issues/1453
+    versions = "[1.5.0,)"
   }
 }
 
@@ -35,11 +35,5 @@ dependencies {
   testLibrary group: 'io.grpc', name: 'grpc-protobuf', version: grpcVersion
   testLibrary group: 'io.grpc', name: 'grpc-stub', version: grpcVersion
 
-  // see https://github.com/open-telemetry/opentelemetry-java-instrumentation/issues/1453
-  latestDepTestLibrary group: 'io.grpc', name: 'grpc-core', version: '1.32.+'
-  latestDepTestLibrary group: 'io.grpc', name: 'grpc-netty', version: '1.32.+'
-  latestDepTestLibrary group: 'io.grpc', name: 'grpc-protobuf', version: '1.32.+'
-  latestDepTestLibrary group: 'io.grpc', name: 'grpc-stub', version: '1.32.+'
-
   testImplementation project(':instrumentation:grpc-1.5:testing')
 }
\ No newline at end of file
diff --git a/instrumentation/grpc-1.5/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/grpc/v1_5/GrpcServerBuilderInstrumentation.java b/instrumentation/grpc-1.5/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/grpc/v1_5/GrpcServerBuilderInstrumentation.java
index 21b7fe46af09..d5c21f08d662 100644
--- a/instrumentation/grpc-1.5/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/grpc/v1_5/GrpcServerBuilderInstrumentation.java
+++ b/instrumentation/grpc-1.5/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/grpc/v1_5/GrpcServerBuilderInstrumentation.java
@@ -5,14 +5,18 @@
 
 package io.opentelemetry.javaagent.instrumentation.grpc.v1_5;
 
-import static java.util.Collections.singletonMap;
+import static io.opentelemetry.javaagent.tooling.ClassLoaderMatcher.hasClassesNamed;
+import static io.opentelemetry.javaagent.tooling.bytebuddy.matcher.AgentElementMatchers.safeHasSuperType;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.isPublic;
 import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
 
-import io.grpc.ServerInterceptor;
+import io.grpc.ServerBuilder;
 import io.opentelemetry.instrumentation.grpc.v1_5.server.TracingServerInterceptor;
+import io.opentelemetry.javaagent.instrumentation.api.CallDepthThreadLocalMap;
 import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
-import java.util.List;
+import java.util.Collections;
 import java.util.Map;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.method.MethodDescription;
@@ -21,33 +25,36 @@
 
 public class GrpcServerBuilderInstrumentation implements TypeInstrumentation {
 
+  @Override
+  public ElementMatcher<ClassLoader> classLoaderOptimization() {
+    return hasClassesNamed("io.grpc.ServerBuilder");
+  }
+
   @Override
   public ElementMatcher<TypeDescription> typeMatcher() {
-    return named("io.grpc.internal.AbstractServerImplBuilder");
+    return safeHasSuperType(named("io.grpc.ServerBuilder"));
   }
 
   @Override
   public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
-    return singletonMap(
-        isMethod().and(named("build")),
-        GrpcServerBuilderInstrumentation.class.getName() + "$AddInterceptorAdvice");
+    return Collections.singletonMap(
+        isMethod().and(isPublic()).and(named("build")).and(takesArguments(0)),
+        GrpcServerBuilderInstrumentation.class.getName() + "$BuildAdvice");
   }
 
-  public static class AddInterceptorAdvice {
+  public static class BuildAdvice {
 
     @Advice.OnMethodEnter(suppress = Throwable.class)
-    public static void addInterceptor(
-        @Advice.FieldValue("interceptors") List<ServerInterceptor> interceptors) {
-      boolean shouldRegister = true;
-      for (ServerInterceptor interceptor : interceptors) {
-        if (interceptor instanceof TracingServerInterceptor) {
-          shouldRegister = false;
-          break;
-        }
-      }
-      if (shouldRegister) {
-        interceptors.add(TracingServerInterceptor.newInterceptor());
+    public static void onEnter(@Advice.This ServerBuilder<?> serverBuilder) {
+      int callDepth = CallDepthThreadLocalMap.incrementCallDepth(ServerBuilder.class);
+      if (callDepth == 0) {
+        serverBuilder.intercept(TracingServerInterceptor.newInterceptor());
       }
     }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void onExit(@Advice.This ServerBuilder<?> serverBuilder) {
+      CallDepthThreadLocalMap.decrementCallDepth(ServerBuilder.class);
+    }
   }
 }