Use CapturedContextProviders API

Author: ygree

dd-java-agent/instrumentation/servicetalk/build.gradle

@@ -1,3 +1 @@
 apply from: "$rootDir/gradle/java.gradle"
-
-//TODO generate lock files

dd-java-agent/instrumentation/servicetalk/gradle.lockfile

@@ -23,7 +23,7 @@ public Reference[] additionalMuzzleReferences() {
     return new Reference[] {
       // This check prevents older instrumentation from being applied to ServiceTalk v0.42.56+
       new Reference.Builder("io.servicetalk.concurrent.api.DelegatingExecutor")
-          // This field existed was removed in v0.42.56
+          // Removed in v0.42.56
           .withField(new String[0], 0, "delegate", "Lio/servicetalk/concurrent/api/Executor;")
           .build(),
     };

dd-java-agent/instrumentation/servicetalk/servicetalk-0.42.0/src/main/java/datadog/trace/instrumentation/servicetalk0_42_0/ServiceTalkInstrumentation.java

@@ -0,0 +1,47 @@
+package datadog.trace.instrumentation.servicetalk0_42_56;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.InstrumenterModule;
+import io.servicetalk.concurrent.api.CapturedContextProvider;
+import java.util.ArrayList;
+import java.util.List;
+import net.bytebuddy.asm.Advice;
+
+@AutoService(InstrumenterModule.class)
+public class CapturedContextProvidersInstrumentation extends InstrumenterModule.Tracing
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
+
+  public CapturedContextProvidersInstrumentation() {
+    super("servicetalk", "servicetalk-concurrent");
+  }
+
+  @Override
+  public String instrumentedType() {
+    return "io.servicetalk.concurrent.api.CapturedContextProviders";
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(named("loadProviders"), getClass().getName() + "$LoadProvidersAdvice");
+  }
+
+  private static final class LoadProvidersAdvice {
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void exit(@Advice.Return(readOnly = false) java.util.List loadedProviders) {
+      List<CapturedContextProvider> providers = new ArrayList<>(loadedProviders.size() + 1);
+      providers.addAll(loadedProviders);
+      providers.add(new DatadogCapturedContextProvider());
+      loadedProviders = providers;
+    }
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {
+      packageName + ".DatadogCapturedContextProvider",
+    };
+  }
+}

dd-java-agent/instrumentation/servicetalk/servicetalk-0.42.56/src/main/java/datadog/trace/instrumentation/servicetalk0_42_56/CapturedContextProvidersInstrumentation.java

@@ -0,0 +1,46 @@
+package datadog.trace.instrumentation.servicetalk0_42_56;
+
+import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
+import io.servicetalk.concurrent.api.CapturedContext;
+import io.servicetalk.concurrent.api.CapturedContextProvider;
+import io.servicetalk.concurrent.api.Scope;
+import io.servicetalk.context.api.ContextMap;
+
+public final class DatadogCapturedContextProvider implements CapturedContextProvider {
+  @Override
+  public CapturedContext captureContext(CapturedContext underlying) {
+    return captureContextCopy(underlying);
+  }
+
+  @Override
+  public CapturedContext captureContextCopy(CapturedContext underlying) {
+    return new WithDatadogCapturedContext(AgentTracer.activeSpan(), underlying);
+  }
+
+  private static final class WithDatadogCapturedContext implements CapturedContext {
+    private final AgentSpan agentSpan;
+    private final CapturedContext underlying;
+
+    public WithDatadogCapturedContext(AgentSpan agentSpan, CapturedContext underlying) {
+      this.agentSpan = agentSpan;
+      this.underlying = underlying;
+    }
+
+    @Override
+    public ContextMap captured() {
+      return underlying.captured();
+    }
+
+    @Override
+    public Scope attachContext() {
+      Scope stScope = underlying.attachContext();
+      AgentScope ddScope = AgentTracer.activateSpan(agentSpan);
+      return () -> {
+        ddScope.close();
+        stScope.close();
+      };
+    }
+  }
+}