Support latest Azure SDK shim (#15068)

Co-authored-by: otelbot <[email protected]>

Author: trask

.fossa.yml

@@ -388,6 +388,9 @@ targets:
     - type: gradle
       path: ./
       target: ':instrumentation:azure-core:azure-core-1.36:javaagent'
+    - type: gradle
+      path: ./
+      target: ':instrumentation:azure-core:azure-core-1.53:javaagent'
     - type: gradle
       path: ./
       target: ':instrumentation:cassandra:cassandra-3.0:javaagent'

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.53: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.53: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.53.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.53:javaagent"))
 }
 
 val latestDepTest = findProperty("testLatestDeps") as Boolean
@@ -46,11 +47,11 @@ testing {
     val testAzure by registering(JvmTestSuite::class) {
       dependencies {
         if (latestDepTest) {
-          implementation("com.azure:azure-core:latest.release")
-          implementation("com.azure:azure-core-test:latest.release")
+          implementation("com.azure:azure-core:1.52.0")
+          implementation("com.azure:azure-core-test:1.26.2")
         } else {
           implementation("com.azure:azure-core:1.36.0")
-          implementation("com.azure:azure-core-test:1.16.2")
+          implementation("com.azure:azure-core-test:1.14.1")
         }
       }
     }

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

@@ -57,6 +57,7 @@ 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")
+        .and(not(hasClassesNamed("com.azure.core.util.LibraryTelemetryOptions")))
         .and(not(hasClassesNamed("com.azure.core.tracing.opentelemetry.OpenTelemetryTracer")));
   }
 

instrumentation/azure-core/azure-core-1.36/library-instrumentation-shaded/build.gradle.kts

@@ -6,8 +6,10 @@ plugins {
 group = "io.opentelemetry.javaagent.instrumentation"
 
 dependencies {
-  // this is the last good version that works with indy build
-  // update to 1.49 or latest once https://github.com/Azure/azure-sdk-for-java/pull/42586 is released.
+  // the latest version that works with azure-core 1.36.0 is 1.0.0-beta.49
+  // but the latest version that works with indy build is 1.0.0-beta.45
+  // (indy build was fixed by https://github.com/Azure/azure-sdk-for-java/pull/42586
+  // in 1.0.0-beta.51)
   implementation("com.azure:azure-core-tracing-opentelemetry:1.0.0-beta.45")
 }
 

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

@@ -0,0 +1,65 @@
+plugins {
+  id("otel.javaagent-instrumentation")
+}
+
+muzzle {
+  pass {
+    group.set("com.azure")
+    module.set("azure-core")
+    versions.set("[1.53.0,)")
+    assertInverse.set(true)
+  }
+}
+
+sourceSets {
+  main {
+    val shadedDep = project(":instrumentation:azure-core:azure-core-1.53:library-instrumentation-shaded")
+    output.dir(
+      shadedDep.file("build/extracted/shadow"),
+      "builtBy" to ":instrumentation:azure-core:azure-core-1.53:library-instrumentation-shaded:extractShadowJar"
+    )
+  }
+}
+
+dependencies {
+  compileOnly(project(":instrumentation:azure-core:azure-core-1.53:library-instrumentation-shaded", configuration = "shadow"))
+
+  library("com.azure:azure-core:1.53.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
+
+tasks {
+  withType<Test>().configureEach {
+    systemProperty("testLatestDeps", 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:latest.release")
+          implementation("com.azure:azure-core-test:latest.release")
+        } else {
+          implementation("com.azure:azure-core:1.53.0")
+          implementation("com.azure:azure-core-test:1.26.2")
+        }
+      }
+    }
+  }
+}
+
+tasks {
+  check {
+    dependsOn(testing.suites)
+  }
+}

instrumentation/azure-core/azure-core-1.53/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/azurecore/v1_53/AzureHttpClientInstrumentation.java

@@ -0,0 +1,80 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.azurecore.v1_53;
+
+import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.implementsInterface;
+import static io.opentelemetry.javaagent.instrumentation.azurecore.v1_53.SuppressNestedClientHelper.disallowNestedClientSpanMono;
+import static io.opentelemetry.javaagent.instrumentation.azurecore.v1_53.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 static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+
+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 javax.annotation.Nullable;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.asm.Advice.AssignReturned;
+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(takesArgument(1, named("com.azure.core.util.Context")))
+            .and(returns(named("reactor.core.publisher.Mono"))),
+        this.getClass().getName() + "$SuppressNestedClientMonoAdvice");
+    transformer.applyAdviceToMethod(
+        isMethod()
+            .and(isPublic())
+            .and(named("sendSync"))
+            .and(takesArgument(1, named("com.azure.core.util.Context")))
+            .and(returns(named("com.azure.core.http.HttpResponse"))),
+        this.getClass().getName() + "$SuppressNestedClientSyncAdvice");
+  }
+
+  @SuppressWarnings("unused")
+  public static class SuppressNestedClientMonoAdvice {
+    @AssignReturned.ToReturned
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static Mono<HttpResponse> asyncSendExit(
+        @Advice.Argument(1) com.azure.core.util.Context azContext,
+        @Advice.Return Mono<HttpResponse> mono) {
+      return disallowNestedClientSpanMono(mono, azContext);
+    }
+  }
+
+  @SuppressWarnings("unused")
+  public static class SuppressNestedClientSyncAdvice {
+
+    @Nullable
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static Scope syncSendEnter(@Advice.Argument(1) com.azure.core.util.Context azContext) {
+      return disallowNestedClientSpanSync(azContext);
+    }
+
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void syncSendExit(@Advice.Enter @Nullable Scope scope) {
+      if (scope != null) {
+        scope.close();
+      }
+    }
+  }
+}

