Implement @WithSpan support for kotlin coroutines

dependencyManagement/build.gradle.kts

@@ -64,6 +64,7 @@ val CORE_DEPENDENCIES = listOf(
   "net.bytebuddy:byte-buddy-gradle-plugin:${byteBuddyVersion}",
   "org.ow2.asm:asm:${asmVersion}",
   "org.ow2.asm:asm-tree:${asmVersion}",
+  "org.ow2.asm:asm-util:${asmVersion}",
   "org.openjdk.jmh:jmh-core:${jmhVersion}",
   "org.openjdk.jmh:jmh-generator-bytecode:${jmhVersion}",
   "org.mockito:mockito-core:${mockitoVersion}",

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

@@ -10,25 +10,35 @@ muzzle {
     group.set("org.jetbrains.kotlinx")
     module.set("kotlinx-coroutines-core")
     versions.set("[1.0.0,1.3.8)")
+    extraDependency(project(":instrumentation-annotations"))
+    extraDependency("io.opentelemetry:opentelemetry-api:1.27.0")
   }
   // 1.3.9 (and beyond?) have changed how artifact names are resolved due to multiplatform variants
   pass {
     group.set("org.jetbrains.kotlinx")
     module.set("kotlinx-coroutines-core-jvm")
     versions.set("[1.3.9,)")
+    extraDependency(project(":instrumentation-annotations"))
+    extraDependency("io.opentelemetry:opentelemetry-api:1.27.0")
   }
 }
 
 dependencies {
   compileOnly("io.opentelemetry:opentelemetry-extension-kotlin")
   compileOnly("org.jetbrains.kotlin:kotlin-stdlib-jdk8")
+  compileOnly(project(":opentelemetry-instrumentation-annotations-shaded-for-instrumenting", configuration = "shadow"))
+
+  implementation("org.ow2.asm:asm-tree")
+  implementation("org.ow2.asm:asm-util")
+  implementation(project(":instrumentation:opentelemetry-instrumentation-annotations-1.16:javaagent"))
 
   testInstrumentation(project(":instrumentation:opentelemetry-extension-kotlin-1.0:javaagent"))
   testInstrumentation(project(":instrumentation:reactor:reactor-3.1:javaagent"))
 
   testImplementation("io.opentelemetry:opentelemetry-extension-kotlin")
   testImplementation("org.jetbrains.kotlin:kotlin-stdlib-jdk8")
   testImplementation(project(":instrumentation:reactor:reactor-3.1:library"))
+  testImplementation(project(":instrumentation-annotations"))
 
   // Use first version with flow support since we have tests for it.
   testLibrary("org.jetbrains.kotlinx:kotlinx-coroutines-core:1.3.0")
@@ -39,6 +49,7 @@ tasks {
   withType(KotlinCompile::class).configureEach {
     kotlinOptions {
       jvmTarget = "1.8"
+      javaParameters = true
     }
   }
 }

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

@@ -0,0 +1,169 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.kotlinxcoroutines.instrumentationannotations;
+
+import static io.opentelemetry.javaagent.instrumentation.kotlinxcoroutines.instrumentationannotations.AnnotationSingletons.instrumenter;
+
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.SpanKind;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.Scope;
+import io.opentelemetry.instrumentation.api.util.VirtualField;
+import kotlin.coroutines.Continuation;
+import kotlin.coroutines.intrinsics.IntrinsicsKt;
+
+public final class AnnotationInstrumentationHelper {
+
+  private static final VirtualField<Continuation<?>, Context> contextField =
+      VirtualField.find(Continuation.class, Context.class);
+
+  public static MethodRequest createMethodRequest(
+      Class<?> declaringClass, String methodName, String withSpanValue, String spanKindString) {
+    SpanKind spanKind = SpanKind.INTERNAL;
+    if (spanKindString != null) {
+      try {
+        spanKind = SpanKind.valueOf(spanKindString);
+      } catch (IllegalArgumentException exception) {
+        // ignore
+      }
+    }
+
+    return MethodRequest.create(declaringClass, methodName, withSpanValue, spanKind);
+  }
+
+  public static Context enterCoroutine(
+      int label, Continuation<?> continuation, MethodRequest request) {
+    // label 0 means that coroutine is started, any other label means that coroutine is resumed
+    if (label == 0) {
+      Context context = instrumenter().start(Context.current(), request);
+      // null continuation means that this method is not going to be resumed, and we don't need to
+      // store the context
+      if (continuation != null) {
+        contextField.set(continuation, context);
+      }
+      return context;
+    } else {
+      return continuation != null ? contextField.get(continuation) : null;
+    }
+  }
+
+  public static Scope openScope(Context context) {
+    return context != null ? context.makeCurrent() : null;
+  }
+
+  public static void exitCoroutine(
+      Object result,
+      MethodRequest request,
+      Continuation<?> continuation,
+      Context context,
+      Scope scope) {
+    exitCoroutine(null, result, request, continuation, context, scope);
+  }
+
+  public static void exitCoroutine(
+      Throwable error,
+      Object result,
+      MethodRequest request,
+      Continuation<?> continuation,
+      Context context,
+      Scope scope) {
+    if (scope == null) {
+      return;
+    }
+    scope.close();
+
+    // end the span when this method can not be resumed (coroutine is null) or if it has reached
+    // final state (returns anything else besides COROUTINE_SUSPENDED)
+    if (continuation == null || result != IntrinsicsKt.getCOROUTINE_SUSPENDED()) {
+      instrumenter().end(context, request, null, error);
+    }
+  }
+
+  public static void setSpanAttribute(int label, String name, boolean value) {
+    // only add the attribute when coroutine is started
+    if (label == 0) {
+      Span.current().setAttribute(name, value);
+    }
+  }
+
+  public static void setSpanAttribute(int label, String name, byte value) {
+    // only add the attribute when coroutine is started
+    if (label == 0) {
+      Span.current().setAttribute(name, value);
+    }
+  }
+
+  public static void setSpanAttribute(int label, String name, char value) {
+    // only add the attribute when coroutine is started
+    if (label == 0) {
+      Span.current().setAttribute(name, String.valueOf(value));
+    }
+  }
+
+  public static void setSpanAttribute(int label, String name, double value) {
+    // only add the attribute when coroutine is started
+    if (label == 0) {
+      Span.current().setAttribute(name, value);
+    }
+  }
+
+  public static void setSpanAttribute(int label, String name, float value) {
+    // only add the attribute when coroutine is started
+    if (label == 0) {
+      Span.current().setAttribute(name, value);
+    }
+  }
+
+  public static void setSpanAttribute(int label, String name, int value) {
+    // only add the attribute when coroutine is started
+    if (label == 0) {
+      Span.current().setAttribute(name, value);
+    }
+  }
+
+  public static void setSpanAttribute(int label, String name, long value) {
+    // only add the attribute when coroutine is started
+    if (label == 0) {
+      Span.current().setAttribute(name, value);
+    }
+  }
+
+  public static void setSpanAttribute(int label, String name, short value) {
+    // only add the attribute when coroutine is started
+    if (label == 0) {
+      Span.current().setAttribute(name, value);
+    }
+  }
+
+  public static void setSpanAttribute(int label, String name, Object value) {
+    // only add the attribute when coroutine is started
+    if (label != 0) {
+      return;
+    }
+    if (value instanceof String) {
+      Span.current().setAttribute(name, (String) value);
+    } else if (value instanceof Boolean) {
+      Span.current().setAttribute(name, (Boolean) value);
+    } else if (value instanceof Byte) {
+      Span.current().setAttribute(name, (Byte) value);
+    } else if (value instanceof Character) {
+      Span.current().setAttribute(name, (Character) value);
+    } else if (value instanceof Double) {
+      Span.current().setAttribute(name, (Double) value);
+    } else if (value instanceof Float) {
+      Span.current().setAttribute(name, (Float) value);
+    } else if (value instanceof Integer) {
+      Span.current().setAttribute(name, (Integer) value);
+    } else if (value instanceof Long) {
+      Span.current().setAttribute(name, (Long) value);
+    }
+    // TODO: arrays and List not supported see AttributeBindingFactoryTest
+  }
+
+  public static void init() {}
+
+  private AnnotationInstrumentationHelper() {}
+}

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

@@ -0,0 +1,46 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.kotlinxcoroutines.instrumentationannotations;
+
+import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasClassesNamed;
+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;
+import net.bytebuddy.matcher.ElementMatcher;
+
+/** Instrumentation for methods annotated with {@code WithSpan} annotation. */
+@AutoService(InstrumentationModule.class)
+public class AnnotationInstrumentationModule extends InstrumentationModule {
+
+  public AnnotationInstrumentationModule() {
+    super(
+        "kotlinx-coroutines-opentelemetry-instrumentation-annotations",
+        "kotlinx-coroutines",
+        "opentelemetry-instrumentation-annotations");
+  }
+
+  @Override
+  public int order() {
+    // Run first to ensure other automatic instrumentation is added after and therefore is executed
+    // earlier in the instrumented method and create the span to attach attributes to.
+    return -1000;
+  }
+
+  @Override
+  public ElementMatcher.Junction<ClassLoader> classLoaderMatcher() {
+    return hasClassesNamed(
+        "application.io.opentelemetry.instrumentation.annotations.WithSpan",
+        "kotlinx.coroutines.CoroutineContextKt");
+  }
+
+  @Override
+  public List<TypeInstrumentation> typeInstrumentations() {
+    return singletonList(new WithSpanInstrumentation());
+  }
+}

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

@@ -0,0 +1,42 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.kotlinxcoroutines.instrumentationannotations;
+
+import io.opentelemetry.api.GlobalOpenTelemetry;
+import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
+import io.opentelemetry.instrumentation.api.instrumenter.code.CodeAttributesExtractor;
+import io.opentelemetry.instrumentation.api.instrumenter.util.SpanNames;
+
+public final class AnnotationSingletons {
+
+  private static final String INSTRUMENTATION_NAME = "io.opentelemetry.kotlinx-coroutines";
+
+  private static final Instrumenter<MethodRequest, Object> INSTRUMENTER = createInstrumenter();
+
+  public static Instrumenter<MethodRequest, Object> instrumenter() {
+    return INSTRUMENTER;
+  }
+
+  private static Instrumenter<MethodRequest, Object> createInstrumenter() {
+    return Instrumenter.builder(
+            GlobalOpenTelemetry.get(),
+            INSTRUMENTATION_NAME,
+            AnnotationSingletons::spanNameFromMethodRequest)
+        .addAttributesExtractor(
+            CodeAttributesExtractor.create(MethodRequestCodeAttributesGetter.INSTANCE))
+        .buildInstrumenter(MethodRequest::getSpanKind);
+  }
+
+  private static String spanNameFromMethodRequest(MethodRequest request) {
+    String spanName = request.getWithSpanValue();
+    if (spanName == null || spanName.isEmpty()) {
+      spanName = SpanNames.fromMethod(request.getDeclaringClass(), request.getMethodName());
+    }
+    return spanName;
+  }
+
+  private AnnotationSingletons() {}
+}

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

@@ -0,0 +1,44 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.kotlinxcoroutines.instrumentationannotations;
+
+import io.opentelemetry.api.trace.SpanKind;
+
+public final class MethodRequest {
+  private final Class<?> declaringClass;
+  private final String methodName;
+  private final String withSpanValue;
+  private final SpanKind spanKind;
+
+  private MethodRequest(
+      Class<?> declaringClass, String methodName, String withSpanValue, SpanKind spanKind) {
+    this.declaringClass = declaringClass;
+    this.methodName = methodName;
+    this.withSpanValue = withSpanValue;
+    this.spanKind = spanKind;
+  }
+
+  public static MethodRequest create(
+      Class<?> declaringClass, String methodName, String withSpanValue, SpanKind spanKind) {
+    return new MethodRequest(declaringClass, methodName, withSpanValue, spanKind);
+  }
+
+  public Class<?> getDeclaringClass() {
+    return declaringClass;
+  }
+
+  public String getMethodName() {
+    return methodName;
+  }
+
+  public String getWithSpanValue() {
+    return withSpanValue;
+  }
+
+  public SpanKind getSpanKind() {
+    return spanKind;
+  }
+}

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

@@ -0,0 +1,22 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.kotlinxcoroutines.instrumentationannotations;
+
+import io.opentelemetry.instrumentation.api.instrumenter.code.CodeAttributesGetter;
+
+enum MethodRequestCodeAttributesGetter implements CodeAttributesGetter<MethodRequest> {
+  INSTANCE;
+
+  @Override
+  public Class<?> getCodeClass(MethodRequest methodRequest) {
+    return methodRequest.getDeclaringClass();
+  }
+
+  @Override
+  public String getMethodName(MethodRequest methodRequest) {
+    return methodRequest.getMethodName();
+  }
+}