Add a naive Kotlin coroutines instrumentation

Signed-off-by: monosoul <Kloz.Klaud@gmail.com>

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

@@ -0,0 +1,51 @@
+package datadog.trace.core;
+
+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 NoOpContextElement implements ThreadContextElement<Object> {
+  static final Key<NoOpContextElement> KEY = new Key<NoOpContextElement>() {};
+
+  public NoOpContextElement() {}
+
+  @Override
+  public void restoreThreadContext(@NotNull CoroutineContext coroutineContext, Object oldState) {}
+
+  @Override
+  public Object updateThreadContext(@NotNull CoroutineContext coroutineContext) {
+    return null;
+  }
+
+  @Nullable
+  @Override
+  public <E extends Element> E get(@NotNull Key<E> key) {
+    return CoroutineContext.Element.DefaultImpls.get(this, key);
+  }
+
+  @NotNull
+  @Override
+  public CoroutineContext minusKey(@NotNull Key<?> key) {
+    return CoroutineContext.Element.DefaultImpls.minusKey(this, key);
+  }
+
+  @NotNull
+  @Override
+  public CoroutineContext plus(@NotNull 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;
+  }
+}

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

@@ -0,0 +1,36 @@
+package datadog.trace.core;
+
+import datadog.trace.bootstrap.instrumentation.api.AgentScopeManager;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
+import datadog.trace.core.scopemanager.ContinuableScopeManager;
+import datadog.trace.core.scopemanager.ScopeStackCoroutineContext;
+import kotlin.coroutines.CoroutineContext;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class ScopeStackCoroutineContextHelper {
+
+  private static final NoOpContextElement NO_OP_CONTEXT_ELEMENT = new NoOpContextElement();
+  private static final Logger logger =
+      LoggerFactory.getLogger(ScopeStackCoroutineContextHelper.class);
+
+  public static CoroutineContext addScopeStackContext(final CoroutineContext other) {
+    final AgentTracer.TracerAPI agentTracer = AgentTracer.get();
+    final AgentScopeManager agentScopeManager =
+        agentTracer instanceof CoreTracer ? ((CoreTracer) agentTracer).scopeManager : null;
+
+    if (agentScopeManager instanceof ContinuableScopeManager) {
+      return other.plus(
+          new ScopeStackCoroutineContext((ContinuableScopeManager) agentScopeManager));
+    }
+
+    logger.warn(
+        "Unexpected Tracer or Scope Manager implementation. Tracer[expected={}, got={}], ScopeManager[expected={}, got={}]",
+        CoreTracer.class.getName(),
+        agentTracer.getClass().getName(),
+        ContinuableScopeManager.class.getName(),
+        agentScopeManager != null ? agentScopeManager.getClass() : "null");
+
+    return NO_OP_CONTEXT_ELEMENT;
+  }
+}

dd-java-agent/instrumentation/kotlin-coroutines/src/main/java/datadog/trace/core/scopemanager/ScopeStackCoroutineContext.java

@@ -0,0 +1,82 @@
+package datadog.trace.core.scopemanager;
+
+import static datadog.trace.bootstrap.instrumentation.api.ScopeSource.INSTRUMENTATION;
+
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.core.scopemanager.ContinuableScopeManager.ScopeStack;
+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 ScopeStackCoroutineContext implements ThreadContextElement<ScopeStack> {
+
+  static final Key<ScopeStackCoroutineContext> KEY = new Key<ScopeStackCoroutineContext>() {};
+  private final ContinuableScopeManager scopeManager;
+  private final ScopeStack scopeStack;
+  @Nullable private final AgentSpan span;
+
+  public ScopeStackCoroutineContext(ContinuableScopeManager scopeManager) {
+    this.scopeManager = scopeManager;
+    this.span = scopeManager.activeSpan();
+    /*
+     * initial scope stack for the context element should be empty to prevent the spans created within the coroutine
+     * from having a wrong parent span
+     */
+    this.scopeStack = scopeManager.tlsScopeStack.initialValue();
+  }
+
+  @Override
+  public void restoreThreadContext(
+      @NotNull CoroutineContext coroutineContext, ScopeStack oldState) {
+    scopeManager.tlsScopeStack.set(oldState);
+  }
+
+  @Override
+  public ScopeStack updateThreadContext(@NotNull CoroutineContext coroutineContext) {
+    final ScopeStack oldScopeStack = scopeManager.tlsScopeStack.get();
+    scopeManager.tlsScopeStack.set(scopeStack);
+
+    if (scopeStack.depth() == 0 && span != null) {
+      /*
+       * This is necessary for the spans created within the coroutine to properly inherit the spans hierarchy.
+       * It's not necessary to close the scope created here as it will be destroyed along with the context element and
+       * the scope stack.
+       */
+      scopeManager.activate(span, INSTRUMENTATION);
+    }
+
+    return oldScopeStack;
+  }
+
+  @Nullable
+  @Override
+  public <E extends Element> E get(@NotNull Key<E> key) {
+    return CoroutineContext.Element.DefaultImpls.get(this, key);
+  }
+
+  @NotNull
+  @Override
+  public CoroutineContext minusKey(@NotNull Key<?> key) {
+    return CoroutineContext.Element.DefaultImpls.minusKey(this, key);
+  }
+
+  @NotNull
+  @Override
+  public CoroutineContext plus(@NotNull 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;
+  }
+}

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

@@ -0,0 +1,15 @@
+package datadog.trace.instrumentation.kotlin.coroutines;
+
+import datadog.trace.core.ScopeStackCoroutineContextHelper;
+import kotlin.coroutines.CoroutineContext;
+import net.bytebuddy.asm.Advice;
+
+public class CoroutineContextAdvice {
+  @Advice.OnMethodEnter
+  public static void enter(
+      @Advice.Argument(value = 1, readOnly = false) CoroutineContext coroutineContext) {
+    if (coroutineContext != null) {
+      coroutineContext = ScopeStackCoroutineContextHelper.addScopeStackContext(coroutineContext);
+    }
+  }
+}

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

@@ -0,0 +1,40 @@
+package datadog.trace.instrumentation.kotlin.coroutines;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static datadog.trace.util.Strings.getPackageName;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.core.CoreTracer;
+
+@AutoService(Instrumenter.class)
+public class KotlinCoroutinesInstrumentation extends Instrumenter.Tracing
+    implements Instrumenter.ForSingleType {
+
+  public KotlinCoroutinesInstrumentation() {
+    super("kotlin-coroutines");
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {
+      getPackageName(CoreTracer.class.getName()) + ".ScopeStackCoroutineContextHelper",
+    };
+  }
+
+  @Override
+  public String instrumentedType() {
+    return "kotlinx.coroutines.CoroutineContextKt";
+  }
+
+  @Override
+  public void adviceTransformations(AdviceTransformation transformation) {
+    transformation.applyAdvice(
+        isMethod()
+            .and(named("newCoroutineContext"))
+            .and(takesArgument(1, named("kotlin.coroutines.CoroutineContext"))),
+        packageName + ".CoroutineContextAdvice");
+  }
+}