Add a test case for kotlin coroutines tracing with suspension

monosoul · monosoul · commit df6d7224ad00 · 2022-12-09T10:22:04.000+01:00
Signed-off-by: monosoul &lt;Kloz.Klaud@gmail.com&gt;
diff --git a/dd-java-agent/instrumentation/kotlin-coroutines/src/test/groovy/KotlinCoroutineInstrumentationTest.groovy b/dd-java-agent/instrumentation/kotlin-coroutines/src/test/groovy/KotlinCoroutineInstrumentationTest.groovy
@@ -1,5 +1,5 @@
-import datadog.trace.core.DDSpan
 import datadog.trace.agent.test.AgentTestRunner
+import datadog.trace.core.DDSpan
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.ThreadPoolDispatcherKt
 
@@ -100,6 +100,41 @@ class KotlinCoroutineInstrumentationTest extends AgentTestRunner {
     dispatcher << dispatchersToTest
   }
 
+  def "kotlin suspension should not mess up traces"() {
+    setup:
+    KotlinCoroutineTests kotlinTest = new KotlinCoroutineTests(
+      ThreadPoolDispatcherKt.newSingleThreadContext("Single-Thread")
+      )
+    int expectedNumberOfSpans = kotlinTest.tracedWithSuspendingCoroutines()
+
+    expect:
+    assertTraces(1) {
+      trace(expectedNumberOfSpans, true) {
+        span(4) {
+          operationName "trace.annotation"
+          parent()
+        }
+        def topLevelSpan = span(3)
+        span(3) {
+          operationName "top-level"
+          childOf span(4)
+        }
+        span(2) {
+          operationName "synchronous-child"
+          childOf topLevelSpan
+        }
+        span(1) {
+          operationName "second-span"
+          childOf topLevelSpan
+        }
+        span(0) {
+          operationName "first-span"
+          childOf topLevelSpan
+        }
+      }
+    }
+  }
+
   private static DDSpan findSpan(List<DDSpan> trace, String opName) {
     for (DDSpan span : trace) {
       if (span.getOperationName() == opName) {
diff --git a/dd-java-agent/instrumentation/kotlin-coroutines/src/test/kotlin/KotlinCoroutineTests.kt b/dd-java-agent/instrumentation/kotlin-coroutines/src/test/kotlin/KotlinCoroutineTests.kt
@@ -1,20 +1,28 @@
 import datadog.trace.api.Trace
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeScope
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer.get
+import datadog.trace.bootstrap.instrumentation.api.ScopeSource.INSTRUMENTATION
 import edu.umd.cs.findbugs.annotations.SuppressFBWarnings
 import kotlinx.coroutines.CompletableDeferred
 import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.CoroutineName
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.CoroutineStart
+import kotlinx.coroutines.Deferred
 import kotlinx.coroutines.async
 import kotlinx.coroutines.awaitAll
 import kotlinx.coroutines.channels.actor
 import kotlinx.coroutines.channels.consumeEach
 import kotlinx.coroutines.channels.produce
 import kotlinx.coroutines.channels.toChannel
+import kotlinx.coroutines.delay
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.runBlocking
 import kotlinx.coroutines.selects.select
+import kotlinx.coroutines.sync.Mutex
+import kotlinx.coroutines.sync.withLock
 import kotlinx.coroutines.withTimeout
 import kotlinx.coroutines.yield
 import java.util.concurrent.TimeUnit
@@ -142,6 +150,68 @@ class KotlinCoroutineTests(private val dispatcher: CoroutineDispatcher) {
     activeSpan().setSpanName(opName)
   }
 
+  fun childSpan(opName: String): AgentSpan = get().buildSpan(opName)
+    .withResourceName("coroutines-test-span")
+    .start()
+
+  /**
+   * --- First job starts -------------------------- First job completes ---
+   *
+   * -------------------- Second job starts ---------------------------- Second job completes ---
+   */
+  @Trace
+  fun tracedWithSuspendingCoroutines(): Int {
+    fun jobContext(jobName: String) = CoroutineName(jobName)
+
+    return runTest {
+      val jobs = mutableListOf<Deferred<Unit>>()
+
+      val beforeFirstJobStartedMutex = Mutex(locked = true)
+      val afterFirstJobStartedMutex = Mutex(locked = true)
+
+      val beforeFirstJobCompletedMutex = Mutex(locked = true)
+      val afterFirstJobCompletedMutex = Mutex(locked = true)
+
+      childSpan("top-level").activateAndUse {
+        childSpan("synchronous-child").activateAndUse {
+          delay(5)
+        }
+
+        // this coroutine starts before the second one starts and completes before the second one
+        async(jobContext("first")) {
+          beforeFirstJobStartedMutex.lock()
+          childSpan("first-span").activateAndUse {
+            afterFirstJobStartedMutex.unlock()
+            beforeFirstJobCompletedMutex.lock()
+          }
+          afterFirstJobCompletedMutex.unlock()
+        }.run(jobs::add)
+
+        // this coroutine starts after the first one and completes after the first one
+        async(jobContext("second")) {
+          afterFirstJobStartedMutex.withLock {
+            childSpan("second-span").activateAndUse {
+              beforeFirstJobCompletedMutex.unlock()
+              afterFirstJobCompletedMutex.lock()
+            }
+          }
+        }.run(jobs::add)
+      }
+      beforeFirstJobStartedMutex.unlock()
+
+      jobs.awaitAll()
+
+      5
+    }
+  }
+
+  private suspend fun AgentSpan.activateAndUse(block: suspend () -> Unit) {
+    get().activateSpan(this, INSTRUMENTATION).use {
+      block()
+      finish()
+    }
+  }
+
   private fun <T> runTest(asyncPropagation: Boolean = true, block: suspend CoroutineScope.() -> T): T {
     activeScope().setAsyncPropagation(asyncPropagation)
     return runBlocking(dispatcher, block = block)
