fix(rt): allow root spans to inherit their parent from the context when not already set (#757)

Author: aajtodd

.changes/0b232d1e-06b5-4557-a5f5-90245238f2ab.json

@@ -0,0 +1,8 @@
+{
+    "id": "0b232d1e-06b5-4557-a5f5-90245238f2ab",
+    "type": "bugfix",
+    "description": "Allow root trace spans to inherit their parent from current context",
+    "issues": [
+        "awslabs/aws-sdk-kotlin#759"
+    ]
+}

runtime/tracing/tracing-core/build.gradle.kts

@@ -8,6 +8,7 @@ extra["displayName"] = "Smithy :: Kotlin :: Tracing :: Core"
 extra["moduleName"] = "aws.smithy.kotlin.runtime.tracing"
 
 val kotlinLoggingVersion: String by project
+val coroutinesVersion: String by project
 
 kotlin {
     sourceSets {
@@ -29,6 +30,12 @@ kotlin {
             }
         }
 
+        commonTest {
+            dependencies {
+                implementation("org.jetbrains.kotlinx:kotlinx-coroutines-test:$coroutinesVersion")
+            }
+        }
+
         all {
             languageSettings.optIn("aws.smithy.kotlin.runtime.util.InternalApi")
         }

runtime/tracing/tracing-core/common/src/aws/smithy/kotlin/runtime/tracing/CoroutineContextUtils.kt

@@ -45,18 +45,31 @@ public suspend inline fun <R> CoroutineContext.withChildTraceSpan(
     }
 }
 
+@InternalApi
+public class WrappedRootSpan(
+    private val rootSpan: TraceSpan,
+    override val parent: TraceSpan,
+) : TraceSpan by rootSpan
+
 @InternalApi
 public suspend inline fun <R> CoroutineContext.withRootTraceSpan(
     rootSpan: TraceSpan,
     crossinline block: suspend CoroutineScope.() -> R,
 ): R =
     try {
         val existingSpan = get(TraceSpanContextElement)?.traceSpan
-        check(existingSpan == null || existingSpan == rootSpan.parent) {
+        check(existingSpan == null || rootSpan.parent == null || existingSpan == rootSpan.parent) {
             "This method may only be called when no current span exists or the new span is a child of the active span"
         }
 
-        withContext(TraceSpanContextElement(rootSpan)) {
+        val newRoot = if (existingSpan != null && rootSpan.parent == null) {
+            // we are in some nested context, the existing span becomes this root's parent
+            WrappedRootSpan(rootSpan, existingSpan)
+        } else {
+            rootSpan
+        }
+
+        withContext(TraceSpanContextElement(newRoot)) {
             block()
         }
     } finally {

runtime/tracing/tracing-core/common/test/aws/smithy/kotlin/runtime/tracing/TraceSpanCoroutineUtilsTest.kt

@@ -0,0 +1,81 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package aws.smithy.kotlin.runtime.tracing
+
+import io.kotest.matchers.string.shouldContain
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.test.runTest
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlin.test.assertFailsWith
+import kotlin.test.assertNull
+
+@OptIn(ExperimentalCoroutinesApi::class)
+class TraceSpanCoroutineUtilsTest {
+    @Test
+    fun testRootSpanNoExisting() = runTest {
+        val tracer = DefaultTracer(NoOpTraceProbe, "test")
+        val rootSpan = tracer.createRootSpan("root")
+        coroutineContext.withRootTraceSpan(rootSpan) {
+            val actualSpan = coroutineContext.traceSpan
+            assertEquals(rootSpan, actualSpan)
+        }
+    }
+
+    @Test
+    fun testRootSpanExistingIsParent() = runTest {
+        val tracer = DefaultTracer(NoOpTraceProbe, "test")
+        val rootSpan = tracer.createRootSpan("root")
+
+        coroutineContext.withRootTraceSpan(rootSpan) {
+            val currSpan = coroutineContext.traceSpan
+            assertNull(currSpan.parent)
+            val nestedTracer = currSpan.asNestedTracer("nested")
+            val nestedRoot = nestedTracer.createRootSpan("root2")
+            // this should be allowed since the existing span is the parent
+            coroutineContext.withRootTraceSpan(nestedRoot) {
+                val actualSpan = coroutineContext.traceSpan
+                assertEquals(currSpan, actualSpan.parent)
+            }
+        }
+    }
+
+    @Test
+    fun testRootSpanExistingIsNotParent() = runTest {
+        val tracer = DefaultTracer(NoOpTraceProbe, "test")
+        val rootSpan1 = tracer.createRootSpan("root1")
+        val rootSpan2 = tracer.createRootSpan("root2")
+        val illegalRoot = object : TraceSpan {
+            override val parent: TraceSpan = rootSpan2
+            override val id: String = "illegal span"
+            override fun postEvent(event: TraceEvent) {}
+            override fun child(id: String): TraceSpan { error("not needed for test") }
+            override fun close() {}
+        }
+
+        coroutineContext.withRootTraceSpan(rootSpan1) {
+            assertFailsWith<IllegalStateException> {
+                coroutineContext.withRootTraceSpan(illegalRoot) { Unit }
+            }.message.shouldContain("when no current span exists or the new span is a child of the active span")
+        }
+    }
+
+    @Test
+    fun testRootSpanExistingNoParent() = runTest {
+        // existing span with a new root created with no lineage, should become a child of the outer span automatically
+        val tracer = DefaultTracer(NoOpTraceProbe, "test")
+        val rootSpan1 = tracer.createRootSpan("root1")
+        val rootSpan2 = tracer.createRootSpan("root2")
+
+        coroutineContext.withRootTraceSpan(rootSpan1) {
+            val outerSpan = coroutineContext.traceSpan
+            coroutineContext.withRootTraceSpan(rootSpan2) {
+                val innerSpan = coroutineContext.traceSpan
+                assertEquals(outerSpan, innerSpan.parent)
+            }
+        }
+    }
+}