ZIO 2.0 instrumentation

docs/supported-libraries.md

@@ -130,6 +130,7 @@ These are the supported libraries and frameworks:
 | [Vert.x Kafka Client](https://vertx.io/docs/vertx-kafka-client/java/)                                                                       | 3.6+                          | N/A                                                                                                                                                                                                                                                                                                                                                                                     | [Messaging Spans]                                                                      |
 | [Vert.x RxJava2](https://vertx.io/docs/vertx-rx/java2/)                                                                                     | 3.5+                          | N/A                                                                                                                                                                                                                                                                                                                                                                                     | context propagation only                                                               |
 | [Vibur DBCP](https://www.vibur.org/)                                                                                                        | 11.0+                         | [opentelemetry-vibur-dbcp-11.0](../instrumentation/vibur-dbcp-11.0/library)                                                                                                                                                                                                                                                                                                             | [Database Pool Metrics]                                                                |
+| [ZIO](https://zio.dev/)                                                                                                                     | 2.0.0+                        | N/A                                                                                                                                                                                                                                                                                                                                                                                     | Context propagation                                                                    |
 
 **[1]** Standalone library instrumentation refers to instrumentation that can be used without the Java agent.
 

instrumentation/zio/zio-2.0/javaagent/build.gradle.kts

@@ -0,0 +1,43 @@
+plugins {
+  id("otel.javaagent-instrumentation")
+  id("otel.nullaway-conventions")
+  id("otel.scala-conventions")
+}
+
+val zioVersion = "2.0.0"
+val scalaVersion = "2.12"
+
+muzzle {
+  pass {
+    group.set("dev.zio")
+    module.set("zio_2.12")
+    versions.set("[$zioVersion,)")
+    assertInverse.set(true)
+  }
+  pass {
+    group.set("dev.zio")
+    module.set("zio_2.13")
+    versions.set("[$zioVersion,)")
+    assertInverse.set(true)
+  }
+  pass {
+    group.set("dev.zio")
+    module.set("zio_3")
+    versions.set("[$zioVersion,)")
+    assertInverse.set(true)
+  }
+}
+
+dependencies {
+  compileOnly("dev.zio:zio_$scalaVersion:$zioVersion")
+
+  testImplementation("dev.zio:zio_$scalaVersion:$zioVersion")
+
+  latestDepTestLibrary("dev.zio:zio_$scalaVersion:+")
+}
+
+tasks {
+  withType<Test>().configureEach {
+    jvmArgs("-Dotel.instrumentation.zio.enabled=true")
+  }
+}

instrumentation/zio/zio-2.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/zio/v2_0/FiberContext.java

@@ -0,0 +1,41 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.zio.v2_0;
+
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.context.Scope;
+import javax.annotation.Nullable;
+
+public class FiberContext {
+
+  private Span span;
+  @Nullable private Scope scope;
+
+  private FiberContext(Span span) {
+    this.span = span;
+  }
+
+  public static FiberContext create() {
+    return new FiberContext(Span.current());
+  }
+
+  public void onEnd() {
+    if (this.scope != null) {
+      this.scope.close();
+    }
+  }
+
+  public void onSuspend() {
+    this.span = Span.current();
+    if (this.scope != null) {
+      this.scope.close();
+    }
+  }
+
+  public void onResume() {
+    this.scope = this.span.makeCurrent();
+  }
+}

instrumentation/zio/zio-2.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/zio/v2_0/TracingSupervisor.java

@@ -0,0 +1,69 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.zio.v2_0;
+
+import io.opentelemetry.instrumentation.api.util.VirtualField;
+import scala.Option;
+import zio.Exit;
+import zio.Fiber;
+import zio.Supervisor;
+import zio.Unsafe;
+import zio.ZEnvironment;
+import zio.ZIO;
+import zio.ZIO$;
+
+@SuppressWarnings("unchecked")
+public final class TracingSupervisor extends Supervisor<Object> {
+
+  @SuppressWarnings("rawtypes")
+  private final VirtualField<Fiber.Runtime, FiberContext> virtualField;
+
+  @SuppressWarnings("rawtypes")
+  public TracingSupervisor(VirtualField<Fiber.Runtime, FiberContext> virtualField) {
+    this.virtualField = virtualField;
+  }
+
+  @Override
+  @SuppressWarnings("rawtypes")
+  public ZIO value(Object trace) {
+    return ZIO$.MODULE$.unit();
+  }
+
+  @Override
+  public <R, E, A1> void onStart(
+      ZEnvironment<R> environment,
+      ZIO<R, E, A1> effect,
+      Option<Fiber.Runtime<Object, Object>> parent,
+      Fiber.Runtime<E, A1> fiber,
+      Unsafe unsafe) {
+    FiberContext context = FiberContext.create();
+    virtualField.set(fiber, context);
+  }
+
+  @Override
+  public <R, E, A1> void onEnd(Exit<E, A1> value, Fiber.Runtime<E, A1> fiber, Unsafe unsafe) {
+    FiberContext context = virtualField.get(fiber);
+    if (context != null) {
+      context.onEnd();
+    }
+  }
+
+  @Override
+  public <E, A1> void onSuspend(Fiber.Runtime<E, A1> fiber, Unsafe unsafe) {
+    FiberContext context = virtualField.get(fiber);
+    if (context != null) {
+      context.onSuspend();
+    }
+  }
+
+  @Override
+  public <E, A1> void onResume(Fiber.Runtime<E, A1> fiber, Unsafe unsafe) {
+    FiberContext context = virtualField.get(fiber);
+    if (context != null) {
+      context.onResume();
+    }
+  }
+}

instrumentation/zio/zio-2.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/zio/v2_0/ZioInstrumentationModule.java

@@ -0,0 +1,34 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.zio.v2_0;
+
+import static java.util.Arrays.asList;
+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 ZioInstrumentationModule extends InstrumentationModule {
+
+  public ZioInstrumentationModule() {
+    super("zio");
+  }
+
+  @Override
+  public List<TypeInstrumentation> typeInstrumentations() {
+    return singletonList(new ZioRuntimeInstrumentation());
+  }
+
+  @Override
+  public List<String> getAdditionalHelperClassNames() {
+    return asList(
+        "io.opentelemetry.javaagent.instrumentation.zio.v2_0.FiberContext",
+        "io.opentelemetry.javaagent.instrumentation.zio.v2_0.TracingSupervisor");
+  }
+}

instrumentation/zio/zio-2.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/zio/v2_0/ZioRuntimeInstrumentation.java

@@ -0,0 +1,51 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.zio.v2_0;
+
+import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasClassesNamed;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+
+import io.opentelemetry.instrumentation.api.util.VirtualField;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+import zio.Fiber;
+import zio.Supervisor;
+
+public class ZioRuntimeInstrumentation implements TypeInstrumentation {
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return named("zio.Runtime$");
+  }
+
+  @Override
+  public ElementMatcher<ClassLoader> classLoaderOptimization() {
+    return hasClassesNamed("zio.Runtime$");
+  }
+
+  @Override
+  public void transform(TypeTransformer transformer) {
+    transformer.applyAdviceToMethod(
+        isMethod().and(named("defaultSupervisor")), getClass().getName() + "$DefaultSupervisor");
+  }
+
+  public static final class DefaultSupervisor {
+
+    private DefaultSupervisor() {}
+
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void onExit(@Advice.Return(readOnly = false) Supervisor<?> supervisor) {
+      @SuppressWarnings("rawtypes")
+      VirtualField<Fiber.Runtime, FiberContext> virtualField =
+          VirtualField.find(Fiber.Runtime.class, FiberContext.class);
+      supervisor = supervisor.$plus$plus(new TracingSupervisor(virtualField));
+    }
+  }
+}

instrumentation/zio/zio-2.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/zio/v2_0/package-info.java

@@ -0,0 +1,8 @@
+@DefaultQualifier(
+    value = NonNull.class,
+    locations = {TypeUseLocation.FIELD, TypeUseLocation.PARAMETER, TypeUseLocation.RETURN})
+package io.opentelemetry.javaagent.instrumentation.zio.v2_0;
+
+import org.checkerframework.checker.nullness.qual.NonNull;
+import org.checkerframework.framework.qual.DefaultQualifier;
+import org.checkerframework.framework.qual.TypeUseLocation;

instrumentation/zio/zio-2.0/javaagent/src/test/scala/io/opentelemetry/javaagent/instrumentation/zio/v2_0/ZioRuntimeInstrumentationTest.scala

@@ -0,0 +1,76 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.zio.v2_0
+
+import io.opentelemetry.instrumentation.testing.junit._
+import io.opentelemetry.javaagent.instrumentation.zio.v2_0.ZioTestFixtures._
+import io.opentelemetry.sdk.testing.assertj.{SpanDataAssert, TraceAssert}
+import org.junit.jupiter.api.extension.RegisterExtension
+import org.junit.jupiter.api.{Test, TestInstance}
+
+import java.util.function.Consumer
+
+@TestInstance(TestInstance.Lifecycle.PER_CLASS)
+class ZioRuntimeInstrumentationTest {
+
+  @RegisterExtension
+  val testing: InstrumentationExtension = AgentInstrumentationExtension.create()
+
+  @Test
+  def traceIsPropagatedToChildFiber(): Unit = {
+    runNestedFibers()
+
+    testing.waitAndAssertTraces(
+      assertTrace { trace =>
+        trace.hasSpansSatisfyingExactly(
+          assertSpan(_.hasName("fiber_1_span_1").hasNoParent),
+          assertSpan(_.hasName("fiber_2_span_1").hasParent(trace.getSpan(0)))
+        )
+      }
+    )
+  }
+
+  @Test
+  def traceIsPreservedWhenFiberIsInterrupted(): Unit = {
+    runInterruptedFiber()
+
+    testing.waitAndAssertTraces(
+      assertTrace { trace =>
+        trace.hasSpansSatisfyingExactly(
+          assertSpan(_.hasName("fiber_1_span_1").hasNoParent),
+          assertSpan(_.hasName("fiber_2_span_1").hasParent(trace.getSpan(0)))
+        )
+      }
+    )
+  }
+
+  @Test
+  def concurrentFibersDoNotInterfereWithEachOthersTraces(): Unit = {
+    runConcurrentFibers()
+
+    testing.waitAndAssertTraces(
+      assertTrace { trace =>
+        trace.hasSpansSatisfyingExactly(
+          assertSpan(_.hasName("fiber_1_span_1").hasNoParent),
+          assertSpan(_.hasName("fiber_1_span_2").hasParent(trace.getSpan(0)))
+        )
+      },
+      assertTrace { trace =>
+        trace.hasSpansSatisfyingExactly(
+          assertSpan(_.hasName("fiber_2_span_1").hasNoParent),
+          assertSpan(_.hasName("fiber_2_span_2").hasParent(trace.getSpan(0)))
+        )
+      }
+    )
+  }
+
+  private def assertTrace(f: TraceAssert => Any): Consumer[TraceAssert] =
+    (t: TraceAssert) => f(t)
+
+  private def assertSpan(f: SpanDataAssert => Any): Consumer[SpanDataAssert] =
+    (t: SpanDataAssert) => f(t)
+
+}