Restore parent scope stack when completing ScopeCoroutines

Author: mcculls

dd-java-agent/instrumentation/kotlin-coroutines/coroutines-1.5/src/test/groovy/KotlinCoroutineInstrumentationTest.groovy

@@ -1,7 +1,6 @@
 import datadog.trace.core.DDSpan
 import datadog.trace.instrumentation.kotlin.coroutines.AbstractKotlinCoroutineInstrumentationTest
 import kotlinx.coroutines.CoroutineDispatcher
-import spock.lang.Ignore
 
 class KotlinCoroutineInstrumentationTest extends AbstractKotlinCoroutineInstrumentationTest<KotlinCoroutineTests> {
 
@@ -44,7 +43,6 @@ class KotlinCoroutineInstrumentationTest extends AbstractKotlinCoroutineInstrume
     [dispatcherName, dispatcher] << dispatchersToTest
   }
 
-  @Ignore("Not working: disconnected trace")
   def "kotlin trace consistent after flow"() {
     setup:
     KotlinCoroutineTests kotlinTest = new KotlinCoroutineTests(dispatcher)

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

@@ -15,6 +15,7 @@
 import java.util.Collections;
 import java.util.Map;
 import kotlinx.coroutines.AbstractCoroutine;
+import kotlinx.coroutines.Job;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
@@ -112,8 +113,18 @@ public static void onStartInvocation(@Advice.This final AbstractCoroutine<?> cor
   public static class JobSupportAfterCompletionInternalAdvice {
     @Advice.OnMethodEnter
     public static void onCompletionInternal(@Advice.This final AbstractCoroutine<?> coroutine) {
-      // close the scope if needed
-      closeScopeStateContext(coroutine);
+      Job parent = null;
+      try {
+        // extract parent when completing scope coroutines so we can restore their scope stack
+        if (coroutine instanceof kotlinx.coroutines.internal.ScopeCoroutine) {
+          parent =
+              ((kotlinx.coroutines.internal.ScopeCoroutine) coroutine)
+                  .getParent$kotlinx_coroutines_core();
+        }
+      } catch (Throwable ignore) {
+        // fallback to restoring the current active scope stack after completing the coroutine
+      }
+      closeScopeStateContext(coroutine, parent);
     }
   }
 }

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

@@ -36,11 +36,12 @@ public static void initializeScopeStateContext(final AbstractCoroutine<?> corout
     }
   }
 
-  public static void closeScopeStateContext(final AbstractCoroutine<?> coroutine) {
+  public static void closeScopeStateContext(
+      final AbstractCoroutine<?> coroutine, final Job parent) {
     final ScopeStateCoroutineContext scopeStackContext =
         getScopeStateContext(coroutine.getContext());
     if (scopeStackContext != null) {
-      scopeStackContext.maybeCloseScopeAndCancelContinuation(coroutine);
+      scopeStackContext.maybeCloseScopeAndCancelContinuation(coroutine, parent);
     }
   }
 }

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

@@ -52,11 +52,20 @@ public ScopeState updateThreadContext(@NotNull final CoroutineContext coroutineC
   }
 
   /** If there's a context item for the coroutine then try to close it */
-  public void maybeCloseScopeAndCancelContinuation(final Job coroutine) {
+  public void maybeCloseScopeAndCancelContinuation(final Job coroutine, final Job parent) {
     final ScopeStateCoroutineContextItem contextItem = contextItemPerCoroutine.get(coroutine);
     if (contextItem != null) {
-      final ScopeState currentThreadScopeState = AgentTracer.get().newScopeState();
-      currentThreadScopeState.fetchFromActive();
+      ScopeState currentThreadScopeState = null;
+      if (parent != null) {
+        final ScopeStateCoroutineContextItem parentItem = contextItemPerCoroutine.get(parent);
+        if (parentItem != null) {
+          currentThreadScopeState = parentItem.getScopeState();
+        }
+      }
+      if (currentThreadScopeState == null) {
+        currentThreadScopeState = AgentTracer.get().newScopeState();
+        currentThreadScopeState.fetchFromActive();
+      }
 
       contextItem.maybeCloseScopeAndCancelContinuation();
       contextItemPerCoroutine.remove(coroutine);
@@ -107,6 +116,10 @@ public ScopeStateCoroutineContextItem() {
       coroutineScopeState = AgentTracer.get().newScopeState();
     }
 
+    public ScopeState getScopeState() {
+      return coroutineScopeState;
+    }
+
     public void activate() {
       coroutineScopeState.activate();
 

dd-java-agent/instrumentation/kotlin-coroutines/src/testFixtures/groovy/datadog/trace/instrumentation/kotlin/coroutines/AbstractKotlinCoroutineInstrumentationTest.groovy

@@ -406,7 +406,6 @@ abstract class AbstractKotlinCoroutineInstrumentationTest<T extends CoreKotlinCo
     [dispatcherName, dispatcher] << dispatchersToTest
   }
 
-  @Ignore("Not working: disconnected trace")
   def "kotlin trace consistent after delay"() {
     setup:
     CoreKotlinCoroutineTests kotlinTest = getCoreKotlinCoroutineTestsInstance(dispatcher)