Add Azure SDK instrumentation

docs/supported-libraries.md

@@ -42,6 +42,7 @@ These are the supported libraries and frameworks:
 | [AsyncHttpClient](https://github.com/AsyncHttpClient/async-http-client)                                                           | 1.9+                           |
 | [AWS Lambda](https://docs.aws.amazon.com/lambda/latest/dg/java-handler.html)                                                      | 1.0+                           |
 | [AWS SDK](https://aws.amazon.com/sdk-for-java/)                                                                                   | 1.11.x and 2.2.0+              |
+| [Azure Core](https://docs.microsoft.com/en-us/java/api/overview/azure/core-readme)                                                | 1.14+                          |
 | [Cassandra Driver](https://github.com/datastax/java-driver)                                                                       | 3.0+                           |
 | [Couchbase Client](https://github.com/couchbase/couchbase-java-client)                                                            | 2.0+ and 3.1+                  |
 | [Dropwizard Views](https://www.dropwizard.io/en/latest/manual/views.html)                                                         | 0.7+                           |

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

@@ -0,0 +1,28 @@
+plugins {
+  id("otel.javaagent-instrumentation")
+}
+
+muzzle {
+  pass {
+    group.set("com.azure")
+    module.set("azure-core")
+    versions.set("[1.14.0,1.19.0)")
+    assertInverse.set(true)
+  }
+}
+
+sourceSets {
+  main {
+    val shadedDep = project(":instrumentation:azure-core:azure-core-1.14:library-instrumentation-shaded")
+    output.dir(shadedDep.file("build/extracted/shadow"), "builtBy" to ":instrumentation:azure-core:azure-core-1.14:library-instrumentation-shaded:extractShadowJar")
+  }
+}
+
+dependencies {
+  compileOnly(project(path = ":instrumentation:azure-core:azure-core-1.14:library-instrumentation-shaded", configuration = "shadow"))
+
+  library("com.azure:azure-core:1.14.0")
+
+  // Ensure no cross interference
+  testInstrumentation(project(":instrumentation:azure-core:azure-core-1.19:javaagent"))
+}

instrumentation/azure-core/azure-core-1.14/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/azurecore/v1_14/AzureHttpClientInstrumentation.java

@@ -0,0 +1,46 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.azurecore.v1_14;
+
+import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.implementsInterface;
+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.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() + "$SuppressNestedClientAdvice");
+  }
+
+  public static class SuppressNestedClientAdvice {
+
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void methodExit(@Advice.Return(readOnly = false) Mono<HttpResponse> mono) {
+      mono = new SuppressNestedClientMono<>(mono);
+    }
+  }
+}

instrumentation/azure-core/azure-core-1.14/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/azurecore/v1_14/AzureSdkInstrumentationModule.java

@@ -0,0 +1,76 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.azurecore.v1_14;
+
+import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasClassesNamed;
+import static java.util.Arrays.asList;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+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 java.util.List;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+@AutoService(InstrumentationModule.class)
+public class AzureSdkInstrumentationModule extends InstrumentationModule {
+  public AzureSdkInstrumentationModule() {
+    super("azure-core", "azure-core-1.14");
+  }
+
+  @Override
+  public void registerHelperResources(HelperResourceBuilder helperResourceBuilder) {
+    helperResourceBuilder.register(
+        "META-INF/services/com.azure.core.http.policy.AfterRetryPolicyProvider",
+        "azure-core-1.14/META-INF/services/com.azure.core.http.policy.AfterRetryPolicyProvider");
+    helperResourceBuilder.register(
+        "META-INF/services/com.azure.core.util.tracing.Tracer",
+        "azure-core-1.14/META-INF/services/com.azure.core.util.tracing.Tracer");
+  }
+
+  @Override
+  public ElementMatcher.Junction<ClassLoader> classLoaderMatcher() {
+    return hasClassesNamed("com.azure.core.util.tracing.Tracer")
+        // this is needed to prevent this instrumentation from being applied to azure-core 1.19+
+        .and(not(hasClassesNamed("com.azure.core.util.tracing.StartSpanOptions")))
+        .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() {
+      // we cannot use com.azure.core.http.policy.AfterRetryPolicyProvider
+      // or com.azure.core.util.tracing.Tracer here because we inject classes that implement these
+      // interfaces, causing the first one of these interfaces to be transformed to cause itself to
+      // be loaded (again), which leads to duplicate class definition error after the interface is
+      // transformed and the triggering class loader tries to load it.
+      //
+      // this is a list of all classes that call one of these:
+      // * ServiceLoader.load(AfterRetryPolicyProvider.class)
+      // * ServiceLoader.load(Tracer.class)
+      return named("com.azure.core.http.policy.HttpPolicyProviders")
+          .or(named("com.azure.core.util.tracing.TracerProxy"))
+          .or(named("com.azure.cosmos.CosmosAsyncClient"))
+          .or(named("com.azure.messaging.eventhubs.EventHubClientBuilder"))
+          .or(named("com.azure.messaging.eventhubs.EventProcessorClientBuilder"))
+          .or(named("com.azure.messaging.servicebus.ServiceBusClientBuilder"));
+    }
+
+    @Override
+    public void transform(TypeTransformer transformer) {
+      // Nothing to instrument, no methods to match
+    }
+  }
+}

instrumentation/azure-core/azure-core-1.14/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/azurecore/v1_14/SuppressNestedClientMono.java

@@ -0,0 +1,37 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.azurecore.v1_14;
+
+import static io.opentelemetry.javaagent.instrumentation.api.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 reactor.core.CoreSubscriber;
+import reactor.core.publisher.Mono;
+
+public class SuppressNestedClientMono<T> extends Mono<T> {
+
+  private final Mono<T> delegate;
+
+  public SuppressNestedClientMono(Mono<T> delegate) {
+    this.delegate = delegate;
+  }
+
+  @Override
+  public void subscribe(CoreSubscriber<? super T> actual) {
+    Context parentContext = currentContext();
+    if (SpanKey.HTTP_CLIENT.fromContextOrNull(parentContext) == null) {
+      try (Scope ignored =
+          SpanKey.HTTP_CLIENT.storeInContext(parentContext, Span.getInvalid()).makeCurrent()) {
+        delegate.subscribe(actual);
+      }
+    } else {
+      delegate.subscribe(actual);
+    }
+  }
+}

instrumentation/azure-core/azure-core-1.14/javaagent/src/main/resources/azure-core-1.14/META-INF/services/com.azure.core.http.policy.AfterRetryPolicyProvider

@@ -0,0 +1 @@
+io.opentelemetry.javaagent.instrumentation.azurecore.v1_14.shaded.com.azure.core.tracing.opentelemetry.OpenTelemetryHttpPolicy

instrumentation/azure-core/azure-core-1.14/javaagent/src/main/resources/azure-core-1.14/META-INF/services/com.azure.core.util.tracing.Tracer

@@ -0,0 +1 @@
+io.opentelemetry.javaagent.instrumentation.azurecore.v1_14.shaded.com.azure.core.tracing.opentelemetry.OpenTelemetryTracer

instrumentation/azure-core/azure-core-1.14/javaagent/src/test/groovy/AzureSdkTest.groovy

@@ -0,0 +1,43 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import com.azure.core.http.policy.HttpPolicyProviders
+import com.azure.core.util.Context
+import com.azure.core.util.tracing.TracerProxy
+import io.opentelemetry.api.trace.StatusCode
+import io.opentelemetry.instrumentation.test.AgentInstrumentationSpecification
+
+class AzureSdkTest extends AgentInstrumentationSpecification {
+
+  def "test helper classes injected"() {
+    expect:
+    TracerProxy.isTracingEnabled()
+
+    def list = new ArrayList()
+    HttpPolicyProviders.addAfterRetryPolicies(list)
+
+    list.size() == 1
+    list.get(0).getClass().getName() == "io.opentelemetry.javaagent.instrumentation.azurecore.v1_14.shaded" +
+      ".com.azure.core.tracing.opentelemetry.OpenTelemetryHttpPolicy"
+  }
+
+  def "test span"() {
+    when:
+    Context context = TracerProxy.start("hello", Context.NONE)
+    TracerProxy.end(200, null, context)
+
+    then:
+    assertTraces(1) {
+      trace(0, 1) {
+        span(0) {
+          name "hello"
+          status StatusCode.OK
+          attributes {
+          }
+        }
+      }
+    }
+  }
+}

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

@@ -0,0 +1,30 @@
+plugins {
+  id("com.github.johnrengelman.shadow")
+
+  id("otel.java-conventions")
+}
+
+group = "io.opentelemetry.javaagent.instrumentation"
+
+dependencies {
+  // this is the latest version that works with azure-core 1.14
+  implementation("com.azure:azure-core-tracing-opentelemetry:1.0.0-beta.12")
+}
+
+tasks {
+  shadowJar {
+    exclude("META-INF/services/*")
+
+    dependencies {
+      // including only azure-core-tracing-opentelemetry excludes its transitive dependencies
+      include(dependency("com.azure:azure-core-tracing-opentelemetry"))
+    }
+    relocate("com.azure.core.tracing.opentelemetry", "io.opentelemetry.javaagent.instrumentation.azurecore.v1_14.shaded.com.azure.core.tracing.opentelemetry")
+  }
+
+  val extractShadowJar by registering(Copy::class) {
+    dependsOn(shadowJar)
+    from(zipTree(shadowJar.get().archiveFile))
+    into("build/extracted/shadow")
+  }
+}

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

@@ -0,0 +1,28 @@
+plugins {
+  id("otel.javaagent-instrumentation")
+}
+
+muzzle {
+  pass {
+    group.set("com.azure")
+    module.set("azure-core")
+    versions.set("[1.19.0,)")
+    assertInverse.set(true)
+  }
+}
+
+sourceSets {
+  main {
+    val shadedDep = project(":instrumentation:azure-core:azure-core-1.19:library-instrumentation-shaded")
+    output.dir(shadedDep.file("build/extracted/shadow"), "builtBy" to ":instrumentation:azure-core:azure-core-1.19:library-instrumentation-shaded:extractShadowJar")
+  }
+}
+
+dependencies {
+  compileOnly(project(path = ":instrumentation:azure-core:azure-core-1.19:library-instrumentation-shaded", configuration = "shadow"))
+
+  library("com.azure:azure-core:1.19.0")
+
+  // Ensure no cross interference
+  testInstrumentation(project(":instrumentation:azure-core:azure-core-1.14:javaagent"))
+}

instrumentation/azure-core/azure-core-1.19/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/azurecore/v1_19/AzureSdkInstrumentationModule.java

@@ -0,0 +1,75 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.azurecore.v1_19;
+
+import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasClassesNamed;
+import static java.util.Collections.singletonList;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+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 java.util.List;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+@AutoService(InstrumentationModule.class)
+public class AzureSdkInstrumentationModule extends InstrumentationModule {
+  public AzureSdkInstrumentationModule() {
+    super("azure-core", "azure-core-1.19");
+  }
+
+  @Override
+  public void registerHelperResources(HelperResourceBuilder helperResourceBuilder) {
+    helperResourceBuilder.register(
+        "META-INF/services/com.azure.core.http.policy.AfterRetryPolicyProvider",
+        "azure-core-1.19/META-INF/services/com.azure.core.http.policy.AfterRetryPolicyProvider");
+    helperResourceBuilder.register(
+        "META-INF/services/com.azure.core.util.tracing.Tracer",
+        "azure-core-1.19/META-INF/services/com.azure.core.util.tracing.Tracer");
+  }
+
+  @Override
+  public ElementMatcher.Junction<ClassLoader> classLoaderMatcher() {
+    // this class was introduced in azure-core 1.19
+    return hasClassesNamed("com.azure.core.util.tracing.StartSpanOptions")
+        .and(not(hasClassesNamed("com.azure.core.tracing.opentelemetry.OpenTelemetryTracer")));
+  }
+
+  @Override
+  public List<TypeInstrumentation> typeInstrumentations() {
+    return singletonList(new EmptyTypeInstrumentation());
+  }
+
+  public static class EmptyTypeInstrumentation implements TypeInstrumentation {
+    @Override
+    public ElementMatcher<TypeDescription> typeMatcher() {
+      // we cannot use com.azure.core.http.policy.AfterRetryPolicyProvider
+      // or com.azure.core.util.tracing.Tracer here because we inject classes that implement these
+      // interfaces, causing the first one of these interfaces to be transformed to cause itself to
+      // be loaded (again), which leads to duplicate class definition error after the interface is
+      // transformed and the triggering class loader tries to load it.
+      //
+      // this is a list of all classes that call one of these:
+      // * ServiceLoader.load(AfterRetryPolicyProvider.class)
+      // * ServiceLoader.load(Tracer.class)
+      return named("com.azure.core.http.policy.HttpPolicyProviders")
+          .or(named("com.azure.core.util.tracing.TracerProxy"))
+          .or(named("com.azure.cosmos.CosmosAsyncClient"))
+          .or(named("com.azure.messaging.eventhubs.EventHubClientBuilder"))
+          .or(named("com.azure.messaging.eventhubs.EventProcessorClientBuilder"))
+          .or(named("com.azure.messaging.servicebus.ServiceBusClientBuilder"));
+    }
+
+    @Override
+    public void transform(TypeTransformer transformer) {
+      // Nothing to instrument, no methods to match
+    }
+  }
+}

instrumentation/azure-core/azure-core-1.19/javaagent/src/main/resources/azure-core-1.19/META-INF/services/com.azure.core.http.policy.AfterRetryPolicyProvider

@@ -0,0 +1 @@
+io.opentelemetry.javaagent.instrumentation.azurecore.v1_19.shaded.com.azure.core.tracing.opentelemetry.OpenTelemetryHttpPolicy

instrumentation/azure-core/azure-core-1.19/javaagent/src/main/resources/azure-core-1.19/META-INF/services/com.azure.core.util.tracing.Tracer

@@ -0,0 +1 @@
+io.opentelemetry.javaagent.instrumentation.azurecore.v1_19.shaded.com.azure.core.tracing.opentelemetry.OpenTelemetryTracer