End span on cancellation of subscription to reactive publishers

instrumentation/lettuce/lettuce-5.1/library/src/test/groovy/io/opentelemetry/instrumentation/lettuce/v5_1/LettuceReactiveClientTest.groovy

@@ -9,8 +9,12 @@ import io.lettuce.core.RedisClient
 import io.lettuce.core.resource.ClientResources
 import io.opentelemetry.instrumentation.reactor.TracingOperator
 import io.opentelemetry.instrumentation.test.LibraryTestTrait
+import spock.lang.Shared
 
 class LettuceReactiveClientTest extends AbstractLettuceReactiveClientTest implements LibraryTestTrait {
+  @Shared
+  TracingOperator tracingOperator = TracingOperator.create()
+
   @Override
   RedisClient createClient(String uri) {
     return RedisClient.create(
@@ -21,10 +25,10 @@ class LettuceReactiveClientTest extends AbstractLettuceReactiveClientTest implem
   }
 
   def setupSpec() {
-    TracingOperator.registerOnEachOperator()
+    tracingOperator.registerOnEachOperator()
   }
 
   def cleanupSpec() {
-    TracingOperator.resetOnEachOperator()
+    tracingOperator.resetOnEachOperator()
   }
 }

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

@@ -9,6 +9,11 @@ muzzle {
   }
 }
 
+tasks.withType(Test).configureEach {
+  // TODO run tests both with and without experimental span attributes
+  jvmArgs "-Dotel.instrumentation.reactor.experimental-span-attributes=true"
+}
+
 dependencies {
   implementation project(':instrumentation:reactor-3.1:library')
 

instrumentation/reactor-3.1/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/reactor/HooksInstrumentation.java

@@ -8,6 +8,7 @@
 import static net.bytebuddy.matcher.ElementMatchers.isTypeInitializer;
 import static net.bytebuddy.matcher.ElementMatchers.named;
 
+import io.opentelemetry.instrumentation.api.config.Config;
 import io.opentelemetry.instrumentation.reactor.TracingOperator;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
@@ -31,7 +32,13 @@ public void transform(TypeTransformer transformer) {
   public static class ResetOnEachOperatorAdvice {
     @Advice.OnMethodExit(suppress = Throwable.class)
     public static void postStaticInitializer() {
-      TracingOperator.registerOnEachOperator();
+      TracingOperator.newBuilder()
+          .setCaptureExperimentalSpanAttributes(
+              Config.get()
+                  .getBooleanProperty(
+                      "otel.instrumentation.reactor.experimental-span-attributes", false))
+          .build()
+          .registerOnEachOperator();
     }
   }
 }

instrumentation/reactor-3.1/javaagent/src/test/groovy/ReactorWithSpanInstrumentationTest.groovy

@@ -130,6 +130,35 @@ class ReactorWithSpanInstrumentationTest extends AgentInstrumentationSpecificati
     }
   }
 
+  def "should capture span for canceled Mono"() {
+    setup:
+    def source = UnicastProcessor.<String>create()
+    def mono = source.singleOrEmpty()
+    def result = new TracedWithSpan()
+      .mono(mono)
+    def verifier = StepVerifier.create(result)
+      .expectSubscription()
+
+    expect:
+    Thread.sleep(500) // sleep a bit just to make sure no span is captured
+    assertTraces(0) {}
+
+    verifier.thenCancel().verify()
+
+    assertTraces(1) {
+      trace(0, 1) {
+        span(0) {
+          name "TracedWithSpan.mono"
+          kind SpanKind.INTERNAL
+          hasNoParent()
+          attributes {
+            "reactor.canceled" true
+          }
+        }
+      }
+    }
+  }
+
   def "should capture span for already completed Flux"() {
     setup:
     def source = Flux.just("Value")
@@ -242,4 +271,35 @@ class ReactorWithSpanInstrumentationTest extends AgentInstrumentationSpecificati
       }
     }
   }
