Clear context before flux retry (#8456)

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

@@ -32,6 +32,8 @@
  * 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 static final Class<?> fluxRetryWhenSubscriberClass = getFluxRetryWhenSubscriberClass();
  private final io.opentelemetry.context.Context traceContext;
  private final Subscriber<? super T> subscriber;
  private final Context context;
@@ -64,7 +66,15 @@ public void onNext(T o) {
 
  @Override
  public void onError(Throwable throwable) {
- withActiveSpan(() -> subscriber.onError(throwable));
+ if (!hasContextToPropagate
+ && (fluxRetrySubscriberClass == subscriber.getClass()
+ || fluxRetryWhenSubscriberClass == 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 +88,32 @@ 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;
+ }
+ }
+
+ private static Class<?> getFluxRetryWhenSubscriberClass() {
+ try {
+ return Class.forName("reactor.core.publisher.FluxRetryWhen$RetryWhenMainSubscriber");
+ } 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;
@@ -20,12 +21,17 @@
 import io.opentelemetry.instrumentation.testing.junit.LibraryInstrumentationExtension;
 import java.lang.invoke.MethodHandle;
 import java.lang.invoke.MethodHandles;
+import java.lang.reflect.Method;
 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;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+import org.reactivestreams.Publisher;
 import reactor.core.Scannable;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
@@ -406,6 +412,133 @@ void doesNotOverrideInnerCurrentSpansWithThereIsOuterCurrent() {
  .hasAttributes(attributeEntry("onNext", true))));
  }
 
+ @ParameterizedTest
+ @ValueSource(strings = {"retry", "retryWhen"})
+ void doesNotLeakContextOnRetry(String retryKind) {
+ // 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++) {
+ int index = i;
+ testing.runWithSpan(
+ "produce " + (beforeRetry.get() ? "before" : "after") + " retry " + i,
+ () -> sink.next(index));
+ }
+ });
+
+ Flux<Integer> flux =
+ Flux.defer(() -> publish.delaySubscription(Duration.ofMillis(1)))
+ .doOnNext(message -> testing.runWithSpan("process " + message, () -> {}))
+ .handle(
+ (message, sink) -> {
+ if (message == 1 && beforeRetry.compareAndSet(true, false)) {
+ sink.error(new RuntimeException("Message has error"));
+ } else {
+ sink.next(message);
+ }
+ });
+
+ switch (retryKind) {
+ case "retry":
+ flux = flux.retry();
+ break;
+ case "retryWhen":
+ flux = retryWhen(flux);
+ break;
+ default:
+ throw new IllegalStateException("Unsupported retry kind " + retryKind);
+ }
+
+ flux.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))));
+ }
+
+ @Test
+ void retryWithParentSpan() {
+ AtomicBoolean beforeRetry = new AtomicBoolean(true);
+ Flux<Integer> publish =
+ Flux.create(
+ sink ->
+ testing.runWithSpan(
+ "produce " + (beforeRetry.get() ? "before" : "after") + " retry",
+ () -> sink.next(0)));
+
+ Flux<Object> flux =
+ Flux.defer(() -> publish.delaySubscription(Duration.ofMillis(1)))
+ .doOnNext(message -> testing.runWithSpan("process", () -> {}))
+ .handle(
+ (message, sink) -> {
+ if (beforeRetry.compareAndSet(true, false)) {
+ sink.error(new RuntimeException("Message has error"));
+ } else {
+ sink.next(message);
+ }
+ })
+ .retry();
+
+ testing.runWithSpan("parent", () -> flux.subscribe());
+
+ testing.waitAndAssertTraces(
+ trace ->
+ trace.hasSpansSatisfyingExactly(
+ span -> span.hasName("parent").hasNoParent(),
+ span -> span.hasName("produce before retry").hasParent(trace.getSpan(0)),
+ span -> span.hasName("process").hasParent(trace.getSpan(0)),
+ span -> span.hasName("produce after retry").hasParent(trace.getSpan(0)),
+ span -> span.hasName("process").hasParent(trace.getSpan(0))));
+ }
+
+ @SuppressWarnings("unchecked")
+ private static <T> Flux<T> retryWhen(Flux<T> flux) {
+ try {
+ Method method = Flux.class.getMethod("retryWhen", Function.class);
+ Function<Flux<Throwable>, ? extends Publisher<?>> function =
+ err -> Flux.create(sink -> sink.next(-1));
+ return (Flux<T>) method.invoke(flux, function);
+ } catch (NoSuchMethodException exception) {
+ // ignore
+ } catch (Exception exception) {
+ throw new IllegalStateException(exception);
+ }
+
+ try {
+ Class<?> retryClass = Class.forName("reactor.util.retry.Retry");
+ Method retryWhenMethod = Flux.class.getMethod("retryWhen", retryClass);
+ Method retrySpecMethod = retryClass.getMethod("indefinitely");
+ return (Flux<T>) retryWhenMethod.invoke(flux, retrySpecMethod.invoke(retryClass));
+ } catch (ClassNotFoundException | NoSuchMethodException exception) {
+ // ignore
+ } catch (Exception exception) {
+ throw new IllegalStateException(exception);
+ }
+ throw new IllegalStateException("Could not find retryWhen method");
+ }
+
  @SuppressWarnings("unchecked")
  private <T> Mono<T> monoSpan(Mono<T> mono, String spanName) {