Add instrumentation for opentelemetry-extension-kotlin (#7341)

Hopefully resolves
https://github.com/open-telemetry/opentelemetry-java-instrumentation/issues/7124
Our kotlin coroutine instrumentation relies on a shaded copy of
`opentelemetry-extension-kotlin`. This doesn't work well when
application also uses `opentelemetry-extension-kotlin`, because the
shaded and unshaded copy store opentelemery context under different key.
This pr attempts to fix this by instrumenting
`opentelemetry-extension-kotlin` provided by the application so that it
would delegate to the one shaded inside the agent.

Co-authored-by: Mateusz Rzeszutek <mrzeszutek@splunk.com>

instrumentation/kotlinx-coroutines/javaagent/build.gradle.kts

@@ -23,6 +23,7 @@ dependencies {
   compileOnly("io.opentelemetry:opentelemetry-extension-kotlin")
   compileOnly("org.jetbrains.kotlin:kotlin-stdlib-jdk8")
 
+  testInstrumentation(project(":instrumentation:opentelemetry-extension-kotlin-1.0:javaagent"))
   testInstrumentation(project(":instrumentation:reactor:reactor-3.1:javaagent"))
 
   testImplementation("io.opentelemetry:opentelemetry-extension-kotlin")

instrumentation/kotlinx-coroutines/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/kotlinxcoroutines/KotlinCoroutinesInstrumentationHelper.java

@@ -14,7 +14,7 @@ public final class KotlinCoroutinesInstrumentationHelper {
   public static CoroutineContext addOpenTelemetryContext(CoroutineContext coroutineContext) {
     Context current = Context.current();
     Context inCoroutine = ContextExtensionsKt.getOpenTelemetryContext(coroutineContext);
-    if (current == inCoroutine) {
+    if (current == inCoroutine || inCoroutine != Context.root()) {
       return coroutineContext;
     }
     return coroutineContext.plus(ContextExtensionsKt.asContextElement(current));

instrumentation/kotlinx-coroutines/javaagent/src/test/kotlin/io/opentelemetry/javaagent/instrumentation/kotlinxcoroutines/KotlinCoroutinesInstrumentationTest.kt

@@ -6,11 +6,12 @@
 package io.opentelemetry.javaagent.instrumentation.kotlinxcoroutines
 
 import io.opentelemetry.context.Context
+import io.opentelemetry.context.ContextKey
 import io.opentelemetry.extension.kotlin.asContextElement
+import io.opentelemetry.extension.kotlin.getOpenTelemetryContext
 import io.opentelemetry.instrumentation.reactor.ContextPropagationOperator
 import io.opentelemetry.instrumentation.testing.junit.AgentInstrumentationExtension
 import io.opentelemetry.instrumentation.testing.util.TelemetryDataUtil.orderByRootSpanName
-import io.opentelemetry.sdk.testing.assertj.OpenTelemetryAssertions.assertThat
 import io.opentelemetry.sdk.testing.assertj.TraceAssert
 import kotlinx.coroutines.CompletableDeferred
 import kotlinx.coroutines.CoroutineDispatcher
@@ -38,6 +39,7 @@ import kotlinx.coroutines.selects.select
 import kotlinx.coroutines.withContext
 import kotlinx.coroutines.withTimeout
 import kotlinx.coroutines.yield
+import org.assertj.core.api.Assertions.assertThat
 import org.awaitility.Awaitility.await
 import org.junit.jupiter.api.AfterAll
 import org.junit.jupiter.api.Test
@@ -92,41 +94,33 @@ class KotlinCoroutinesInstrumentationTest {
 
     testing.waitAndAssertTraces(
       { trace ->
-        // TODO(anuraaga): Need hasSpansSatisfyingExactlyInAnyOrder sometimes
-        trace.satisfiesExactlyInAnyOrder(
-          Consumer {
-            assertThat(it)
-              .hasName("parent")
+        trace.hasSpansSatisfyingExactlyInAnyOrder(
+          {
+            it.hasName("parent")
               .hasNoParent()
           },
-          Consumer {
-            assertThat(it)
-              .hasName("produce_0")
+          {
+            it.hasName("produce_0")
               .hasParent(trace.getSpan(0))
           },
-          Consumer {
-            assertThat(it)
-              .hasName("consume_0")
+          {
+            it.hasName("consume_0")
               .hasParent(trace.getSpan(0))
           },
-          Consumer {
-            assertThat(it)
-              .hasName("produce_1")
+          {
+            it.hasName("produce_1")
               .hasParent(trace.getSpan(0))
           },
-          Consumer {
-            assertThat(it)
-              .hasName("consume_1")
+          {
+            it.hasName("consume_1")
               .hasParent(trace.getSpan(0))
           },
-          Consumer {
-            assertThat(it)
-              .hasName("produce_2")
+          {
+            it.hasName("produce_2")
               .hasParent(trace.getSpan(0))
           },
-          Consumer {
-            assertThat(it)
-              .hasName("consume_2")
+          {
+            it.hasName("consume_2")
               .hasParent(trace.getSpan(0))
           }
         )
@@ -225,25 +219,25 @@ class KotlinCoroutinesInstrumentationTest {
 
     testing.waitAndAssertTraces(
       { trace ->
-        trace.satisfiesExactlyInAnyOrder(
-          Consumer {
-            assertThat(it).hasName("parent")
+        trace.hasSpansSatisfyingExactlyInAnyOrder(
+          {
+            it.hasName("parent")
               .hasNoParent()
           },
-          Consumer {
-            assertThat(it).hasName("future1")
+          {
+            it.hasName("future1")
               .hasParent(trace.getSpan(0))
           },
-          Consumer {
-            assertThat(it).hasName("keptPromise")
+          {
+            it.hasName("keptPromise")
               .hasParent(trace.getSpan(0))
           },
-          Consumer {
-            assertThat(it).hasName("keptPromise2")
+          {
+            it.hasName("keptPromise2")
               .hasParent(trace.getSpan(0))
           },
-          Consumer {
-            assertThat(it).hasName("brokenPromise")
+          {
+            it.hasName("brokenPromise")
               .hasParent(trace.getSpan(0))
           }
         )
@@ -280,25 +274,24 @@ class KotlinCoroutinesInstrumentationTest {
 
     testing.waitAndAssertTraces(
       { trace ->
-        // TODO(anuraaga): Need hasSpansSatisfyingExactlyInAnyOrder sometimes
-        trace.satisfiesExactlyInAnyOrder(
-          Consumer {
-            assertThat(it)
+        trace.hasSpansSatisfyingExactlyInAnyOrder(
+          {
+            it
               .hasName("parent")
               .hasNoParent()
           },
-          Consumer {
-            assertThat(it)
+          {
+            it
               .hasName("timeout1")
               .hasParent(trace.getSpan(0))
           },
-          Consumer {
-            assertThat(it)
+          {
+            it
               .hasName("timeout2")
               .hasParent(trace.getSpan(0))
           },
-          Consumer {
-            assertThat(it)
+          {
+            it
               .hasName("timeout3")
               .hasParent(trace.getSpan(0))
           }
@@ -459,6 +452,45 @@ class KotlinCoroutinesInstrumentationTest {
     )
   }
 
+  private val ANIMAL: ContextKey<String> = ContextKey.named("animal")
+
+  @ParameterizedTest
+  @ArgumentsSource(DispatchersSource::class)
+  fun `context contains expected value`(dispatcher: DispatcherWrapper) {
+    runTest(dispatcher) {
+      val context1 = Context.current().with(ANIMAL, "cat")
+      runBlocking(context1.asContextElement()) {
+        assertThat(Context.current().get(ANIMAL)).isEqualTo("cat")
+        assertThat(coroutineContext.getOpenTelemetryContext().get(ANIMAL)).isEqualTo("cat")
+        tracedChild("nested1")
+        withContext(context1.with(ANIMAL, "dog").asContextElement()) {
+          assertThat(Context.current().get(ANIMAL)).isEqualTo("dog")
+          assertThat(coroutineContext.getOpenTelemetryContext().get(ANIMAL)).isEqualTo("dog")
+          tracedChild("nested2")
+        }
+      }
+    }
+
+    testing.waitAndAssertTraces(
+      { trace ->
+        trace.hasSpansSatisfyingExactly(
+          {
+            it.hasName("parent")
+              .hasNoParent()
+          },
+          {
+            it.hasName("nested1")
+              .hasParent(trace.getSpan(0))
+          },
+          {
+            it.hasName("nested2")
+              .hasParent(trace.getSpan(0))
+          }
+        )
+      }
+    )
+  }
+
   private fun tracedChild(opName: String) {
     tracer.spanBuilder(opName).startSpan().end()
   }

instrumentation/opentelemetry-extension-kotlin-1.0/javaagent/build.gradle.kts

@@ -0,0 +1,51 @@
+plugins {
+  id("org.jetbrains.kotlin.jvm")
+  id("otel.javaagent-instrumentation")
+}
+
+muzzle {
+  pass {
+    group.set("io.opentelemetry")
+    module.set("opentelemetry-extension-kotlin")
+    versions.set("[0.17.0,)")
+    assertInverse.set(true)
+    skip("0.13.0") // has a bad dependency on non-alpha api-metric 0.13.0
+    extraDependency("org.jetbrains.kotlinx:kotlinx-coroutines-core:1.3.0")
+  }
+}
+
+dependencies {
+  compileOnly("org.jetbrains.kotlin:kotlin-stdlib-jdk8")
+  library("io.opentelemetry:opentelemetry-extension-kotlin")
+  // see the comment in opentelemetry-api-1.0.gradle for more details
+  compileOnly(project(":opentelemetry-api-shaded-for-instrumenting", configuration = "shadow"))
+
+  implementation(project(":instrumentation:opentelemetry-api:opentelemetry-api-1.0:javaagent"))
+
+  testImplementation("org.jetbrains.kotlin:kotlin-stdlib-jdk8")
+  testImplementation("org.jetbrains.kotlinx:kotlinx-coroutines-core:1.3.0")
+}
+
+if (!(findProperty("testLatestDeps") as Boolean)) {
+  // run tests against an early version of opentelemetry-extension-kotlin, latest dep tests will use
+  // the current version
+  configurations.configureEach {
+    if (!name.contains("muzzle")) {
+      resolutionStrategy {
+        eachDependency {
+          if (requested.group == "io.opentelemetry" && requested.name == "opentelemetry-extension-kotlin") {
+            useVersion("1.0.0")
+          }
+        }
+      }
+    }
+  }
+}
+
+tasks {
+  withType(org.jetbrains.kotlin.gradle.tasks.KotlinCompile::class).configureEach {
+    kotlinOptions {
+      jvmTarget = "1.8"
+    }
+  }
+}

instrumentation/opentelemetry-extension-kotlin-1.0/javaagent/gradle.properties

@@ -0,0 +1 @@
+kotlin.stdlib.default.dependency=false

instrumentation/opentelemetry-extension-kotlin-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/extensionkotlin/ContextExtensionInstrumentation.java

@@ -0,0 +1,118 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.extensionkotlin;
+
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+
+import application.io.opentelemetry.context.Context;
+import io.opentelemetry.extension.kotlin.ContextExtensionsKt;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import io.opentelemetry.javaagent.instrumentation.opentelemetryapi.context.AgentContextStorage;
+import kotlin.coroutines.CoroutineContext;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+public class ContextExtensionInstrumentation implements TypeInstrumentation {
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return named("application.io.opentelemetry.extension.kotlin.ContextExtensionsKt");
+  }
+
+  @Override
+  public void transform(TypeTransformer transformer) {
+    transformer.applyAdviceToMethod(
+        named("asContextElement")
+            .and(takesArgument(0, named("application.io.opentelemetry.context.Context"))),
+        this.getClass().getName() + "$ContextAdvice");
+
+    transformer.applyAdviceToMethod(
+        named("asContextElement")
+            .and(
+                takesArgument(
+                    0, named("application.io.opentelemetry.context.ImplicitContextKeyed"))),
+        this.getClass().getName() + "$ImplicitContextKeyedAdvice");
+
+    transformer.applyAdviceToMethod(
+        named("getOpenTelemetryContext")
+            .and(takesArgument(0, named("kotlin.coroutines.CoroutineContext"))),
+        this.getClass().getName() + "$GetOpenTelemetryContextAdvice");
+  }
+
+  @SuppressWarnings("unused")
+  public static class ContextAdvice {
+
+    @Advice.OnMethodEnter(skipOn = Advice.OnNonDefaultValue.class)
+    public static CoroutineContext enter(@Advice.Argument(0) Context applicationContext) {
+      if (applicationContext != null) {
+        io.opentelemetry.context.Context agentContext =
+            AgentContextStorage.getAgentContext(applicationContext);
+        return ContextExtensionsKt.asContextElement(agentContext);
+      }
+      return null;
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class)
+    public static void onExit(
+        @Advice.Return(readOnly = false) CoroutineContext result,
+        @Advice.Enter CoroutineContext coroutineContext) {
+      if (coroutineContext != null) {
+        result = coroutineContext;
+      }
+    }
+  }
+
+  @SuppressWarnings("unused")
+  public static class ImplicitContextKeyedAdvice {
+
+    @Advice.OnMethodEnter(skipOn = Advice.OnNonDefaultValue.class)
+    public static CoroutineContext enter(
+        @Advice.Argument(0)
+            application.io.opentelemetry.context.ImplicitContextKeyed implicitContextKeyed) {
+      if (implicitContextKeyed != null) {
+        Context applicationContext = Context.current().with(implicitContextKeyed);
+        io.opentelemetry.context.Context agentContext =
+            AgentContextStorage.getAgentContext(applicationContext);
+        return ContextExtensionsKt.asContextElement(agentContext);
+      }
+      return null;
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class)
+    public static void onExit(
+        @Advice.Return(readOnly = false) CoroutineContext result,
+        @Advice.Enter CoroutineContext coroutineContext) {
+      if (coroutineContext != null) {
+        result = coroutineContext;
+      }
+    }
+  }
+
+  @SuppressWarnings("unused")
+  public static class GetOpenTelemetryContextAdvice {
+
+    @Advice.OnMethodEnter(skipOn = Advice.OnNonDefaultValue.class)
+    public static Context enter(@Advice.Argument(0) CoroutineContext coroutineContext) {
+      if (coroutineContext != null) {
+        io.opentelemetry.context.Context agentContext =
+            ContextExtensionsKt.getOpenTelemetryContext(coroutineContext);
+        return AgentContextStorage.toApplicationContext(agentContext);
+      }
+      return null;
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class)
+    public static void onExit(
+        @Advice.Return(readOnly = false) Context result, @Advice.Enter Context context) {
+      if (context != null) {
+        result = context;
+      }
+    }
+  }
+}

instrumentation/opentelemetry-extension-kotlin-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/extensionkotlin/ContextExtensionInstrumentationModule.java

@@ -0,0 +1,31 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.extensionkotlin;
+
+import static java.util.Collections.singletonList;
+
+import com.google.auto.service.AutoService;
+import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import java.util.List;
+
+@AutoService(InstrumentationModule.class)
+public class ContextExtensionInstrumentationModule extends InstrumentationModule {
+
+  public ContextExtensionInstrumentationModule() {
+    super("opentelemetry-extension-kotlin");
+  }
+
+  @Override
+  public boolean isHelperClass(String className) {
+    return className.startsWith("io.opentelemetry.extension.kotlin.");
+  }
+
+  @Override
+  public List<TypeInstrumentation> typeInstrumentations() {
+    return singletonList(new ContextExtensionInstrumentation());
+  }
+}

instrumentation/opentelemetry-extension-kotlin-1.0/javaagent/src/test/kotlin/io/opentelemetry/javaagent/instrumentation/extensionkotlin/ContextExtensionInstrumentationTest.kt

@@ -0,0 +1,31 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.extensionkotlin
+
+import io.opentelemetry.context.Context
+import io.opentelemetry.context.ContextKey
+import io.opentelemetry.extension.kotlin.asContextElement
+import io.opentelemetry.extension.kotlin.getOpenTelemetryContext
+import org.assertj.core.api.Assertions.assertThat
+import org.junit.jupiter.api.Test
+
+class ContextExtensionInstrumentationTest {
+  private val ANIMAL: ContextKey<String> = ContextKey.named("animal")
+
+  @Test
+  fun `is instrumented`() {
+    val context1 = Context.root().with(ANIMAL, "cat")
+    val contextElement = context1.asContextElement()
+    // check that the context element is from the opentelemetry-extension-kotlin that is shaded
+    // inside the agent
+    assertThat(contextElement.javaClass.name).startsWith("io.opentelemetry.javaagent.shaded")
+    val context2 = contextElement.getOpenTelemetryContext()
+    assertThat(context2.get(ANIMAL)).isEqualTo("cat")
+    // instrumentation does not preserve context identity due to conversion between application and
+    // agent context
+    assert(context1 != context2) { "Not instrumented" }
+  }
+}

settings.gradle.kts

@@ -361,6 +361,7 @@ hideFromDependabot(":instrumentation:opentelemetry-api:opentelemetry-api-1.4:jav
 hideFromDependabot(":instrumentation:opentelemetry-api:opentelemetry-api-1.10:javaagent")
 hideFromDependabot(":instrumentation:opentelemetry-api:opentelemetry-api-logs-1.19:javaagent")
 hideFromDependabot(":instrumentation:opentelemetry-extension-annotations-1.0:javaagent")
+hideFromDependabot(":instrumentation:opentelemetry-extension-kotlin-1.0:javaagent")
 hideFromDependabot(":instrumentation:opentelemetry-instrumentation-annotations-1.16:javaagent")
 hideFromDependabot(":instrumentation:opentelemetry-instrumentation-api:javaagent")
 hideFromDependabot(":instrumentation:opentelemetry-instrumentation-api:testing")