+
+  def "should capture span for canceled Flux"() {
+    setup:
+    def error = new IllegalArgumentException("Boom")
+    def source = UnicastProcessor.<String>create()
+    def result = new TracedWithSpan()
+      .flux(source)
+    def verifier = StepVerifier.create(result)
+      .expectSubscription()
+
+    expect:
+    Thread.sleep(500) // sleep a bit just to make sure no span is captured
+    assertTraces(0) {}
+
+    source.onError(error)
+
+    verifier.thenCancel().verify()
+
+    assertTraces(1) {
+      trace(0, 1) {
+        span(0) {
+          name "TracedWithSpan.flux"
+          kind SpanKind.INTERNAL
+          hasNoParent()
+          attributes {
+            "reactor.canceled" true
+          }
+        }
+      }
+    }
+  }
 }

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

@@ -5,6 +5,8 @@
 
 package io.opentelemetry.instrumentation.reactor;
 
+import io.opentelemetry.api.common.AttributeKey;
+import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.context.Context;
 import io.opentelemetry.instrumentation.api.tracer.BaseTracer;
 import io.opentelemetry.instrumentation.api.tracer.async.AsyncSpanEndStrategy;
@@ -14,8 +16,23 @@
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
-public enum ReactorAsyncSpanEndStrategy implements AsyncSpanEndStrategy {
-  INSTANCE;
+public final class ReactorAsyncSpanEndStrategy implements AsyncSpanEndStrategy {
+  private static final AttributeKey<Boolean> CANCELED_ATTRIBUTE_KEY =
+      AttributeKey.booleanKey("reactor.canceled");
+
+  public static ReactorAsyncSpanEndStrategy create() {
+    return newBuilder().build();
+  }
+
+  public static ReactorAsyncSpanEndStrategyBuilder newBuilder() {
+    return new ReactorAsyncSpanEndStrategyBuilder();
+  }
+
+  private final boolean captureExperimentalSpanAttributes;
+
+  ReactorAsyncSpanEndStrategy(boolean captureExperimentalSpanAttributes) {
+    this.captureExperimentalSpanAttributes = captureExperimentalSpanAttributes;
+  }
 
   @Override
   public boolean supports(Class<?> returnType) {
@@ -29,10 +46,14 @@ public Object end(BaseTracer tracer, Context context, Object returnValue) {
         new EndOnFirstNotificationConsumer(tracer, context);
     if (returnValue instanceof Mono) {
       Mono<?> mono = (Mono<?>) returnValue;
-      return mono.doOnError(notificationConsumer).doOnSuccess(notificationConsumer::onSuccess);
+      return mono.doOnError(notificationConsumer)
+          .doOnSuccess(notificationConsumer::onSuccess)
+          .doOnCancel(notificationConsumer::onCancel);
     } else {
       Flux<?> flux = Flux.from((Publisher<?>) returnValue);
-      return flux.doOnError(notificationConsumer).doOnComplete(notificationConsumer);
+      return flux.doOnError(notificationConsumer)
+          .doOnComplete(notificationConsumer)
+          .doOnCancel(notificationConsumer::onCancel);
     }
   }
 
@@ -41,7 +62,7 @@ public Object end(BaseTracer tracer, Context context, Object returnValue) {
    * OnError notifications are received. Multiple notifications can happen anytime multiple
    * subscribers subscribe to the same publisher.
    */
-  private static final class EndOnFirstNotificationConsumer extends AtomicBoolean
+  private final class EndOnFirstNotificationConsumer extends AtomicBoolean
       implements Runnable, Consumer<Throwable> {
 
     private final BaseTracer tracer;
@@ -57,6 +78,15 @@ public <T> void onSuccess(T ignored) {
       accept(null);
     }
 
+    public void onCancel() {
+      if (compareAndSet(false, true)) {
+        if (captureExperimentalSpanAttributes) {
+          Span.fromContext(context).setAttribute(CANCELED_ATTRIBUTE_KEY, true);
+        }
+        tracer.end(context);
+      }
+    }
+
     @Override
     public void run() {
       accept(null);

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

@@ -0,0 +1,22 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.reactor;
+
+public final class ReactorAsyncSpanEndStrategyBuilder {
+  private boolean captureExperimentalSpanAttributes;
+
+  ReactorAsyncSpanEndStrategyBuilder() {}
+
+  public ReactorAsyncSpanEndStrategyBuilder setCaptureExperimentalSpanAttributes(
+      boolean captureExperimentalSpanAttributes) {
+    this.captureExperimentalSpanAttributes = captureExperimentalSpanAttributes;
+    return this;
+  }
+
+  public ReactorAsyncSpanEndStrategy build() {
+    return new ReactorAsyncSpanEndStrategy(captureExperimentalSpanAttributes);
+  }
+}

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

@@ -36,21 +36,39 @@
 /** Based on Spring Sleuth's Reactor instrumentation. */
 public final class TracingOperator {
 
+  public static TracingOperator create() {
+    return newBuilder().build();
+  }
+
+  public static TracingOperatorBuilder newBuilder() {
+    return new TracingOperatorBuilder();
+  }
+
+  private final boolean captureExperimentalSpanAttributes;
+
+  TracingOperator(boolean captureExperimentalSpanAttributes) {
+    this.captureExperimentalSpanAttributes = captureExperimentalSpanAttributes;
+  }
+
   /**
    * Registers a hook that applies to every operator, propagating {@link Context} to downstream
    * callbacks to ensure spans in the {@link Context} are available throughout the lifetime of a
    * reactive stream. This should generally be called in a static initializer block in your
    * application.
    */
-  public static void registerOnEachOperator() {
+  public void registerOnEachOperator() {
     Hooks.onEachOperator(TracingSubscriber.class.getName(), tracingLift());
-    AsyncSpanEndStrategies.getInstance().registerStrategy(ReactorAsyncSpanEndStrategy.INSTANCE);
+    AsyncSpanEndStrategies.getInstance()
+        .registerStrategy(
+            ReactorAsyncSpanEndStrategy.newBuilder()
+                .setCaptureExperimentalSpanAttributes(captureExperimentalSpanAttributes)
+                .build());
   }
 
   /** Unregisters the hook registered by {@link #registerOnEachOperator()}. */
-  public static void resetOnEachOperator() {
+  public void resetOnEachOperator() {
     Hooks.resetOnEachOperator(TracingSubscriber.class.getName());
-    AsyncSpanEndStrategies.getInstance().unregisterStrategy(ReactorAsyncSpanEndStrategy.INSTANCE);
+    AsyncSpanEndStrategies.getInstance().unregisterStrategy(ReactorAsyncSpanEndStrategy.class);
   }
 
   private static <T> Function<? super Publisher<T>, ? extends Publisher<T>> tracingLift() {
@@ -69,6 +87,4 @@ public CoreSubscriber<? super T> apply(Scannable publisher, CoreSubscriber<? sup
       return new TracingSubscriber<>(sub, sub.currentContext());
     }
   }
-
-  private TracingOperator() {}
 }

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

@@ -0,0 +1,22 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.reactor;
+
+public final class TracingOperatorBuilder {
+  private boolean captureExperimentalSpanAttributes;
+
+  TracingOperatorBuilder() {}
+
+  public TracingOperatorBuilder setCaptureExperimentalSpanAttributes(
+      boolean captureExperimentalSpanAttributes) {
+    this.captureExperimentalSpanAttributes = captureExperimentalSpanAttributes;
+    return this;
+  }
+
+  public TracingOperator build() {
+    return new TracingOperator(captureExperimentalSpanAttributes);
+  }
+}

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

@@ -14,6 +14,7 @@ class HooksTest extends LibraryInstrumentationSpecification {
 
   def "can reset out hooks"() {
     setup:
+    def underTest = TracingOperator.create()
     AtomicReference<CoreSubscriber> subscriber = new AtomicReference<>()
 
     when: "no hook registered"
@@ -23,14 +24,14 @@ class HooksTest extends LibraryInstrumentationSpecification {
     !(subscriber.get() instanceof TracingSubscriber)
 
     when: "hook registered"
-    TracingOperator.registerOnEachOperator()
+    underTest.registerOnEachOperator()
     new CapturingMono(subscriber).map { it + 1 }.subscribe()
 
     then:
     subscriber.get() instanceof TracingSubscriber
 
     when: "hook reset"
-    TracingOperator.resetOnEachOperator()
+    underTest.resetOnEachOperator()
     new CapturingMono(subscriber).map { it + 1 }.subscribe()
 
     then: