Tracing: don't make http span current and ignore current if parent context was provided (Azure#33971)

* don't make http span current
* don't use current as parent if invalid context is extracted
* fail span on cancellation

Author: Liudmila Molkova

eng/code-quality-reports/src/main/resources/spotbugs/spotbugs-exclude.xml

@@ -566,8 +566,8 @@
 
   <!-- Incorrect flagging, the surrounding hasValue() checks that get() isn't null -->
   <Match>
-    <Class name="com.azure.core.http.rest.RestProxy"/>
-    <Method name="endTracingSpan"/>
+    <Class name="com.azure.core.implementation.http.rest.AsyncRestProxy"/>
+    <Method name="~(.*)\$endSpanWhenDone\$(.*)"/>
     <Bug pattern="NP_NULL_ON_SOME_PATH_FROM_RETURN_VALUE"/>
   </Match>
 

sdk/core/azure-core-tracing-opentelemetry/src/main/java/com/azure/core/tracing/opentelemetry/OpenTelemetryTracer.java

@@ -185,7 +185,7 @@ private SpanBuilder createSpanBuilder(String spanName,
         Context remoteParentContext = options.getRemoteParent();
         SpanContext remoteSpanContext = remoteParentContext == null ? null : getOrNull(remoteParentContext, SPAN_CONTEXT_KEY, SpanContext.class);
 
-        if (remoteSpanContext != null && remoteSpanContext.isValid()) {
+        if (remoteSpanContext != null) {
             if (parentContext == null) {
                 parentContext = io.opentelemetry.context.Context.root();
             }

sdk/core/azure-core-tracing-opentelemetry/src/test/java/com/azure/core/tracing/opentelemetry/OpenTelemetryHttpPolicyTests.java

@@ -23,6 +23,7 @@
 import io.opentelemetry.api.trace.SpanKind;
 import io.opentelemetry.api.trace.StatusCode;
 import io.opentelemetry.api.trace.Tracer;
+import io.opentelemetry.context.Scope;
 import io.opentelemetry.sdk.OpenTelemetrySdk;
 import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter;
 import io.opentelemetry.sdk.trace.ReadableSpan;
@@ -43,6 +44,7 @@
 import reactor.core.publisher.Mono;
 import reactor.test.StepVerifier;
 
+import java.time.Duration;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashMap;
@@ -62,6 +64,7 @@
 /**
  * Unit tests for {com.azure.core.implementation.http.policy.InstrumentationPolicy.
  */
+@SuppressWarnings("try")
 public class OpenTelemetryHttpPolicyTests {
 
     private static final String X_MS_REQUEST_ID_1 = "response id 1";
@@ -85,7 +88,7 @@ public void setUp(TestInfo testInfo) {
 
     @Test
     public void openTelemetryHttpPolicyTest() {
-        // Start user parent span.
+        // Start parent span.
         Span parentSpan = tracer.spanBuilder(SPAN_NAME).startSpan();
 
         // Add parent span to tracingContext
@@ -95,8 +98,10 @@ public void openTelemetryHttpPolicyTest() {
         // Act
         HttpRequest request = new HttpRequest(HttpMethod.POST, "https://httpbin.org/hello?there#otel");
         request.setHeader("User-Agent", "user-agent");
-        HttpResponse response =  createHttpPipeline(azTracer).send(request, tracingContext).block();
 
+        try (Scope scope = parentSpan.makeCurrent()) {
+            createHttpPipeline(azTracer).send(request, tracingContext).block();
+        }
         // Assert
         List<SpanData> exportedSpans = exporter.getFinishedSpanItems();
         // rest proxy span is not exported as global otel is not configured
@@ -150,10 +155,12 @@ public String getDescription() {
 
         // Act
         HttpRequest request = new HttpRequest(HttpMethod.DELETE, "https://httpbin.org/hello?there#otel");
-        HttpResponse response =  createHttpPipeline(new OpenTelemetryTracer("test", null, null,
+        try (Scope scope = tracer.spanBuilder("test").startSpan().makeCurrent()) {
+            createHttpPipeline(new OpenTelemetryTracer("test", null, null,
                 new OpenTelemetryTracingOptions().setProvider(providerWithSampler)))
-            .send(request).block();
-
+                .send(request)
+                .block();
+        }
         // Assert
         List<SpanData> exportedSpans = exporter.getFinishedSpanItems();
         assertEquals(0, exportedSpans.size());
@@ -162,33 +169,33 @@ public String getDescription() {
 
     @Test
     public void clientRequestIdIsStamped() {
-        HttpRequest request = new HttpRequest(HttpMethod.PUT, "https://httpbin.org/hello?there#otel");
-        HttpResponse response =  createHttpPipeline(azTracer, new RequestIdPolicy()).send(request).block();
+        try (Scope scope = tracer.spanBuilder("test").startSpan().makeCurrent()) {
+            HttpRequest request = new HttpRequest(HttpMethod.PUT, "https://httpbin.org/hello?there#otel");
+            HttpResponse response = createHttpPipeline(azTracer, new RequestIdPolicy()).send(request).block();
 
-        // Assert
-        List<SpanData> exportedSpans = exporter.getFinishedSpanItems();
-        assertEquals(1, exportedSpans.size());
+            // Assert
+            List<SpanData> exportedSpans = exporter.getFinishedSpanItems();
+            assertEquals(1, exportedSpans.size());
 
-        assertEquals("HTTP PUT", exportedSpans.get(0).getName());
+            assertEquals("HTTP PUT", exportedSpans.get(0).getName());
 
-        Map<String, Object> httpAttributes = getAttributes(exportedSpans.get(0));
-        assertEquals(5, httpAttributes.size());
+            Map<String, Object> httpAttributes = getAttributes(exportedSpans.get(0));
+            assertEquals(5, httpAttributes.size());
 
-        assertEquals(response.getRequest().getHeaders().getValue("x-ms-client-request-id"), httpAttributes.get("az.client_request_id"));
-        assertEquals(X_MS_REQUEST_ID_1, httpAttributes.get("az.service_request_id"));
+            assertEquals(response.getRequest().getHeaders().getValue("x-ms-client-request-id"), httpAttributes.get("az.client_request_id"));
+            assertEquals(X_MS_REQUEST_ID_1, httpAttributes.get("az.service_request_id"));
 
-        assertEquals("https://httpbin.org/hello?there#otel", httpAttributes.get("http.url"));
-        assertEquals("PUT", httpAttributes.get("http.method"));
-        assertEquals(Long.valueOf(RESPONSE_STATUS_CODE), httpAttributes.get("http.status_code"));
+            assertEquals("https://httpbin.org/hello?there#otel", httpAttributes.get("http.url"));
+            assertEquals("PUT", httpAttributes.get("http.method"));
+            assertEquals(Long.valueOf(RESPONSE_STATUS_CODE), httpAttributes.get("http.status_code"));
+        }
     }
 
     @Test
     public void everyTryIsTraced() {
         AtomicInteger attemptCount = new AtomicInteger();
         AtomicReference<String> traceparentTry503 = new AtomicReference<>();
         AtomicReference<String> traceparentTry200 = new AtomicReference<>();
-        AtomicReference<Span> currentSpanTry503 = new AtomicReference<>();
-        AtomicReference<Span> currentSpanTry200 = new AtomicReference<>();
 
         OpenTelemetryTracingOptions options = new OpenTelemetryTracingOptions().setProvider(tracerProvider);
 
@@ -205,12 +212,10 @@ public void everyTryIsTraced() {
                 int count = attemptCount.getAndIncrement();
                 if (count == 0) {
                     traceparentTry503.set(request.getHeaders().getValue("traceparent"));
-                    currentSpanTry503.set(Span.current());
                     headers.set("x-ms-request-id", X_MS_REQUEST_ID_1);
                     return Mono.just(new MockHttpResponse(request, 503, headers));
                 } else if (count == 1) {
                     traceparentTry200.set(request.getHeaders().getValue("traceparent"));
-                    currentSpanTry200.set(Span.current());
                     headers.set("x-ms-request-id", X_MS_REQUEST_ID_2);
                     return Mono.just(new MockHttpResponse(request, 200, headers));
                 } else {
@@ -240,9 +245,6 @@ public void everyTryIsTraced() {
         assertEquals(traceparentTry503.get(), String.format("00-%s-%s-01", try503.getTraceId(), try503.getSpanId()));
         assertEquals(traceparentTry200.get(), String.format("00-%s-%s-01", try200.getTraceId(), try200.getSpanId()));
 
-        assertEquals(currentSpanTry503.get().getSpanContext().getSpanId(), try503.getSpanId());
-        assertEquals(currentSpanTry200.get().getSpanContext().getSpanId(), try200.getSpanId());
-
         assertEquals("HTTP GET", try503.getName());
         Map<String, Object> httpAttributes503 = getAttributes(try503);
         assertEquals(5, httpAttributes503.size());
@@ -370,6 +372,36 @@ public void timeoutIsTraced() {
         assertEquals(TimeoutException.class.getName(), events.get(0).getAttributes().get(AttributeKey.stringKey("exception.type")));
     }
 
+    @Test
+    public void cancelIsTraced() {
+        OpenTelemetryTracingOptions options = new OpenTelemetryTracingOptions().setProvider(tracerProvider);
+
+        com.azure.core.util.tracing.Tracer azTracer = new OpenTelemetryTracer("test", null, null, options);
+
+        List<HttpPipelinePolicy> policies = new ArrayList<>(Arrays.asList(new RetryPolicy()));
+        HttpPolicyProviders.addAfterRetryPolicies(policies);
+
+        HttpPipeline pipeline = new HttpPipelineBuilder()
+            .policies(policies.toArray(new HttpPipelinePolicy[0]))
+            .httpClient(request ->
+                Mono.delay(Duration.ofSeconds(10)).map(l -> new MockHttpResponse(request, 200)))
+            .tracer(azTracer)
+            .build();
+
+        pipeline.send(new HttpRequest(HttpMethod.GET, "http://localhost/hello"), Context.NONE)
+                .toFuture()
+            .cancel(true);
+
+        List<SpanData> exportedSpans = exporter.getFinishedSpanItems();
+        assertEquals(1, exportedSpans.size());
+
+        SpanData cancelled = exportedSpans.get(0);
+        Map<String, Object> httpAttributesTimeout = getAttributes(cancelled);
+        assertNull(httpAttributesTimeout.get("http.status_code"));
+        assertEquals(StatusCode.ERROR, cancelled.getStatus().getStatusCode());
+        assertEquals("cancel", cancelled.getStatus().getDescription());
+    }
+
     private Map<String, Object> getAttributes(SpanData span) {
         Map<String, Object> attributes = new HashMap<>();
         span.getAttributes().forEach((k, v) -> attributes.put(k.getKey(), v));
@@ -411,13 +443,10 @@ public Mono<HttpResponse> send(HttpRequest request) {
             HttpHeaders headers = new HttpHeaders()
                 .set("x-ms-request-id", X_MS_REQUEST_ID_1);
 
+            // parent span
             SpanContext currentContext = Span.current().getSpanContext();
-            String expectedTraceparent = String.format("00-%s-%s-%s", currentContext.getTraceId(), currentContext.getSpanId(), currentContext.getTraceFlags().toString());
-            if (currentContext.isValid()) {
-                assertEquals(expectedTraceparent, request.getHeaders().getValue("traceparent"));
-            } else {
-                assertNull(request.getHeaders().getValue("traceparent"));
-            }
+            assertTrue(currentContext.isValid());
+            assertEquals(currentContext.getTraceId(), request.getHeaders().getValue("traceparent").substring(3, 35));
 
             return Mono.just(new MockHttpResponse(request, RESPONSE_STATUS_CODE, headers));
         }

sdk/core/azure-core-tracing-opentelemetry/src/test/java/com/azure/core/tracing/opentelemetry/OpenTelemetryTracerTest.java

@@ -1143,9 +1143,7 @@ public void invalidRemoteParent() {
         assertEquals(2, testExporter.getFinishedSpanItems().size());
 
         SpanData innerSpan = testExporter.getFinishedSpanItems().get(0);
-        SpanData outerSpan = testExporter.getFinishedSpanItems().get(1);
-        assertEquals(innerSpan.getSpanContext().getTraceId(), outerSpan.getSpanContext().getTraceId());
-        assertEquals(innerSpan.getParentSpanId(), outerSpan.getSpanContext().getSpanId());
+        assertFalse(innerSpan.getParentSpanContext().isValid());
     }
 
     @Test

sdk/core/azure-core/pom.xml

@@ -68,6 +68,7 @@
       --add-opens com.azure.core/com.azure.core.implementation.logging=ALL-UNNAMED
       --add-opens com.azure.core/com.azure.core.models=ALL-UNNAMED
       --add-opens com.azure.core/com.azure.core.util=ALL-UNNAMED
+      --add-opens com.azure.core/com.azure.core.util.tracing=ALL-UNNAMED
       --add-opens com.azure.core/com.azure.core.util.logging=ALL-UNNAMED
       --add-opens com.azure.core/com.azure.core.util.polling=ALL-UNNAMED
       --add-opens com.azure.core/com.azure.core.util.serializer=ALL-UNNAMED

sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/policy/InstrumentationPolicy.java

@@ -16,10 +16,7 @@
 import com.azure.core.util.tracing.SpanKind;
 import com.azure.core.util.tracing.StartSpanOptions;
 import com.azure.core.util.tracing.Tracer;
-import reactor.core.CoreSubscriber;
 import reactor.core.publisher.Mono;
-import reactor.core.publisher.Operators;
-import reactor.core.publisher.Signal;
 
 import static com.azure.core.util.tracing.Tracer.DISABLE_TRACING_KEY;
 
@@ -37,7 +34,6 @@ public class InstrumentationPolicy implements HttpPipelinePolicy {
     private static final ClientLogger LOGGER = new ClientLogger(InstrumentationPolicy.class);
 
     private Tracer tracer;
-    private ScalarPropagatingMono propagatingMono;
     private static boolean foundLegacyOTelPolicy;
 
     static {
@@ -53,7 +49,6 @@ public InstrumentationPolicy() {
 
     public void initialize(Tracer tracer) {
         this.tracer = tracer;
-        this.propagatingMono = new ScalarPropagatingMono(tracer);
     }
 
     @Override
@@ -65,10 +60,13 @@ public Mono<HttpResponse> process(HttpPipelineCallContext context, HttpPipelineN
         // OpenTelemetry reactor instrumentation needs a bit of help
         // to pick up Azure SDK context. While we're working on explicit
         // context propagation, ScalarPropagatingMono.INSTANCE is the workaround
-        return propagatingMono
-            .flatMap(ignored -> next.process())
-            .doOnEach(this::handleResponse)
-            .contextWrite(reactor.util.context.Context.of(REACTOR_HTTP_TRACE_CONTEXT_KEY, startSpan(context)));
+        return Mono.defer(() -> {
+            Context span = startSpan(context);
+            return next.process()
+                .doOnSuccess(response -> endSpan(response, span))
+                .doOnCancel(() -> tracer.end("cancel", null, span))
+                .doOnError(exception -> tracer.end(null, exception, span));
+        });
     }
 
     @SuppressWarnings("try")
@@ -79,19 +77,16 @@ public HttpResponse processSync(HttpPipelineCallContext context, HttpPipelineNex
         }
 
         Context span = startSpan(context);
-        Throwable exception = null;
-        HttpResponse response = null;
         try (AutoCloseable scope = tracer.makeSpanCurrent(span)) {
-            response = next.processSync();
+            HttpResponse response = next.processSync();
+            endSpan(response, span);
             return response;
         } catch (RuntimeException ex) {
-            exception = ex;
+            tracer.end(null, ex, span);
             throw ex;
         } catch (Exception ex) {
-            exception = ex;
+            tracer.end(null, ex, span);
             throw LOGGER.logExceptionAsWarning(new RuntimeException(ex));
-        } finally {
-            endSpan(response, exception, span);
         }
     }
 
@@ -123,23 +118,7 @@ private void addPostSamplingAttributes(Context span, HttpRequest request) {
         }
     }
 
-    /**
-     * Handles retrieving the information from the service response and ending the span.
-     *
-     * @param signal Reactive Stream signal fired by Reactor.
-     */
-    private void handleResponse(Signal<? extends HttpResponse> signal) {
-        // Ignore the on complete and on subscribe events, they don't contain the information needed to end the span.
-        if (signal.isOnComplete() || signal.isOnSubscribe()) {
-            return;
-        }
-
-        // Get the context that was added to the mono, this will contain the information needed to end the span.
-        Context span = signal.getContextView().getOrDefault(REACTOR_HTTP_TRACE_CONTEXT_KEY, null);
-        endSpan(signal.get(), signal.getThrowable(), span);
-    }
-
-    private void endSpan(HttpResponse response, Throwable error, Context span) {
+    private void endSpan(HttpResponse response, Context span) {
         if (response != null) {
             int statusCode = response.getStatusCode();
             tracer.setAttribute(HTTP_STATUS_CODE, statusCode, span);
@@ -151,43 +130,11 @@ private void endSpan(HttpResponse response, Throwable error, Context span) {
             tracer.end((statusCode >= 400) ? "error" : null, null, span);
         }
 
-        tracer.end(null, error, span);
+        tracer.end("", null, span);
     }
 
     private boolean isTracingEnabled(HttpPipelineCallContext context) {
         return tracer != null && tracer.isEnabled() && !foundLegacyOTelPolicy
             && !((boolean) context.getData(DISABLE_TRACING_KEY).orElse(false));
     }
-
-    /**
-     * Helper class allowing to run Mono subscription and any hot path
-     * in scope of trace context. This enables OpenTelemetry auto-collection
-     * to pick it up and correlate lower levels of instrumentation and logs
-     * to logical/HTTP spans.
-     *
-     * OpenTelemetry reactor auto-instrumentation will take care of the cold path.
-     */
-    static final class ScalarPropagatingMono extends Mono<Object> {
-        private final Object value = new Object();
-        private final Tracer tracer;
-
-        private ScalarPropagatingMono(Tracer tracer) {
-            this.tracer = tracer;
-        }
-
-        @Override
-        @SuppressWarnings("try")
-        public void subscribe(CoreSubscriber<? super Object> actual) {
-            Context traceContext = actual.currentContext().getOrDefault(REACTOR_HTTP_TRACE_CONTEXT_KEY, null);
-            if (tracer.isEnabled() && traceContext != null) {
-                try (AutoCloseable scope = tracer.makeSpanCurrent(traceContext)) {
-                    actual.onSubscribe(Operators.scalarSubscription(actual, value));
-                } catch (Exception e) {
-                    LOGGER.verbose("Error closing scope", e);
-                }
-            } else {
-                actual.onSubscribe(Operators.scalarSubscription(actual, value));
-            }
-        }
-    }
 }