instrumentation/azure-core/azure-core-1.53/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/azurecore/v1_53/AzureSdkInstrumentationModule.java

@@ -0,0 +1,85 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.azurecore.v1_53;
+
+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.53");
+  }
+
+  @Override
+  public void registerHelperResources(HelperResourceBuilder helperResourceBuilder) {
+    helperResourceBuilder.register(
+        "META-INF/services/com.azure.core.util.tracing.TracerProvider",
+        "azure-core-1.53/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.53/META-INF/services/com.azure.core.util.tracing.Tracer");
+  }
+
+  @Override
+  public void injectClasses(ClassInjector injector) {
+    injector
+        .proxyBuilder(
+            "io.opentelemetry.javaagent.instrumentation.azurecore.v1_53.shaded.com.azure.core.tracing.opentelemetry.OpenTelemetryTracer")
+        .inject(InjectionMode.CLASS_ONLY);
+    injector
+        .proxyBuilder(
+            "io.opentelemetry.javaagent.instrumentation.azurecore.v1_53.shaded.com.azure.core.tracing.opentelemetry.OpenTelemetryTracerProvider")
+        .inject(InjectionMode.CLASS_ONLY);
+  }
+
+  @Override
+  public ElementMatcher.Junction<ClassLoader> classLoaderMatcher() {
+    // LibraryTelemetryOptions was introduced in azure-core 1.53
+    return hasClassesNamed("com.azure.core.util.LibraryTelemetryOptions")
+        .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
+    }
+  }
+
+  @Override
+  public boolean isIndyReady() {
+    return true;
+  }
+}

instrumentation/azure-core/azure-core-1.53/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/azurecore/v1_53/SuppressNestedClientHelper.java

@@ -0,0 +1,61 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.azurecore.v1_53;
+
+import static io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge.currentContext;
+
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.Scope;
+import io.opentelemetry.instrumentation.api.internal.SpanKey;
+import javax.annotation.Nullable;
+import reactor.core.CoreSubscriber;
+import reactor.core.publisher.Mono;
+
+public class SuppressNestedClientHelper {
+
+  @Nullable
+  public static Scope disallowNestedClientSpanSync(com.azure.core.util.Context azContext) {
+    Context parentContext = currentContext();
+    boolean hasAzureClientSpan = azContext.getData("client-method-call-flag").isPresent();
+    if (doesNotHaveClientSpan(parentContext) && hasAzureClientSpan) {
+      return disallowNestedClientSpan(parentContext).makeCurrent();
+    }
+    return null;
+  }
+
+  public static <T> Mono<T> disallowNestedClientSpanMono(
+      Mono<T> delegate, com.azure.core.util.Context azContext) {
+    return new Mono<T>() {
+      @Override
+      public void subscribe(CoreSubscriber<? super T> coreSubscriber) {
+        Context parentContext = currentContext();
+
+        boolean hasAzureClientSpan = azContext.getData("client-method-call-flag").isPresent();
+        if (doesNotHaveClientSpan(parentContext) && hasAzureClientSpan) {
+          try (Scope ignored = disallowNestedClientSpan(parentContext).makeCurrent()) {
+            delegate.subscribe(coreSubscriber);
+          }
+        } else {
+          delegate.subscribe(coreSubscriber);
+        }
+      }
+    };
+  }
+
+  private static boolean doesNotHaveClientSpan(Context parentContext) {
+    return SpanKey.KIND_CLIENT.fromContextOrNull(parentContext) == null
+        && SpanKey.HTTP_CLIENT.fromContextOrNull(parentContext) == null;
+  }
+
+  private static Context disallowNestedClientSpan(Context parentContext) {
+    Span span = Span.getInvalid();
+    return SpanKey.HTTP_CLIENT.storeInContext(
+        SpanKey.KIND_CLIENT.storeInContext(parentContext, span), span);
+  }
+
+  private SuppressNestedClientHelper() {}
+}