Use new reactor contextWrite when available (from reactor 3.4.0)

dependencyManagement/build.gradle.kts

@@ -33,7 +33,7 @@ val DEPENDENCY_BOMS = listOf(
   "io.opentelemetry:opentelemetry-bom-alpha:1.22.0-alpha",
   "org.junit:junit-bom:5.9.2",
   "org.testcontainers:testcontainers-bom:1.17.6",
-  "org.spockframework:spock-bom:2.3-groovy-4.0"
+  "org.spockframework:spock-bom:2.4-M1-groovy-4.0"
 )
 
 val autoServiceVersion = "1.0.1"

docs/supported-libraries.md

@@ -118,7 +118,7 @@ These are the supported libraries and frameworks:
 | [Spring RestTemplate](https://docs.spring.io/spring-framework/docs/current/javadoc-api/org/springframework/web/client/package-summary.html) | 3.1+                          | [opentelemetry-spring-web-3.1](../instrumentation/spring/spring-web/spring-web-3.1/library)                                                                                                                                                                                                                                                                                             | [HTTP Client Spans], [HTTP Client Metrics]                                             |
 | [Spring Web MVC](https://docs.spring.io/spring/docs/current/javadoc-api/org/springframework/web/servlet/mvc/package-summary.html)           | 3.1+                          | [opentelemetry-spring-webmvc-5.3](../instrumentation/spring/spring-webmvc/spring-webmvc-5.3/library)                                                                                                                                                                                                                                                                                    | [HTTP Server Spans], [HTTP Server Metrics]                                             |
 | [Spring Web Services](https://spring.io/projects/spring-ws)                                                                                 | 2.0+                          | N/A                                                                                                                                                                                                                                                                                                                                                                                     | none                                                                                   |
-| [Spring WebFlux](https://docs.spring.io/spring/docs/current/javadoc-api/org/springframework/web/reactive/package-summary.html)              | 5.0+ (not including 6.0+ yet) | [opentelemetry-spring-webflux-5.0](../instrumentation/spring/spring-webflux-5.0/library)                                                                                                                                                                                                                                                                                                | [HTTP Client Spans], [HTTP Client Metrics],                                            |
+| [Spring WebFlux](https://docs.spring.io/spring/docs/current/javadoc-api/org/springframework/web/reactive/package-summary.html)              | 5.0+                          | [opentelemetry-spring-webflux-5.0](../instrumentation/spring/spring-webflux-5.0/library)                                                                                                                                                                                                                                                                                                | [HTTP Client Spans], [HTTP Client Metrics],                                            |
 | [Spymemcached](https://github.com/couchbase/spymemcached)                                                                                   | 2.12+                         | N/A                                                                                                                                                                                                                                                                                                                                                                                     | [Database Client Spans]                                                                |
 | [Tomcat JDBC Pool](https://tomcat.apache.org/tomcat-7.0-doc/jdbc-pool.html)                                                                 | 8.5.0+                        | N/A                                                                                                                                                                                                                                                                                                                                                                                     | [Database Pool Metrics]                                                                |
 | [Twilio](https://github.com/twilio/twilio-java)                                                                                             | 6.6+ (not including 8.x yet)  | N/A                                                                                                                                                                                                                                                                                                                                                                                     | none                                                                                   |

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

@@ -6,7 +6,7 @@ muzzle {
   pass {
     group.set("io.projectreactor")
     module.set("reactor-core")
-    versions.set("[3.1.0.RELEASE,3.5.0)")
+    versions.set("[3.1.0.RELEASE,)")
     extraDependency("io.opentelemetry:opentelemetry-api:1.0.0")
     assertInverse.set(true)
     excludeInstrumentationName("opentelemetry-api")

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

@@ -6,6 +6,7 @@
 package io.opentelemetry.javaagent.instrumentation.reactor;
 
 import static io.opentelemetry.sdk.testing.assertj.OpenTelemetryAssertions.attributeEntry;
+import static java.lang.invoke.MethodType.methodType;
 import static org.assertj.core.api.Assertions.assertThat;
 
 import io.opentelemetry.api.trace.Span;
@@ -14,7 +15,11 @@
 import io.opentelemetry.instrumentation.reactor.ContextPropagationOperator;
 import io.opentelemetry.instrumentation.testing.junit.AgentInstrumentationExtension;
 import io.opentelemetry.instrumentation.testing.junit.InstrumentationExtension;
+import java.lang.invoke.MethodHandle;
+import java.lang.invoke.MethodHandles;
 import java.time.Duration;
+import java.util.function.Function;
+import javax.annotation.Nullable;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 import reactor.core.publisher.Flux;
@@ -26,6 +31,25 @@ class ContextPropagationOperatorInstrumentationTest {
   @RegisterExtension
   static final InstrumentationExtension testing = AgentInstrumentationExtension.create();
 
+  @Nullable
+  private static final MethodHandle MONO_CONTEXT_WRITE_METHOD = getContextWriteMethod(Mono.class);
+
+  @Nullable
+  private static MethodHandle getContextWriteMethod(Class<?> type) {
+    MethodHandles.Lookup lookup = MethodHandles.publicLookup();
+    try {
+      return lookup.findVirtual(type, "contextWrite", methodType(type, Function.class));
+    } catch (NoSuchMethodException | IllegalAccessException e) {
+      // ignore
+    }
+    try {
+      return lookup.findVirtual(type, "subscriberContext", methodType(type, Function.class));
+    } catch (NoSuchMethodException | IllegalAccessException e) {
+      // ignore
+    }
+    return null;
+  }
+
   @Test
   void storeAndGetContext() {
     reactor.util.context.Context reactorContext = reactor.util.context.Context.empty();
@@ -136,18 +160,25 @@ void storeContextForcesItToBecomeCurrent() {
               Span span =
                   testing.getOpenTelemetry().getTracer("test").spanBuilder("parent").startSpan();
 
-              return Mono.delay(Duration.ofMillis(1))
-                  .flatMap(
-                      unused -> {
-                        // usual trick to force this to run under new TracingSubscriber with context
-                        // written in the next call
-                        return new ExtensionAnnotationsTracedWithSpan().mono(Mono.just("Value"));
-                      })
-                  .subscriberContext(
-                      ctx ->
-                          ContextPropagationOperator.storeOpenTelemetryContext(
-                              ctx, Context.current().with(span)))
-                  .doFinally(unused -> span.end());
+              Mono<String> interim =
+                  Mono.delay(Duration.ofMillis(1))
+                      .flatMap(
+                          unused -> {
+                            // usual trick to force this to run under new TracingSubscriber with
+                            // context
+                            // written in the next call
+                            return new ExtensionAnnotationsTracedWithSpan()
+                                .mono(Mono.just("Value"));
+                          });
+              try {
+                interim =
+                    (Mono<String>)
+                        MONO_CONTEXT_WRITE_METHOD.invoke(
+                            interim, new StoreOpenTelemetryContext(span));
+              } catch (Throwable e) {
+                throw new RuntimeException(e);
+              }
+              return interim.doFinally(unused -> span.end());
             });
 
     StepVerifier.create(result).expectNext("Value").verifyComplete();
@@ -161,4 +192,20 @@ void storeContextForcesItToBecomeCurrent() {
                         .hasKind(SpanKind.INTERNAL)
                         .hasParent(trace.getSpan(0))));
   }
+
+  private static class StoreOpenTelemetryContext
+      implements Function<reactor.util.context.Context, reactor.util.context.Context> {
+
+    private final Span span;
+
+    private StoreOpenTelemetryContext(Span span) {
+      this.span = span;
+    }
+
+    @Override
+    public reactor.util.context.Context apply(reactor.util.context.Context context) {
+      return ContextPropagationOperator.storeOpenTelemetryContext(
+          context, Context.current().with(span));
+    }
+  }
 }

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

@@ -22,9 +22,13 @@
 
 package io.opentelemetry.instrumentation.reactor;
 
+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 java.lang.invoke.MethodHandle;
+import java.lang.invoke.MethodHandles;
 import java.util.function.BiFunction;
 import java.util.function.Function;
 import javax.annotation.Nullable;
@@ -42,6 +46,28 @@ public final class ContextPropagationOperator {
 
   private static final Object VALUE = new Object();
 
+  @Nullable
+  private static final MethodHandle MONO_CONTEXT_WRITE_METHOD = getContextWriteMethod(Mono.class);
+
+  @Nullable
+  private static final MethodHandle FLUX_CONTEXT_WRITE_METHOD = getContextWriteMethod(Flux.class);
+
+  @Nullable
+  private static MethodHandle getContextWriteMethod(Class<?> type) {
+    MethodHandles.Lookup lookup = MethodHandles.publicLookup();
+    try {
+      return lookup.findVirtual(type, "contextWrite", methodType(type, Function.class));
+    } catch (NoSuchMethodException | IllegalAccessException e) {
+      // ignore
+    }
+    try {
+      return lookup.findVirtual(type, "subscriberContext", methodType(type, Function.class));
+    } catch (NoSuchMethodException | IllegalAccessException e) {
+      // ignore
+    }
+    return null;
+  }
+
   public static ContextPropagationOperator create() {
     return builder().build();
   }
@@ -133,29 +159,65 @@ public void resetOnEachOperator() {
   }
 
   /** Forces Mono to run in traceContext scope. */
+  @SuppressWarnings("unchecked")
   public static <T> Mono<T> runWithContext(Mono<T> publisher, Context tracingContext) {
-    if (!enabled) {
+    if (!enabled || MONO_CONTEXT_WRITE_METHOD == null) {
       return publisher;
     }
 
     // this hack forces 'publisher' to run in the onNext callback of `TracingSubscriber`
     // (created for this publisher) and with current() span that refers to span created here
     // without the hack, publisher runs in the onAssembly stage, before traceContext is made current
-    return ScalarPropagatingMono.create(publisher)
-        .subscriberContext(ctx -> storeOpenTelemetryContext(ctx, tracingContext));
+    try {
+      return (Mono<T>)
+          MONO_CONTEXT_WRITE_METHOD.invoke(
+              ScalarPropagatingMono.create(publisher),
+              new StoreOpenTelemetryContext(tracingContext));
+    } catch (Throwable t) {
+      // rethrowing without any wrapping to avoid any change to the underlying application behavior
+      throw sneakyThrow(t);
+    }
   }
 
   /** Forces Flux to run in traceContext scope. */
+  @SuppressWarnings("unchecked")
   public static <T> Flux<T> runWithContext(Flux<T> publisher, Context tracingContext) {
-    if (!enabled) {
+    if (!enabled || FLUX_CONTEXT_WRITE_METHOD == null) {
       return publisher;
     }
 
     // this hack forces 'publisher' to run in the onNext callback of `TracingSubscriber`
     // (created for this publisher) and with current() span that refers to span created here
     // without the hack, publisher runs in the onAssembly stage, before traceContext is made current
-    return ScalarPropagatingFlux.create(publisher)
-        .subscriberContext(ctx -> storeOpenTelemetryContext(ctx, tracingContext));
+    try {
+      return (Flux<T>)
+          FLUX_CONTEXT_WRITE_METHOD.invoke(
+              ScalarPropagatingFlux.create(publisher),
+              new StoreOpenTelemetryContext(tracingContext));
+    } catch (Throwable t) {
+      // rethrowing without any wrapping to avoid any change to the underlying application behavior
+      throw sneakyThrow(t);
+    }
+  }
+
+  @SuppressWarnings({"TypeParameterUnusedInFormals", "unchecked"})
+  private static <T extends Throwable> T sneakyThrow(Throwable t) throws T {
+    throw (T) t;
+  }
+
+  private static class StoreOpenTelemetryContext
+      implements Function<reactor.util.context.Context, reactor.util.context.Context> {
+
+    private final Context tracingContext;
+
+    private StoreOpenTelemetryContext(Context tracingContext) {
+      this.tracingContext = tracingContext;
+    }
+
+    @Override
+    public reactor.util.context.Context apply(reactor.util.context.Context context) {
+      return storeOpenTelemetryContext(context, tracingContext);
+    }
   }
 
   public static class Lifter<T>