Kotlin coroutines instrumentation (#4405)

Kotlin coroutines instrumentation

dd-java-agent/instrumentation/kotlin-coroutines/build.gradle

@@ -1,12 +1,68 @@
+muzzle {
+  pass {
+    group = 'org.jetbrains.kotlin'
+    module = 'kotlin-stdlib'
+    versions = "[1.3.72,)"
+    extraDependency "org.jetbrains.kotlinx:kotlinx-coroutines-core:1.3.0"
+  }
+  pass {
+    group = 'org.jetbrains.kotlinx'
+    module = 'kotlinx-coroutines-core'
+    versions = "[1.3.0,1.3.8)"
+    extraDependency "org.jetbrains.kotlin:kotlin-stdlib:1.3.72"
+  }
+  // for 1.3.9 - 1.4.3 muzzle fails to choose the variant
+  pass {
+    group = 'org.jetbrains.kotlinx'
+    module = 'kotlinx-coroutines-core-jvm'
+    versions = "[1.5.0,)"
+    extraDependency "org.jetbrains.kotlin:kotlin-stdlib:1.3.72"
+  }
+}
+
 apply from: "$rootDir/gradle/java.gradle"
 apply from: "$rootDir/gradle/test-with-kotlin.gradle"
 
+// Having Groovy, Kotlin and Java in the same project is a bit problematic
+// this removes Kotlin from main source set to avoid compilation issues
+sourceSets {
+  main {
+    kotlin {
+      srcDirs = []
+    }
+    java {
+      srcDirs = ["src/main/java"]
+    }
+  }
+}
+
+// this creates Kotlin dirs to make JavaCompile tasks work
+def createKotlinDirs = tasks.register("createKotlinDirs") {
+  def dirsToCreate = ["classes/kotlin/main"]
+  doFirst {
+    dirsToCreate.forEach {
+      new File(project.buildDir, it).mkdirs()
+    }
+  }
+
+  outputs.dirs(
+    dirsToCreate.collect {
+      project.layout.buildDirectory.dir(it)
+    }
+    )
+}
+
+tasks.withType(JavaCompile).configureEach {
+  inputs.files(createKotlinDirs)
+}
+
 dependencies {
   implementation project(':dd-java-agent:instrumentation:java-concurrent')
-  testImplementation project(':dd-trace-api')
+  compileOnly deps.kotlin
+  compileOnly deps.coroutines
 
   testImplementation deps.kotlin
   testImplementation deps.coroutines
-
+  testImplementation project(':dd-trace-api')
   testImplementation project(':dd-java-agent:instrumentation:trace-annotation')
 }

dd-java-agent/instrumentation/kotlin-coroutines/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/AbstractCoroutineInstrumentation.java

@@ -0,0 +1,148 @@
+package datadog.trace.instrumentation.kotlin.coroutines;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.extendsClass;
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static datadog.trace.instrumentation.kotlin.coroutines.CoroutineContextHelper.getScopeStateContext;
+import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
+import static net.bytebuddy.matcher.ElementMatchers.isDeclaredBy;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.isOverriddenFrom;
+import static net.bytebuddy.matcher.ElementMatchers.returns;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+import static net.bytebuddy.matcher.ElementMatchers.takesNoArguments;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import kotlin.coroutines.CoroutineContext;
+import kotlinx.coroutines.AbstractCoroutine;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+@AutoService(Instrumenter.class)
+public class AbstractCoroutineInstrumentation extends Instrumenter.Tracing
+    implements Instrumenter.ForTypeHierarchy {
+
+  private static final String ABSTRACT_COROUTINE_CLASS_NAME =
+      "kotlinx.coroutines.AbstractCoroutine";
+
+  private static final String JOB_SUPPORT_CLASS_NAME = "kotlinx.coroutines.JobSupport";
+  private static final String COROUTINE_CONTEXT_CLASS_NAME = "kotlin.coroutines.CoroutineContext";
+
+  public AbstractCoroutineInstrumentation() {
+    super("kotlin_coroutine.experimental");
+  }
+
+  @Override
+  protected final boolean defaultEnabled() {
+    return false;
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {
+      packageName + ".ScopeStateCoroutineContext",
+      packageName + ".ScopeStateCoroutineContext$ContextElementKey",
+      packageName + ".CoroutineContextHelper",
+    };
+  }
+
+  @Override
+  public void adviceTransformations(AdviceTransformation transformation) {
+    transformation.applyAdvice(
+        isConstructor()
+            .and(isDeclaredBy(named(ABSTRACT_COROUTINE_CLASS_NAME)))
+            .and(takesArguments(2))
+            .and(takesArgument(0, named(COROUTINE_CONTEXT_CLASS_NAME)))
+            .and(takesArgument(1, named("boolean"))),
+        AbstractCoroutineInstrumentation.class.getName() + "$AbstractCoroutineConstructorAdvice");
+
+    transformation.applyAdvice(
+        isMethod()
+            .and(isOverriddenFrom(named(ABSTRACT_COROUTINE_CLASS_NAME)))
+            .and(named("onStart"))
+            .and(takesNoArguments())
+            .and(returns(void.class)),
+        AbstractCoroutineInstrumentation.class.getName() + "$AbstractCoroutineOnStartAdvice");
+
+    transformation.applyAdvice(
+        isMethod()
+            .and(isOverriddenFrom(named(JOB_SUPPORT_CLASS_NAME)))
+            .and(named("onCompletionInternal"))
+            .and(takesArguments(1))
+            .and(returns(void.class)),
+        AbstractCoroutineInstrumentation.class.getName()
+            + "$JobSupportAfterCompletionInternalAdvice");
+  }
+
+  @Override
+  public String hierarchyMarkerType() {
+    return ABSTRACT_COROUTINE_CLASS_NAME;
+  }
+
+  @Override
+  public ElementMatcher<TypeDescription> hierarchyMatcher() {
+    return extendsClass(named(hierarchyMarkerType()));
+  }
+
+  /**
+   * Guarantees every coroutine created has a new instance of ScopeStateCoroutineContext, so that it
+   * is never inherited from the parent context.
+   *
+   * @see ScopeStateCoroutineContext
+   * @see AbstractCoroutine#AbstractCoroutine(CoroutineContext, boolean)
+   */
+  public static class AbstractCoroutineConstructorAdvice {
+    @Advice.OnMethodEnter
+    public static void constructorInvocation(
+        @Advice.Argument(value = 0, readOnly = false) CoroutineContext parentContext,
+        @Advice.Argument(value = 1) final boolean active) {
+      final ScopeStateCoroutineContext scopeStackContext = new ScopeStateCoroutineContext();
+      parentContext = parentContext.plus(scopeStackContext);
+      if (active) {
+        // if this is not a lazy coroutine, inherit parent span from the coroutine constructor call
+        // site
+        scopeStackContext.maybeInitialize();
+      }
+    }
+  }
+
+  /**
+   * If/when coroutine is started lazily, initializes ScopeStateCoroutineContext element on
+   * coroutine start
+   *
+   * @see ScopeStateCoroutineContext
+   * @see AbstractCoroutine#onStart()
+   */
+  public static class AbstractCoroutineOnStartAdvice {
+    @Advice.OnMethodEnter
+    public static void onStartInvocation(@Advice.This final AbstractCoroutine<?> coroutine) {
+      final ScopeStateCoroutineContext scopeStackContext =
+          getScopeStateContext(coroutine.getContext());
+      if (scopeStackContext != null) {
+        // try to inherit parent span from the coroutine start call site
+        scopeStackContext.maybeInitialize();
+      }
+    }
+  }
+
+  /**
+   * Guarantees a ScopeStateCoroutineContext element is always closed when coroutine transitions
+   * into a terminal state.
+   *
+   * @see ScopeStateCoroutineContext
+   * @see AbstractCoroutine#onCompletionInternal(Object)
+   */
+  public static class JobSupportAfterCompletionInternalAdvice {
+    @Advice.OnMethodEnter
+    public static void onCompletionInternal(@Advice.This final AbstractCoroutine<?> coroutine) {
+      final ScopeStateCoroutineContext scopeStackContext =
+          getScopeStateContext(coroutine.getContext());
+      if (scopeStackContext != null) {
+        // close the scope if needed
+        scopeStackContext.maybeCloseScopeAndCancelContinuation();
+      }
+    }
+  }
+}

dd-java-agent/instrumentation/kotlin-coroutines/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/CoroutineContextHelper.java

@@ -0,0 +1,22 @@
+package datadog.trace.instrumentation.kotlin.coroutines;
+
+import kotlin.coroutines.CoroutineContext;
+import kotlinx.coroutines.Job;
+import org.jetbrains.annotations.Nullable;
+
+public class CoroutineContextHelper {
+  /*
+  IntelliJ shows a warning here for Job being out of bounds, but that's not true, the class compiles.
+   */
+
+  @Nullable
+  @SuppressWarnings("unchecked")
+  public static Job getJob(final CoroutineContext context) {
+    return context.get((CoroutineContext.Key<Job>) Job.Key);
+  }
+
+  @Nullable
+  public static ScopeStateCoroutineContext getScopeStateContext(final CoroutineContext context) {
+    return context.get(ScopeStateCoroutineContext.KEY);
+  }
+}

dd-java-agent/instrumentation/kotlin-coroutines/src/main/java/datadog/trace/instrumentation/kotlin/coroutines/ScopeStateCoroutineContext.java

@@ -0,0 +1,109 @@
+package datadog.trace.instrumentation.kotlin.coroutines;
+
+import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
+import datadog.trace.bootstrap.instrumentation.api.ScopeState;
+import kotlin.coroutines.CoroutineContext;
+import kotlin.jvm.functions.Function2;
+import kotlinx.coroutines.ThreadContextElement;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+public class ScopeStateCoroutineContext implements ThreadContextElement<ScopeState> {
+
+  public static final Key<ScopeStateCoroutineContext> KEY = new ContextElementKey();
+  private final ScopeState coroutineScopeState;
+  @Nullable private AgentScope.Continuation continuation;
+  @Nullable private AgentScope continuationScope;
+  private boolean isInitialized = false;
+
+  public ScopeStateCoroutineContext() {
+    coroutineScopeState = AgentTracer.get().newScopeState();
+  }
+
+  /**
+   * If there is an active scope at the time of invocation, and it is async propagated, then
+   * captures the scope's continuation
+   */
+  public void maybeInitialize() {
+    if (!isInitialized) {
+      final AgentScope activeScope = AgentTracer.get().activeScope();
+      if (activeScope != null && activeScope.isAsyncPropagating()) {
+        continuation = activeScope.captureConcurrent();
+      }
+      isInitialized = true;
+    }
+  }
+
+  @Override
+  public void restoreThreadContext(
+      @NotNull final CoroutineContext coroutineContext, final ScopeState oldState) {
+    oldState.activate();
+  }
+
+  @Override
+  public ScopeState updateThreadContext(@NotNull final CoroutineContext coroutineContext) {
+    final ScopeState oldScopeState = AgentTracer.get().newScopeState();
+    oldScopeState.fetchFromActive();
+
+    coroutineScopeState.activate();
+
+    if (continuation != null && continuationScope == null) {
+      continuationScope = continuation.activate();
+    }
+
+    return oldScopeState;
+  }
+
+  /**
+   * If the context element has a captured scope continuation and an active scope, then closes the
+   * scope and cancels the continuation.
+   */
+  public void maybeCloseScopeAndCancelContinuation() {
+    final ScopeState currentThreadScopeState = AgentTracer.get().newScopeState();
+    currentThreadScopeState.fetchFromActive();
+
+    coroutineScopeState.activate();
+
+    if (continuationScope != null) {
+      continuationScope.close();
+    }
+    if (continuation != null) {
+      continuation.cancel();
+    }
+
+    currentThreadScopeState.activate();
+  }
+
+  @Nullable
+  @Override
+  public <E extends Element> E get(@NotNull final Key<E> key) {
+    return CoroutineContext.Element.DefaultImpls.get(this, key);
+  }
+
+  @NotNull
+  @Override
+  public CoroutineContext minusKey(@NotNull final Key<?> key) {
+    return CoroutineContext.Element.DefaultImpls.minusKey(this, key);
+  }
+
+  @NotNull
+  @Override
+  public CoroutineContext plus(@NotNull final CoroutineContext coroutineContext) {
+    return CoroutineContext.DefaultImpls.plus(this, coroutineContext);
+  }
+
+  @Override
+  public <R> R fold(
+      R initial, @NotNull Function2<? super R, ? super Element, ? extends R> operation) {
+    return CoroutineContext.Element.DefaultImpls.fold(this, initial, operation);
+  }
+
+  @NotNull
+  @Override
+  public Key<?> getKey() {
+    return KEY;
+  }
+
+  static class ContextElementKey implements Key<ScopeStateCoroutineContext> {}
+}