make tasks wrapped by grpc-context unwrappable (#6977)

richardstartin · web-flow · commit 860baa900bb3 · 2024-05-01T13:24:26.000+01:00
diff --git a/dd-java-agent/instrumentation/grpc-1.5/src/test/groovy/UnwrapGrpcContextForkedTest.groovy b/dd-java-agent/instrumentation/grpc-1.5/src/test/groovy/UnwrapGrpcContextForkedTest.groovy
@@ -0,0 +1,26 @@
+import datadog.trace.agent.test.AgentTestRunner
+import datadog.trace.bootstrap.instrumentation.api.TaskWrapper
+import io.grpc.Context
+
+class UnwrapGrpcContextForkedTest extends AgentTestRunner {
+  @Override
+  protected void configurePreAgent() {
+    super.configurePreAgent()
+    injectSysConfig("dd.profiling.enabled", "true")
+  }
+
+  def "test unwrap gRPC context wrappers"() {
+    // relies on TaskUnwrappingInstrumentation (java-concurrent) targeting grpc context
+    setup:
+    def runnable = {}
+    def callable = {"result"}
+    when:
+    def unwrappedRunnableType = TaskWrapper.getUnwrappedType(Context.ROOT.wrap(runnable))
+    then:
+    unwrappedRunnableType == runnable.getClass()
+    when:
+    def unwrappedCallableType = TaskWrapper.getUnwrappedType(Context.ROOT.wrap(callable))
+    then:
+    unwrappedCallableType == callable.getClass()
+  }
+}
diff --git a/dd-java-agent/instrumentation/java-concurrent/src/main/java/datadog/trace/instrumentation/java/concurrent/TaskUnwrappingInstrumentation.java b/dd-java-agent/instrumentation/java-concurrent/src/main/java/datadog/trace/instrumentation/java/concurrent/TaskUnwrappingInstrumentation.java
@@ -52,6 +52,10 @@ public boolean isEnabled() {
     "task",
     "io.grpc.netty.shaded.io.netty.util.concurrent.PromiseTask",
     "task",
+    "io.grpc.Context$1",
+    "val$r",
+    "io.grpc.Context$2",
+    "val$c",
     "akka.dispatch.TaskInvocation",
     "runnable",
     "scala.concurrent.impl.CallbackRunnable",
