Allow reading otel context from reactor ContextView (#11235)

Co-authored-by: Trask Stalnaker <trask.stalnaker@gmail.com>

instrumentation/reactor/reactor-3.1/javaagent/build.gradle.kts

@@ -14,13 +14,16 @@ muzzle {
 }
 
 tasks.withType<Test>().configureEach {
+  systemProperty("testLatestDeps", findProperty("testLatestDeps") as Boolean)
   // TODO run tests both with and without experimental span attributes
   jvmArgs("-Dotel.instrumentation.reactor.experimental-span-attributes=true")
 }
 
 dependencies {
+  // we compile against 3.4.0, so we could use reactor.util.context.ContextView
+  // instrumentation is tested against 3.1.0.RELEASE
+  compileOnly("io.projectreactor:reactor-core:3.4.0")
   implementation(project(":instrumentation:reactor:reactor-3.1:library"))
-  library("io.projectreactor:reactor-core:3.1.0.RELEASE")
 
   implementation(project(":instrumentation:opentelemetry-api:opentelemetry-api-1.0:javaagent"))
 
@@ -30,14 +33,12 @@ dependencies {
 
   testInstrumentation(project(":instrumentation:opentelemetry-extension-annotations-1.0:javaagent"))
 
+  testLibrary("io.projectreactor:reactor-core:3.1.0.RELEASE")
   testLibrary("io.projectreactor:reactor-test:3.1.0.RELEASE")
   testImplementation(project(":instrumentation-annotations-support-testing"))
   testImplementation(project(":instrumentation:reactor:reactor-3.1:testing"))
   testImplementation(project(":instrumentation-annotations"))
   testImplementation("io.opentelemetry:opentelemetry-extension-annotations")
-
-  latestDepTestLibrary("io.projectreactor:reactor-core:3.4.+")
-  latestDepTestLibrary("io.projectreactor:reactor-test:3.4.+")
 }
 
 testing {
@@ -46,7 +47,11 @@ testing {
       dependencies {
         implementation(project(":instrumentation:reactor:reactor-3.1:library"))
         implementation(project(":instrumentation-annotations"))
-        implementation("io.projectreactor:reactor-test:3.1.0.RELEASE")
+        if (findProperty("testLatestDeps") as Boolean) {
+          implementation("io.projectreactor:reactor-test:+")
+        } else {
+          implementation("io.projectreactor:reactor-test:3.1.0.RELEASE")
+        }
       }
     }
   }

instrumentation/reactor/reactor-3.1/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/reactor/v3_1/BaseMonoWithSpanTest.java

@@ -88,7 +88,8 @@ abstract class BaseMonoWithSpanTest extends AbstractWithSpanTest<Mono<String>, M
                     span ->
                         span.hasName("inner-manual")
                             .hasKind(SpanKind.INTERNAL)
-                            .hasParent(trace.getSpan(1))
+                            // earliest tested and latest version behave differently
+                            .hasParent(trace.getSpan(Boolean.getBoolean("testLatestDeps") ? 0 : 1))
                             .hasAttributes(Attributes.empty())));
   }
 
@@ -130,7 +131,7 @@ abstract class BaseMonoWithSpanTest extends AbstractWithSpanTest<Mono<String>, M
                     span ->
                         span.hasName("inner-manual")
                             .hasKind(SpanKind.INTERNAL)
-                            .hasParent(trace.getSpan(1))
+                            .hasParent(trace.getSpan(Boolean.getBoolean("testLatestDeps") ? 0 : 1))
                             .hasAttributes(Attributes.empty())));
   }
 

instrumentation/reactor/reactor-3.1/library/build.gradle.kts

