Suppress nested SDK public method calls

sdk/core/azure-core-tracing-opentelemetry/pom.xml

@@ -34,6 +34,7 @@
   <properties>
     <javaModulesSurefireArgLine>
       --add-opens com.azure.core.tracing.opentelemetry/com.azure.core.tracing.opentelemetry.implementation=ALL-UNNAMED
+      --add-opens com.azure.core.tracing.opentelemetry/com.azure.core.tracing.opentelemetry=ALL-UNNAMED
     </javaModulesSurefireArgLine>
   </properties>
 

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

@@ -11,7 +11,6 @@
 import com.azure.core.http.policy.AfterRetryPolicyProvider;
 import com.azure.core.http.policy.HttpPipelinePolicy;
 import com.azure.core.tracing.opentelemetry.implementation.HttpTraceUtil;
-import com.azure.core.tracing.opentelemetry.implementation.OpenTelemetrySpanSuppressionHelper;
 import com.azure.core.util.CoreUtils;
 import io.opentelemetry.api.GlobalOpenTelemetry;
 import io.opentelemetry.api.common.AttributeKey;
@@ -228,7 +227,7 @@ private static io.opentelemetry.context.Context getTraceContextOrCurrent(HttpPip
      */
     static final class ScalarPropagatingMono extends Mono<Object> {
         public static final Mono<Object> INSTANCE = new ScalarPropagatingMono();
-
+        private static final AutoCloseable NOOP_CLOSEABLE = () -> { };
         private final Object value = new Object();
 
         private ScalarPropagatingMono() {
@@ -237,16 +236,11 @@ private ScalarPropagatingMono() {
         @Override
         public void subscribe(CoreSubscriber<? super Object> actual) {
             Context traceContext = actual.currentContext().getOrDefault(REACTOR_PARENT_TRACE_CONTEXT_KEY, null);
-            if (traceContext != null) {
-                Object agentContext = OpenTelemetrySpanSuppressionHelper.registerClientSpan(traceContext);
-                AutoCloseable closeable = OpenTelemetrySpanSuppressionHelper.makeCurrent(agentContext, traceContext);
-                actual.onSubscribe(Operators.scalarSubscription(actual, value));
-                try {
-                    closeable.close();
-                } catch (Exception ignored) {
-                }
-            } else {
-                actual.onSubscribe(Operators.scalarSubscription(actual, value));
+            AutoCloseable closeable = traceContext != null ? traceContext.makeCurrent() : NOOP_CLOSEABLE;
+            actual.onSubscribe(Operators.scalarSubscription(actual, value));
+            try {
+                closeable.close();
+            } catch (Exception ignored) {
             }
         }
     }

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

@@ -35,6 +35,7 @@
  * requests.
  */
 public class OpenTelemetryTracer implements com.azure.core.util.tracing.Tracer {
+    private static final StartSpanOptions DEFAULT_OPTIONS = new StartSpanOptions(com.azure.core.util.tracing.SpanKind.INTERNAL);
     private final Tracer tracer;
 
     /**
@@ -64,28 +65,36 @@ public OpenTelemetryTracer() {
 
     private static final ClientLogger LOGGER = new ClientLogger(OpenTelemetryTracer.class);
     private static final AutoCloseable NOOP_CLOSEABLE = () -> { };
+    private static final SpanKind SHARED_SPAN_BUILDER_KIND = SpanKind.CLIENT;
+    private static final String SUPPRESSED_SPAN_FLAG = "suppressed-span-flag";
+    private static final String CLIENT_METHOD_CALL_FLAG = "client-method-call-flag";
 
     /**
      * {@inheritDoc}
      */
     @Override
     public Context start(String spanName, Context context) {
-        Objects.requireNonNull(spanName, "'spanName' cannot be null.");
-        SpanBuilder spanBuilder = createSpanBuilder(spanName, null, SpanKind.INTERNAL, null, context);
-
-        return startSpanInternal(spanBuilder, null, context);
+        return start(spanName, DEFAULT_OPTIONS, context);
     }
 
     /**
      * {@inheritDoc}
      */
     @Override
     public Context start(String spanName, StartSpanOptions options, Context context) {
+        Objects.requireNonNull(spanName, "'spanName' cannot be null.");
         Objects.requireNonNull(options, "'options' cannot be null.");
-        SpanBuilder spanBuilder = createSpanBuilder(spanName, null, convertToOtelKind(options.getSpanKind()),
+
+        SpanKind spanKind = convertToOtelKind(options.getSpanKind());
+        if (shouldSuppress(spanKind, context)) {
+            return startSuppressedSpan(context);
+        }
+
+        context = unsuppress(context);
+        SpanBuilder spanBuilder = createSpanBuilder(spanName, null, spanKind,
             options.getAttributes(), context);
 
-        return startSpanInternal(spanBuilder, null, context);
+        return startSpanInternal(spanBuilder, isClientCall(spanKind), null, context);
     }
 
     /**
@@ -98,28 +107,44 @@ public Context start(String spanName, Context context, ProcessKind processKind)
         Objects.requireNonNull(processKind, "'processKind' cannot be null.");
 
         SpanBuilder spanBuilder;
+        if (shouldSuppress(processKindToSpanKind(processKind), context)) {
+            return startSuppressedSpan(context);
+        }
+
+        context = unsuppress(context);
         switch (processKind) {
             case SEND:
                 // use previously created span builder from the LINK process.
                 spanBuilder = getOrNull(context, SPAN_BUILDER_KEY, SpanBuilder.class);
                 if (spanBuilder == null) {
-                    return context;
+                    // we can't return context here, because called would not know that span was not created.
+                    // will add attributes, events to parent span and end parent span.
+                    LOGGER.atWarning()
+                        .addKeyValue("spanName", spanName)
+                        .addKeyValue("processKind", processKind)
+                        .log("start span is called without builder on the context, creating default builder");
+                    spanBuilder = createSpanBuilder(spanName, null, SHARED_SPAN_BUILDER_KIND, null, context);
                 }
-                return startSpanInternal(spanBuilder, this::addMessagingAttributes, context);
+
+                return startSpanInternal(spanBuilder, isClientCall(SHARED_SPAN_BUILDER_KIND), this::addMessagingAttributes, context);
             case MESSAGE:
                 spanBuilder = createSpanBuilder(spanName, null, SpanKind.PRODUCER, null, context);
-                context = startSpanInternal(spanBuilder, this::addMessagingAttributes, context);
+                context = startSpanInternal(spanBuilder, false, this::addMessagingAttributes, context);
                 return setDiagnosticId(context);
             case PROCESS:
                 SpanContext remoteParentContext = getOrNull(context, SPAN_CONTEXT_KEY, SpanContext.class);
                 spanBuilder = createSpanBuilder(spanName, remoteParentContext, SpanKind.CONSUMER, null, context);
-                context = startSpanInternal(spanBuilder, this::addMessagingAttributes, context);
+                context = startSpanInternal(spanBuilder, false, this::addMessagingAttributes, context);
 
                 // TODO (limolkova) we should do this in the EventHub/ServiceBus SDK instead to make sure scope is
                 //  closed in the same thread where it was started to prevent leaking the context.
                 return context.addData(SCOPE_KEY, makeSpanCurrent(context));
             default:
-                return context;
+                LOGGER.atWarning()
+                    .addKeyValue("spanName", spanName)
+                    .addKeyValue("processKind", processKind)
+                    .log("start span is called with unknown process kind, suppressing the span");
+                return startSuppressedSpan(context);
         }
     }
 
@@ -129,6 +154,7 @@ public Context start(String spanName, Context context, ProcessKind processKind)
     @Override
     public void end(int responseCode, Throwable throwable, Context context) {
         Objects.requireNonNull(context, "'context' cannot be null.");
+
         final Span span = getSpanOrNull(context);
         if (span == null) {
             return;
@@ -146,6 +172,7 @@ public void end(int responseCode, Throwable throwable, Context context) {
     @Override
     public void setAttribute(String key, String value, Context context) {
         Objects.requireNonNull(context, "'context' cannot be null");
+
         if (CoreUtils.isNullOrEmpty(value)) {
             LOGGER.verbose("Failed to set span attribute since value is null or empty.");
             return;
@@ -175,7 +202,6 @@ public Context setSpanName(String spanName, Context context) {
     @Override
     public void end(String statusMessage, Throwable throwable, Context context) {
         Span span = getSpanOrNull(context);
-
         if (span == null) {
             return;
         }
@@ -218,14 +244,18 @@ public Context extractContext(String diagnosticId, Context context) {
     @Override
     public Context getSharedSpanBuilder(String spanName, Context context) {
         // this is used to create messaging send spanBuilder, and it's a CLIENT span
-        return context.addData(SPAN_BUILDER_KEY, createSpanBuilder(spanName, null, SpanKind.CLIENT, null, context));
+        return context.addData(SPAN_BUILDER_KEY, createSpanBuilder(spanName, null, SHARED_SPAN_BUILDER_KIND, null, context));
     }
 
     /**
      * {@inheritDoc}
      */
     @Override
     public AutoCloseable makeSpanCurrent(Context context) {
+        if (getBoolean(SUPPRESSED_SPAN_FLAG, context)) {
+            return NOOP_CLOSEABLE;
+        }
+
         io.opentelemetry.context.Context traceContext = getTraceContextOrDefault(context, null);
         if (traceContext == null) {
             LOGGER.verbose("There is no OpenTelemetry Context on the context, cannot make it current");
@@ -278,6 +308,7 @@ public void addEvent(String eventName, Map<String, Object> traceEventAttributes,
      * @return A {@link Context} with created {@link Span}.
      */
     private Context startSpanInternal(SpanBuilder spanBuilder,
+        boolean isClientMethod,
         java.util.function.BiConsumer<Span, Context> setAttributes,
         Context context) {
         Objects.requireNonNull(spanBuilder, "'spanBuilder' cannot be null.");
@@ -297,7 +328,12 @@ private Context startSpanInternal(SpanBuilder spanBuilder,
             }
         }
 
-        return context.addData(PARENT_TRACE_CONTEXT_KEY, getTraceContextOrDefault(context, io.opentelemetry.context.Context.current()).with(span));
+        if (isClientMethod && !context.getData(CLIENT_METHOD_CALL_FLAG).isPresent()) {
+            context = context.addData(CLIENT_METHOD_CALL_FLAG, true);
+        }
+
+        return context
+            .addData(PARENT_TRACE_CONTEXT_KEY, getTraceContextOrDefault(context, io.opentelemetry.context.Context.current()).with(span));
     }
 
     /**
@@ -307,15 +343,15 @@ private Context startSpanInternal(SpanBuilder spanBuilder,
      * @param spanName The name of the returned Span.
      * @param remoteParentContext Remote parent context if any, or {@code null} otherwise.
      * @param spanKind Kind of the span to create.
-     * @param beforeSaplingAttributes Optional attributes available when span starts and important for sampling.
+     * @param beforeSamplingAttributes Optional attributes available when span starts and important for sampling.
      * @param context The context containing the span and the span name.
      * @return A {@link SpanBuilder} to create and start a new {@link Span}.
      */
     @SuppressWarnings("unchecked")
     private SpanBuilder createSpanBuilder(String spanName,
         SpanContext remoteParentContext,
         SpanKind spanKind,
-        Map<String, Object> beforeSaplingAttributes,
+        Map<String, Object> beforeSamplingAttributes,
         Context context) {
         String spanNameKey = getOrNull(context, USER_SPAN_NAME_KEY, String.class);
 
@@ -335,8 +371,8 @@ private SpanBuilder createSpanBuilder(String spanName,
         }
 
         // if some attributes are provided, set them
-        if (!CoreUtils.isNullOrEmpty(beforeSaplingAttributes)) {
-            Attributes otelAttributes = convertToOtelAttributes(beforeSaplingAttributes);
+        if (!CoreUtils.isNullOrEmpty(beforeSamplingAttributes)) {
+            Attributes otelAttributes = convertToOtelAttributes(beforeSamplingAttributes);
             otelAttributes.forEach(
                 (key, value) -> spanBuilder.setAttribute((AttributeKey<Object>) key, value));
         }
@@ -468,7 +504,7 @@ private void addMessagingAttributes(Span span, Context context) {
      * @return The T type of raw class object
      */
     @SuppressWarnings("unchecked")
-    private <T> T getOrNull(Context context, String key, Class<T> clazz) {
+    private static <T> T getOrNull(Context context, String key, Class<T> clazz) {
         final Optional<Object> optional = context.getData(key);
         final Object result = optional.filter(value -> clazz.isAssignableFrom(value.getClass())).orElseGet(() -> {
             LOGGER.verbose("Could not extract key '{}' of type '{}' from context.", key, clazz);
@@ -484,7 +520,7 @@ private <T> T getOrNull(Context context, String key, Class<T> clazz) {
      * or PARENT_SPAN_KEY (for backward-compatibility) or default value.
      */
     @SuppressWarnings("deprecation")
-    private io.opentelemetry.context.Context getTraceContextOrDefault(Context azContext, io.opentelemetry.context.Context defaultContext) {
+    private static io.opentelemetry.context.Context getTraceContextOrDefault(Context azContext, io.opentelemetry.context.Context defaultContext) {
         io.opentelemetry.context.Context traceContext = getOrNull(azContext,
             PARENT_TRACE_CONTEXT_KEY,
             io.opentelemetry.context.Context.class);
@@ -507,6 +543,10 @@ private io.opentelemetry.context.Context getTraceContextOrDefault(Context azCont
      */
     @SuppressWarnings("deprecation")
     private Span getSpanOrNull(Context azContext) {
+        if (getBoolean(SUPPRESSED_SPAN_FLAG, azContext)) {
+            return null;
+        }
+
         io.opentelemetry.context.Context traceContext = getOrNull(azContext,
             PARENT_TRACE_CONTEXT_KEY,
             io.opentelemetry.context.Context.class);
@@ -523,4 +563,49 @@ private Span getSpanOrNull(Context azContext) {
 
         return traceContext == null ? null : Span.fromContext(traceContext);
     }
+
+    private SpanKind processKindToSpanKind(ProcessKind processKind) {
+        switch (processKind) {
+            case SEND:
+                return SHARED_SPAN_BUILDER_KIND;
+            case MESSAGE:
+                return SpanKind.PRODUCER;
+            case PROCESS:
+                return SpanKind.CONSUMER;
+            default:
+                return SpanKind.INTERNAL;
+        }
+    }
+
+    private Context startSuppressedSpan(Context context) {
+        return context.addData(SUPPRESSED_SPAN_FLAG, true);
+    }
+
+    private static boolean shouldSuppress(SpanKind kind, Context context) {
+        if (isClientCall(kind)) {
+            boolean suppress = getBoolean(CLIENT_METHOD_CALL_FLAG, context);
+            boolean suppressed = getBoolean(SUPPRESSED_SPAN_FLAG, context);
+
+            return suppress && !suppressed;
+        }
+
+        return false;
+    }
+
+    private static Context unsuppress(Context context) {
+        if (getBoolean(SUPPRESSED_SPAN_FLAG, context)) {
+            return context.addData(SUPPRESSED_SPAN_FLAG, false);
+        }
+
+        return context;
+    }
+
+    private static boolean getBoolean(String key, Context context) {
+        Optional<Object> flag = context.getData(key);
+        return flag.isPresent() && Boolean.TRUE.equals(flag.get());
+    }
+
+    private static boolean isClientCall(SpanKind kind) {
+        return kind == SpanKind.CLIENT || kind == SpanKind.INTERNAL;
+    }
 }