Clear context before flux retry

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

@@ -32,6 +32,7 @@
  * https://github.com/opentracing-contrib/java-reactor/blob/master/src/main/java/io/opentracing/contrib/reactor/TracedSubscriber.java
  */
 public class TracingSubscriber<T> implements CoreSubscriber<T> {
+  private static final Class<?> fluxRetrySubscriberClass = getFluxRetrySubscriberClass();
   private final io.opentelemetry.context.Context traceContext;
   private final Subscriber<? super T> subscriber;
   private final Context context;
@@ -64,7 +65,13 @@ public void onNext(T o) {
 
   @Override
   public void onError(Throwable throwable) {
-    withActiveSpan(() -> subscriber.onError(throwable));
+    if (!hasContextToPropagate && fluxRetrySubscriberClass == subscriber.getClass()) {
+      // clear context for retry to avoid having retried operations run with currently active
+      // context as parent context
+      withActiveSpan(io.opentelemetry.context.Context.root(), () -> subscriber.onError(throwable));
+    } else {
+      withActiveSpan(() -> subscriber.onError(throwable));
+    }
   }
 
   @Override
@@ -78,12 +85,24 @@ public Context currentContext() {
   }
 
   private void withActiveSpan(Runnable runnable) {
-    if (hasContextToPropagate) {
-      try (Scope ignored = traceContext.makeCurrent()) {
+    withActiveSpan(hasContextToPropagate ? traceContext : null, runnable);
+  }
+
+  private static void withActiveSpan(io.opentelemetry.context.Context context, Runnable runnable) {
+    if (context != null) {
+      try (Scope ignored = context.makeCurrent()) {
         runnable.run();
       }
     } else {
       runnable.run();
     }
   }
+
+  private static Class<?> getFluxRetrySubscriberClass() {
+    try {
+      return Class.forName("reactor.core.publisher.FluxRetry$RetrySubscriber");
+    } catch (ClassNotFoundException exception) {
+      return null;
+    }
+  }
 }

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

@@ -5,6 +5,7 @@
 
 package io.opentelemetry.instrumentation.reactor.v3_1;
 
+import static io.opentelemetry.instrumentation.testing.util.TelemetryDataUtil.orderByRootSpanName;
 import static io.opentelemetry.sdk.testing.assertj.OpenTelemetryAssertions.attributeEntry;
 import static java.lang.invoke.MethodType.methodType;
 import static org.assertj.core.api.Assertions.assertThat;
@@ -21,6 +22,7 @@
 import java.lang.invoke.MethodHandle;
 import java.lang.invoke.MethodHandles;
 import java.time.Duration;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Function;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.BeforeAll;
@@ -406,6 +408,68 @@ void doesNotOverrideInnerCurrentSpansWithThereIsOuterCurrent() {
                         .hasAttributes(attributeEntry("onNext", true))));
   }
 
+  @Test
+  void doesNotLeakContextOnRetry() {
+    // retry calls subscribe again from onError where we have active context, check that this
+    // context is not used as parent for retried operations
+    AtomicBoolean beforeRetry = new AtomicBoolean(true);
+    Flux<Integer> publish =
+        Flux.create(
+            sink -> {
+              for (int i = 0; i < 2; i++) {
+                Span s =
+                    tracer
+                        .spanBuilder(
+                            "produce " + (beforeRetry.get() ? "before" : "after") + " retry " + i)
+                        .startSpan();
+                try (Scope scope = s.makeCurrent()) {
+                  sink.next(i);
+                } finally {
+                  s.end();
+                }
+              }
+            });
+
+    Flux.defer(() -> publish.delaySubscription(Duration.ofMillis(1)))
+        .doOnNext(
+            message -> {
+              tracer.spanBuilder("process " + message).startSpan().end();
+            })
+        .handle(
+            (message, sink) -> {
+              if (message == 1 && beforeRetry.compareAndSet(true, false)) {
+                sink.error(new RuntimeException("Message has error"));
+              } else {
+                sink.next(message);
+              }
+            })
+        .retry()
+        .subscribe();
+
+    testing.waitAndAssertSortedTraces(
+        orderByRootSpanName(
+            "produce before retry 0",
+            "produce before retry 1",
+            "produce after retry 0",
+            "produce after retry 1"),
+        trace ->
+            trace.hasSpansSatisfyingExactly(
+                span -> span.hasName("produce before retry 0").hasNoParent(),
+                span -> span.hasName("process 0").hasParent(trace.getSpan(0))),
+        trace ->
+            trace.hasSpansSatisfyingExactly(
+                span -> span.hasName("produce before retry 1").hasNoParent(),
+                span -> span.hasName("process 1").hasParent(trace.getSpan(0))),
+        trace ->
+            trace.hasSpansSatisfyingExactly(
+                span -> span.hasName("produce after retry 0").hasNoParent(),
+                span -> span.hasName("process 0").hasParent(trace.getSpan(0))),
+        trace ->
+            trace.hasSpansSatisfyingExactly(
+                span -> span.hasName("produce after retry 1").hasNoParent(),
+                span -> span.hasName("process 1").hasParent(trace.getSpan(0))));
+  }
+
   @SuppressWarnings("unchecked")
   private <T> Mono<T> monoSpan(Mono<T> mono, String spanName) {