@@ -3,12 +3,19 @@ plugins {
 }
 
 dependencies {
-  library("io.projectreactor:reactor-core:3.1.0.RELEASE")
+  // we compile against 3.4.0, so we could use reactor.util.context.ContextView
+  // instrumentation is expected it to work with 3.1.0.RELEASE
+  compileOnly("io.projectreactor:reactor-core:3.4.0")
+  compileOnly(project(":muzzle")) // For @NoMuzzle
   implementation(project(":instrumentation-annotations-support"))
+  testLibrary("io.projectreactor:reactor-core:3.1.0.RELEASE")
   testLibrary("io.projectreactor:reactor-test:3.1.0.RELEASE")
 
   testImplementation(project(":instrumentation:reactor:reactor-3.1:testing"))
+}
 
-  latestDepTestLibrary("io.projectreactor:reactor-core:3.4.+")
-  latestDepTestLibrary("io.projectreactor:reactor-test:3.4.+")
+tasks {
+  withType<Test>().configureEach {
+    systemProperty("testLatestDeps", findProperty("testLatestDeps") as Boolean)
+  }
 }

instrumentation/reactor/reactor-3.1/library/src/main/java/io/opentelemetry/instrumentation/reactor/v3_1/ContextPropagationOperator.java

@@ -27,6 +27,7 @@ import static java.lang.invoke.MethodType.methodType;
 import io.opentelemetry.context.Context;
 import io.opentelemetry.context.Scope;
 import io.opentelemetry.instrumentation.api.annotation.support.async.AsyncOperationEndStrategies;
+import io.opentelemetry.javaagent.tooling.muzzle.NoMuzzle;
 import java.lang.invoke.MethodHandle;
 import java.lang.invoke.MethodHandles;
 import java.util.function.BiFunction;
@@ -134,6 +135,20 @@ public final class ContextPropagationOperator {
     return context.getOrDefault(TRACE_CONTEXT_KEY, defaultTraceContext);
   }
 
