Update azure-core-tracing-opentelemetry version and fix sync suppress…

…ion (#10350) Co-authored-by: Lauri Tulmin <[email protected]>
open-telemetry · Feb 1, 2024 · 9b5e29d · 9b5e29d
1 parent 1d84b3d
commit 9b5e29d
Show file tree

Hide file tree

Showing 4 changed files with 68 additions and 19 deletions.
diff --git a/...entelemetry/javaagent/instrumentation/azurecore/v1_36/AzureHttpClientInstrumentation.java b/...entelemetry/javaagent/instrumentation/azurecore/v1_36/AzureHttpClientInstrumentation.java
@@ -6,12 +6,15 @@
 package io.opentelemetry.javaagent.instrumentation.azurecore.v1_36;
 
 import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.implementsInterface;
+import static io.opentelemetry.javaagent.instrumentation.azurecore.v1_36.SuppressNestedClientHelper.disallowNestedClientSpanMono;
+import static io.opentelemetry.javaagent.instrumentation.azurecore.v1_36.SuppressNestedClientHelper.disallowNestedClientSpanSync;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.isPublic;
 import static net.bytebuddy.matcher.ElementMatchers.named;
 import static net.bytebuddy.matcher.ElementMatchers.returns;
 
 import com.azure.core.http.HttpResponse;
+import io.opentelemetry.context.Scope;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
 import net.bytebuddy.asm.Advice;
@@ -33,15 +36,37 @@ public void transform(TypeTransformer transformer) {
             .and(isPublic())
             .and(named("send"))
             .and(returns(named("reactor.core.publisher.Mono"))),
-        this.getClass().getName() + "$SuppressNestedClientAdvice");
+        this.getClass().getName() + "$SuppressNestedClientMonoAdvice");
+    transformer.applyAdviceToMethod(
+        isMethod()
+            .and(isPublic())
+            .and(named("sendSync"))
+            .and(returns(named("com.azure.core.http.HttpResponse"))),
+        this.getClass().getName() + "$SuppressNestedClientSyncAdvice");
+  }
+
+  @SuppressWarnings("unused")
+  public static class SuppressNestedClientMonoAdvice {
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void asyncSendExit(@Advice.Return(readOnly = false) Mono<HttpResponse> mono) {
+      mono = disallowNestedClientSpanMono(mono);
+    }
   }
 
   @SuppressWarnings("unused")
-  public static class SuppressNestedClientAdvice {
+  public static class SuppressNestedClientSyncAdvice {
+
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void syncSendEnter(@Advice.Local("otelScope") Scope scope) {
+      scope = disallowNestedClientSpanSync();
+    }
 
     @Advice.OnMethodExit(suppress = Throwable.class)
-    public static void methodExit(@Advice.Return(readOnly = false) Mono<HttpResponse> mono) {
-      mono = new SuppressNestedClientMono<>(mono);
+    public static void syncSendExit(
+        @Advice.Return HttpResponse response, @Advice.Local("otelScope") Scope scope) {
+      if (scope != null) {
+        scope.close();
+      }
     }
   }
 }
diff --git a/...ecore/v1_36/SuppressNestedClientMono.java → ...ore/v1_36/SuppressNestedClientHelper.java b/...ecore/v1_36/SuppressNestedClientMono.java → ...ore/v1_36/SuppressNestedClientHelper.java
@@ -14,24 +14,30 @@
 import reactor.core.CoreSubscriber;
 import reactor.core.publisher.Mono;
 
-public class SuppressNestedClientMono<T> extends Mono<T> {
+public class SuppressNestedClientHelper {
 
-  private final Mono<T> delegate;
-
-  public SuppressNestedClientMono(Mono<T> delegate) {
-    this.delegate = delegate;
-  }
-
-  @Override
-  public void subscribe(CoreSubscriber<? super T> actual) {
+  public static Scope disallowNestedClientSpanSync() {
     Context parentContext = currentContext();
     if (doesNotHaveClientSpan(parentContext)) {
-      try (Scope ignored = disallowNestedClientSpan(parentContext).makeCurrent()) {
-        delegate.subscribe(actual);
-      }
-    } else {
-      delegate.subscribe(actual);
+      return disallowNestedClientSpan(parentContext).makeCurrent();
     }
+    return null;
+  }
+
+  public static <T> Mono<T> disallowNestedClientSpanMono(Mono<T> delegate) {
+    return new Mono<T>() {
+      @Override
+      public void subscribe(CoreSubscriber<? super T> coreSubscriber) {
+        Context parentContext = currentContext();
+        if (doesNotHaveClientSpan(parentContext)) {
+          try (Scope ignored = disallowNestedClientSpan(parentContext).makeCurrent()) {
+            delegate.subscribe(coreSubscriber);
+          }
+        } else {
+          delegate.subscribe(coreSubscriber);
+        }
+      }
+    };
   }
 
   private static boolean doesNotHaveClientSpan(Context parentContext) {
@@ -44,4 +50,6 @@ private static Context disallowNestedClientSpan(Context parentContext) {
     return SpanKey.HTTP_CLIENT.storeInContext(
         SpanKey.KIND_CLIENT.storeInContext(parentContext, span), span);
   }
+
+  private SuppressNestedClientHelper() {}
 }
diff --git a/instrumentation/azure-core/azure-core-1.36/library-instrumentation-shaded/build.gradle.kts b/instrumentation/azure-core/azure-core-1.36/library-instrumentation-shaded/build.gradle.kts
@@ -7,7 +7,7 @@ plugins {
 group = "io.opentelemetry.javaagent.instrumentation"
 
 dependencies {
-  implementation("com.azure:azure-core-tracing-opentelemetry:1.0.0-beta.32")
+  implementation("com.azure:azure-core-tracing-opentelemetry:1.0.0-beta.42")
 }
 
 tasks {

diff --git a/...haded/src/main/java/com/azure/core/tracing/opentelemetry/OpenTelemetryTracingOptions.java b/...haded/src/main/java/com/azure/core/tracing/opentelemetry/OpenTelemetryTracingOptions.java
@@ -0,0 +1,16 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package com.azure.core.tracing.opentelemetry;
+
+import com.azure.core.util.TracingOptions;
+
+/**
+ * Replace {@link OpenTelemetryTracingOptions} from com.azure:azure-core-tracing-opentelemetry with
+ * a stub. Auto instrumentation does not use {@link OpenTelemetryTracingOptions}. This is needed
+ * because {@link OpenTelemetryTracingOptions} calls super constructor in {@link TracingOptions}
+ * that does exist in com.azure:azure-core:1.36.0 which triggers muzzle failure.
+ */
+public class OpenTelemetryTracingOptions extends TracingOptions {}