From 324e70c27edb7d38daae3fc1d5dbc9ddd3f80373 Mon Sep 17 00:00:00 2001
From: Lauri Tulmin <ltulmin@splunk.com>
Date: Wed, 22 Feb 2023 20:19:52 +0200
Subject: [PATCH] Fix kotlin coroutine context propagation

---
 .../KotlinCoroutinesInstrumentation.java      |  1 +
 .../KotlinCoroutinesInstrumentationTest.kt    | 39 +++++++++++++++++++
 2 files changed, 40 insertions(+)

diff --git a/instrumentation/kotlinx-coroutines/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/kotlinxcoroutines/KotlinCoroutinesInstrumentation.java b/instrumentation/kotlinx-coroutines/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/kotlinxcoroutines/KotlinCoroutinesInstrumentation.java
index 9910a0c53503..6b9ed930ae1a 100644
--- a/instrumentation/kotlinx-coroutines/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/kotlinxcoroutines/KotlinCoroutinesInstrumentation.java
+++ b/instrumentation/kotlinx-coroutines/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/kotlinxcoroutines/KotlinCoroutinesInstrumentation.java
@@ -25,6 +25,7 @@ public ElementMatcher<TypeDescription> typeMatcher() {
   public void transform(TypeTransformer transformer) {
     transformer.applyAdviceToMethod(
         named("newCoroutineContext")
+            .and(takesArgument(0, named("kotlinx.coroutines.CoroutineScope")))
             .and(takesArgument(1, named("kotlin.coroutines.CoroutineContext"))),
         this.getClass().getName() + "$ContextAdvice");
   }
diff --git a/instrumentation/kotlinx-coroutines/javaagent/src/test/kotlin/io/opentelemetry/javaagent/instrumentation/kotlinxcoroutines/KotlinCoroutinesInstrumentationTest.kt b/instrumentation/kotlinx-coroutines/javaagent/src/test/kotlin/io/opentelemetry/javaagent/instrumentation/kotlinxcoroutines/KotlinCoroutinesInstrumentationTest.kt
index bfa003e45a78..8a0ae43cb8ea 100644
--- a/instrumentation/kotlinx-coroutines/javaagent/src/test/kotlin/io/opentelemetry/javaagent/instrumentation/kotlinxcoroutines/KotlinCoroutinesInstrumentationTest.kt
+++ b/instrumentation/kotlinx-coroutines/javaagent/src/test/kotlin/io/opentelemetry/javaagent/instrumentation/kotlinxcoroutines/KotlinCoroutinesInstrumentationTest.kt
@@ -7,6 +7,7 @@ package io.opentelemetry.javaagent.instrumentation.kotlinxcoroutines
 
 import io.opentelemetry.context.Context
 import io.opentelemetry.context.ContextKey
+import io.opentelemetry.context.Scope
 import io.opentelemetry.extension.kotlin.asContextElement
 import io.opentelemetry.extension.kotlin.getOpenTelemetryContext
 import io.opentelemetry.instrumentation.reactor.v3_1.ContextPropagationOperator
@@ -20,6 +21,7 @@ import kotlinx.coroutines.CoroutineStart
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.GlobalScope
+import kotlinx.coroutines.ThreadContextElement
 import kotlinx.coroutines.asCoroutineDispatcher
 import kotlinx.coroutines.async
 import kotlinx.coroutines.awaitAll
@@ -54,6 +56,7 @@ import java.util.concurrent.Executors
 import java.util.concurrent.TimeUnit
 import java.util.function.Consumer
 import java.util.stream.Stream
+import kotlin.coroutines.CoroutineContext
 
 @TestInstance(TestInstance.Lifecycle.PER_CLASS)
 @ExperimentalCoroutinesApi
@@ -563,4 +566,40 @@ class KotlinCoroutinesInstrumentationTest {
   class DispatcherWrapper(val dispatcher: CoroutineDispatcher) {
     override fun toString(): String = dispatcher.toString()
   }
+
+  // regression test for https://github.com/open-telemetry/opentelemetry-java-instrumentation/issues/7837
+  // tests that a custom ThreadContextElement runs after KotlinContextElement that is used for
+  // context propagation in coroutines
+  @Test
+  fun `test custom context element`() {
+    val testValue = "test-value"
+    val contextKey = ContextKey.named<String>("test-key")
+    val scope = Context.current().with(contextKey, "wrong value").makeCurrent()
+    scope.use {
+      runBlocking {
+        val context = Context.current().with(contextKey, testValue)
+        withContext(TestContextElement(context)) {
+          delay(10)
+          val result = Context.current().get(contextKey)
+          assertThat(result).isEqualTo(testValue)
+        }
+      }
+    }
+  }
+
+  class TestContextElement(private val otelContext: Context) : ThreadContextElement<Scope> {
+    companion object Key : CoroutineContext.Key<TestContextElement> {
+    }
+
+    override val key: CoroutineContext.Key<TestContextElement>
+      get() = Key
+
+    override fun restoreThreadContext(context: CoroutineContext, oldState: Scope) {
+      oldState.close()
+    }
+
+    override fun updateThreadContext(context: CoroutineContext): Scope {
+      return otelContext.makeCurrent()
+    }
+  }
 }