+  /**
+   * Gets Trace {@link Context} from Reactor {@link reactor.util.context.ContextView}.
+   *
+   * @param contextView Reactor's context to get trace context from.
+   * @param defaultTraceContext Default value to be returned if no trace context is found on Reactor
+   *     context.
+   * @return Trace context or default value.
+   */
+  @NoMuzzle
+  public static Context getOpenTelemetryContextFromContextView(
+      reactor.util.context.ContextView contextView, Context defaultTraceContext) {
+    return contextView.getOrDefault(TRACE_CONTEXT_KEY, defaultTraceContext);
+  }
+
   ContextPropagationOperator(boolean captureExperimentalSpanAttributes) {
     this.asyncOperationEndStrategy =
         ReactorAsyncOperationEndStrategy.builder()

instrumentation/reactor/reactor-3.1/library/src/test/java/io/opentelemetry/instrumentation/reactor/v3_1/HooksTest.java

@@ -62,7 +62,7 @@ class HooksTest {
     Disposable disposable =
         Mono.defer(
                 () ->
-                    Mono.fromCallable(callable).publishOn(Schedulers.elastic()).flatMap(Mono::just))
+                    Mono.fromCallable(callable).publishOn(Schedulers.single()).flatMap(Mono::just))
             .subscribeOn(Schedulers.single())
             .subscribe();
 

instrumentation/reactor/reactor-3.1/library/src/test/java/io/opentelemetry/instrumentation/reactor/v3_1/ReactorCoreTest.java

@@ -160,13 +160,18 @@ class ReactorCoreTest extends AbstractReactorCoreTest {
 
   @Test
   void nestedNonBlocking() {
+    boolean testLatestDeps = Boolean.getBoolean("testLatestDeps");
     int result =
         testing.runWithSpan(
             "parent",
             () ->
                 Mono.defer(
                         () -> {
-                          Span.current().setAttribute("middle", "foo");
+                          // earliest tested and latest version behave differently
+                          // in latest dep test current span is "parent" not "middle"
+                          if (!testLatestDeps) {
+                            Span.current().setAttribute("middle", "foo");
+                          }
                           return Mono.fromCallable(
                                   () -> {
                                     Span.current().setAttribute("inner", "bar");
@@ -183,10 +188,12 @@ class ReactorCoreTest extends AbstractReactorCoreTest {
         trace ->
             trace.hasSpansSatisfyingExactly(
                 span -> span.hasName("parent").hasNoParent(),
-                span ->
-                    span.hasName("middle")
-                        .hasParent(trace.getSpan(0))
-                        .hasAttributes(attributeEntry("middle", "foo")),
+                span -> {
+                  span.hasName("middle").hasParent(trace.getSpan(0));
+                  if (!testLatestDeps) {
+                    span.hasAttributes(attributeEntry("middle", "foo"));
+                  }
+                },
                 span ->
                     span.hasName("inner")
                         .hasParent(trace.getSpan(1))

instrumentation/reactor/reactor-3.4/javaagent/build.gradle.kts

@@ -0,0 +1,34 @@
+plugins {
+  id("otel.javaagent-instrumentation")
+}
+
+muzzle {
+  pass {
+    group.set("io.projectreactor")
+    module.set("reactor-core")
+    versions.set("[3.4.0,)")
+    extraDependency("io.opentelemetry:opentelemetry-api:1.0.0")
+    assertInverse.set(true)
+    excludeInstrumentationName("opentelemetry-api")
+  }
+}
+
+dependencies {
+  library("io.projectreactor:reactor-core:3.4.0")
+  implementation(project(":instrumentation:reactor:reactor-3.1:library"))
+
+  implementation(project(":instrumentation:opentelemetry-api:opentelemetry-api-1.0:javaagent"))
+
+  compileOnly(project(":javaagent-tooling"))
+  compileOnly(project(":instrumentation-annotations-support"))
+  compileOnly(project(":opentelemetry-api-shaded-for-instrumenting", configuration = "shadow"))
+
+  testInstrumentation(project(":instrumentation:reactor:reactor-3.1:javaagent"))
+  testInstrumentation(project(":instrumentation:opentelemetry-extension-annotations-1.0:javaagent"))
+
+  testLibrary("io.projectreactor:reactor-test:3.1.0.RELEASE")
+  testImplementation(project(":instrumentation-annotations-support-testing"))
+  testImplementation(project(":instrumentation:reactor:reactor-3.1:testing"))
+  testImplementation(project(":instrumentation-annotations"))
+  testImplementation("io.opentelemetry:opentelemetry-extension-annotations")
+}

instrumentation/reactor/reactor-3.4/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/reactor/v3_4/operator/ContextPropagationOperator34Instrumentation.java

@@ -0,0 +1,66 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.reactor.v3_4.operator;
+
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.isPublic;
+import static net.bytebuddy.matcher.ElementMatchers.isStatic;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.returns;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+
+import application.io.opentelemetry.context.Context;
+import io.opentelemetry.instrumentation.reactor.v3_1.ContextPropagationOperator;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import io.opentelemetry.javaagent.instrumentation.opentelemetryapi.context.AgentContextStorage;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+public class ContextPropagationOperator34Instrumentation implements TypeInstrumentation {
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return named(
+        "application.io.opentelemetry.instrumentation.reactor.v3_1.ContextPropagationOperator");
+  }
+
+  @Override
+  public void transform(TypeTransformer transformer) {
+    transformer.applyAdviceToMethod(
+        isMethod()
+            .and(isPublic())
+            .and(isStatic())
+            .and(named("getOpenTelemetryContextFromContextView"))
+            .and(takesArgument(0, named("reactor.util.context.ContextView")))
+            .and(takesArgument(1, named("application.io.opentelemetry.context.Context")))
+            .and(returns(named("application.io.opentelemetry.context.Context"))),
+        ContextPropagationOperator34Instrumentation.class.getName() + "$GetContextViewAdvice");
+  }
+
+  @SuppressWarnings("unused")
+  public static class GetContextViewAdvice {
+    @Advice.OnMethodEnter(skipOn = Advice.OnDefaultValue.class)
+    public static boolean methodEnter() {
+      return false;
+    }
+
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void methodExit(
+        @Advice.Argument(0) reactor.util.context.ContextView reactorContext,
+        @Advice.Argument(1) Context defaultContext,
+        @Advice.Return(readOnly = false) Context applicationContext) {
+
+      io.opentelemetry.context.Context agentContext =
+          ContextPropagationOperator.getOpenTelemetryContextFromContextView(reactorContext, null);
+      if (agentContext == null) {
+        applicationContext = defaultContext;
+      } else {
+        applicationContext = AgentContextStorage.toApplicationContext(agentContext);
+      }
+    }
+  }
+}

instrumentation/reactor/reactor-3.4/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/reactor/v3_4/operator/ContextPropagationOperator34InstrumentationModule.java

@@ -0,0 +1,34 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.reactor.v3_4.operator;
+
+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;
+
+@AutoService(InstrumentationModule.class)
+public class ContextPropagationOperator34InstrumentationModule extends InstrumentationModule {
+
+  public ContextPropagationOperator34InstrumentationModule() {
+    super("reactor", "reactor-3.4", "reactor-context-propagation-operator");
+  }
+
+  @Override
+  public ElementMatcher.Junction<ClassLoader> classLoaderMatcher() {
+    return hasClassesNamed(
+        "application.io.opentelemetry.context.Context", "reactor.util.context.ContextView");
+  }
+
+  @Override
+  public List<TypeInstrumentation> typeInstrumentations() {
+    return singletonList(new ContextPropagationOperator34Instrumentation());
+  }
+}

instrumentation/reactor/reactor-3.4/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/reactor/v3_4/ContextPropagationOperator34InstrumentationTest.java

@@ -0,0 +1,55 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.reactor.v3_4;
+
+import static io.opentelemetry.sdk.testing.assertj.OpenTelemetryAssertions.attributeEntry;
+import static org.assertj.core.api.Assertions.assertThat;
+
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.SpanKind;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.instrumentation.reactor.v3_1.ContextPropagationOperator;
+import io.opentelemetry.instrumentation.testing.junit.AgentInstrumentationExtension;
+import io.opentelemetry.instrumentation.testing.junit.InstrumentationExtension;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+class ContextPropagationOperator34InstrumentationTest {
+
+  @RegisterExtension
+  static final InstrumentationExtension testing = AgentInstrumentationExtension.create();
+
+  @Test
+  void storeAndGetContext() {
+    reactor.util.context.Context reactorContext = reactor.util.context.Context.empty();
+    testing.runWithSpan(
+        "parent",
+        () -> {
+          reactor.util.context.Context newReactorContext =
+              ContextPropagationOperator.storeOpenTelemetryContext(
+                  reactorContext, Context.current());
+          Context otelContext =
+              ContextPropagationOperator.getOpenTelemetryContextFromContextView(
+                  newReactorContext, null);
+          assertThat(otelContext).isNotNull();
+          Span.fromContext(otelContext).setAttribute("foo", "bar");
+          Context otelContext2 =
+              ContextPropagationOperator.getOpenTelemetryContext(newReactorContext, null);
+          assertThat(otelContext2).isNotNull();
+          Span.fromContext(otelContext2).setAttribute("foo2", "bar2");
+        });
+
+    testing.waitAndAssertTraces(
+        trace ->
+            trace.hasSpansSatisfyingExactly(
+                span ->
+                    span.hasName("parent")
+                        .hasKind(SpanKind.INTERNAL)
+                        .hasNoParent()
+                        .hasAttributes(
+                            attributeEntry("foo", "bar"), attributeEntry("foo2", "bar2"))));
+  }
+}

settings.gradle.kts

@@ -465,6 +465,7 @@ include(":instrumentation:ratpack:ratpack-1.7:library")
 include(":instrumentation:reactor:reactor-3.1:javaagent")
 include(":instrumentation:reactor:reactor-3.1:library")
 include(":instrumentation:reactor:reactor-3.1:testing")
+include(":instrumentation:reactor:reactor-3.4:javaagent")
 include(":instrumentation:reactor:reactor-kafka-1.0:javaagent")
 include(":instrumentation:reactor:reactor-kafka-1.0:testing")
 include(":instrumentation:reactor:reactor-netty:reactor-netty-0.9:javaagent")