Update azure-core-tracing-opentelemetry version and fix sync suppression

instrumentation/azure-core/azure-core-1.14/javaagent/build.gradle.kts

@@ -29,6 +29,7 @@ dependencies {
   // Ensure no cross interference
   testInstrumentation(project(":instrumentation:azure-core:azure-core-1.19:javaagent"))
   testInstrumentation(project(":instrumentation:azure-core:azure-core-1.36:javaagent"))
+  testInstrumentation(project(":instrumentation:azure-core:azure-core-1.40:javaagent"))
 }
 
 val latestDepTest = findProperty("testLatestDeps") as Boolean

instrumentation/azure-core/azure-core-1.19/javaagent/build.gradle.kts

@@ -29,6 +29,7 @@ dependencies {
   // Ensure no cross interference
   testInstrumentation(project(":instrumentation:azure-core:azure-core-1.14:javaagent"))
   testInstrumentation(project(":instrumentation:azure-core:azure-core-1.36:javaagent"))
+  testInstrumentation(project(":instrumentation:azure-core:azure-core-1.40:javaagent"))
 }
 
 val latestDepTest = findProperty("testLatestDeps") as Boolean

instrumentation/azure-core/azure-core-1.36/javaagent/build.gradle.kts

@@ -6,7 +6,7 @@ muzzle {
   pass {
     group.set("com.azure")
     module.set("azure-core")
-    versions.set("[1.36.0,)")
+    versions.set("[1.36.0,1.40.0)")
     assertInverse.set(true)
   }
 }
@@ -29,6 +29,7 @@ dependencies {
   // Ensure no cross interference
   testInstrumentation(project(":instrumentation:azure-core:azure-core-1.14:javaagent"))
   testInstrumentation(project(":instrumentation:azure-core:azure-core-1.19:javaagent"))
+  testInstrumentation(project(":instrumentation:azure-core:azure-core-1.40:javaagent"))
 }
 
 val latestDepTest = findProperty("testLatestDeps") as Boolean
@@ -40,7 +41,7 @@ testing {
     val testAzure by registering(JvmTestSuite::class) {
       dependencies {
         if (latestDepTest) {
-          implementation("com.azure:azure-core:+")
+          implementation("com.azure:azure-core:1.39.0")
         } else {
           implementation("com.azure:azure-core:1.36.0")
         }

instrumentation/azure-core/azure-core-1.36/javaagent/src/main/java/io/opentelemetry/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();
+      }
     }
   }
 }

instrumentation/azure-core/azure-core-1.36/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/azurecore/v1_36/AzureSdkInstrumentationModule.java

@@ -57,6 +57,9 @@ public void injectClasses(ClassInjector injector) {
   public ElementMatcher.Junction<ClassLoader> classLoaderMatcher() {
     // this class was introduced in azure-core 1.36
     return hasClassesNamed("com.azure.core.util.tracing.TracerProvider")
+        // this class was introduced in azure-core 1.40
+        .and(
+            not(hasClassesNamed("com.azure.core.implementation.jackson.ResponseErrorDeserializer")))
         .and(not(hasClassesNamed("com.azure.core.tracing.opentelemetry.OpenTelemetryTracer")));
   }
 

instrumentation/azure-core/azure-core-1.36/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/azurecore/v1_36/SuppressNestedClientMono.java → instrumentation/azure-core/azure-core-1.36/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/azurecore/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() {}
 }

instrumentation/azure-core/azure-core-1.40/javaagent/build.gradle.kts

@@ -0,0 +1,57 @@
+plugins {
+  id("otel.javaagent-instrumentation")
+}
+
+muzzle {
+  pass {
+    group.set("com.azure")
+    module.set("azure-core")
+    versions.set("[1.40.0,)")
+    assertInverse.set(true)
+  }
+}
+
+sourceSets {
+  main {
+    val shadedDep = project(":instrumentation:azure-core:azure-core-1.40:library-instrumentation-shaded")
+    output.dir(
+      shadedDep.file("build/extracted/shadow"),
+      "builtBy" to ":instrumentation:azure-core:azure-core-1.40:library-instrumentation-shaded:extractShadowJar"
+    )
+  }
+}
+
+dependencies {
+  compileOnly(project(":instrumentation:azure-core:azure-core-1.40:library-instrumentation-shaded", configuration = "shadow"))
+
+  library("com.azure:azure-core:1.40.0")
+
+  // Ensure no cross interference
+  testInstrumentation(project(":instrumentation:azure-core:azure-core-1.14:javaagent"))
+  testInstrumentation(project(":instrumentation:azure-core:azure-core-1.19:javaagent"))
+  testInstrumentation(project(":instrumentation:azure-core:azure-core-1.36:javaagent"))
+}
+
+val latestDepTest = findProperty("testLatestDeps") as Boolean
+
+testing {
+  suites {
+    // using a test suite to ensure that classes from library-instrumentation-shaded that were
+    // extracted to the output directory are not available during tests
+    val testAzure by registering(JvmTestSuite::class) {
+      dependencies {
+        if (latestDepTest) {
+          implementation("com.azure:azure-core:+")
+        } else {
+          implementation("com.azure:azure-core:1.40.0")
+        }
+      }
+    }
+  }
+}
+
+tasks {
+  check {
+    dependsOn(testing.suites)
+  }
+}

instrumentation/azure-core/azure-core-1.40/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/azurecore/v1_40/AzureHttpClientInstrumentation.java

@@ -0,0 +1,72 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.azurecore.v1_40;
+
+import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.implementsInterface;
+import static io.opentelemetry.javaagent.instrumentation.azurecore.v1_40.SuppressNestedClientHelper.disallowNestedClientSpanMono;
+import static io.opentelemetry.javaagent.instrumentation.azurecore.v1_40.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;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+import reactor.core.publisher.Mono;
+
+public class AzureHttpClientInstrumentation implements TypeInstrumentation {
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return implementsInterface(named("com.azure.core.http.HttpClient"));
+  }
+
+  @Override
+  public void transform(TypeTransformer transformer) {
+    transformer.applyAdviceToMethod(
+        isMethod()
+            .and(isPublic())
+            .and(named("send"))
+            .and(returns(named("reactor.core.publisher.Mono"))),
+        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 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 syncSendExit(
+        @Advice.Return HttpResponse response, @Advice.Local("otelScope") Scope scope) {
+      if (scope != null) {
+        scope.close();
+      }
+    }
+  }
+}

instrumentation/azure-core/azure-core-1.40/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/azurecore/v1_40/AzureSdkInstrumentationModule.java

@@ -0,0 +1,82 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.azurecore.v1_40;
+
+import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasClassesNamed;
+import static java.util.Arrays.asList;
+import static net.bytebuddy.matcher.ElementMatchers.namedOneOf;
+import static net.bytebuddy.matcher.ElementMatchers.not;
+
+import com.google.auto.service.AutoService;
+import io.opentelemetry.javaagent.extension.instrumentation.HelperResourceBuilder;
+import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import io.opentelemetry.javaagent.extension.instrumentation.internal.ExperimentalInstrumentationModule;
+import io.opentelemetry.javaagent.extension.instrumentation.internal.injection.ClassInjector;
+import io.opentelemetry.javaagent.extension.instrumentation.internal.injection.InjectionMode;
+import java.util.List;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+@AutoService(InstrumentationModule.class)
+public class AzureSdkInstrumentationModule extends InstrumentationModule
+    implements ExperimentalInstrumentationModule {
+  public AzureSdkInstrumentationModule() {
+    super("azure-core", "azure-core-1.40");
+  }
+
+  @Override
+  public void registerHelperResources(HelperResourceBuilder helperResourceBuilder) {
+    helperResourceBuilder.register(
+        "META-INF/services/com.azure.core.util.tracing.TracerProvider",
+        "azure-core-1.40/META-INF/services/com.azure.core.util.tracing.TracerProvider");
+    // some azure sdks (e.g. EventHubs) are still looking up Tracer via service loader
+    // and not yet using the new TracerProvider
+    helperResourceBuilder.register(
+        "META-INF/services/com.azure.core.util.tracing.Tracer",
+        "azure-core-1.40/META-INF/services/com.azure.core.util.tracing.Tracer");
+  }
+
+  @Override
+  public void injectClasses(ClassInjector injector) {
+    injector
+        .proxyBuilder(
+            "io.opentelemetry.javaagent.instrumentation.azurecore.v1_40.shaded.com.azure.core.tracing.opentelemetry.OpenTelemetryTracer")
+        .inject(InjectionMode.CLASS_ONLY);
+    injector
+        .proxyBuilder(
+            "io.opentelemetry.javaagent.instrumentation.azurecore.v1_40.shaded.com.azure.core.tracing.opentelemetry.OpenTelemetryTracerProvider")
+        .inject(InjectionMode.CLASS_ONLY);
+  }
+
+  @Override
+  public ElementMatcher.Junction<ClassLoader> classLoaderMatcher() {
+    // this class was introduced in azure-core 1.36
+    return hasClassesNamed("com.azure.core.util.tracing.TracerProvider")
+        // this class was introduced in azure-core 1.40
+        .and(hasClassesNamed("com.azure.core.implementation.jackson.ResponseErrorDeserializer"))
+        .and(not(hasClassesNamed("com.azure.core.tracing.opentelemetry.OpenTelemetryTracer")));
+  }
+
+  @Override
+  public List<TypeInstrumentation> typeInstrumentations() {
+    return asList(new EmptyTypeInstrumentation(), new AzureHttpClientInstrumentation());
+  }
+
+  public static class EmptyTypeInstrumentation implements TypeInstrumentation {
+    @Override
+    public ElementMatcher<TypeDescription> typeMatcher() {
+      return namedOneOf(
+          "com.azure.core.util.tracing.TracerProvider", "com.azure.core.util.tracing.Tracer");
+    }
+
+    @Override
+    public void transform(TypeTransformer transformer) {
+      // Nothing to instrument, no methods to match
+    }
+  }
+}