Fix azure-core instrumentation, including indy

Author: trask

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

@@ -8,6 +8,9 @@ muzzle {
     module.set("azure-core")
     versions.set("[1.53.0,)")
     assertInverse.set(true)
+    // our advice helper bridges an explicitly supplied application parent context to the agent
+    // context, so it references io.opentelemetry.context.{Context,Scope}
+    extraDependency("io.opentelemetry:opentelemetry-api:1.27.0")
   }
 }
 
@@ -24,6 +27,10 @@ sourceSets {
 dependencies {
   compileOnly(project(":instrumentation:azure-core:azure-core-1.53:library-instrumentation-shaded", configuration = "shadow"))
 
+  // needed to bridge an explicitly supplied application parent context (the unshaded
+  // "application.io.opentelemetry.*" types) to the agent context inside our advice
+  compileOnly(project(":opentelemetry-api-shaded-for-instrumenting", configuration = "shadow"))
+
   library("com.azure:azure-core:1.53.0")
 
   // Ensure no cross interference

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

@@ -0,0 +1,51 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.azurecore.v1_53;
+
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import java.util.Optional;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.asm.Advice.AssignReturned;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+/**
+ * Bridges an explicitly supplied parent context for {@code azure-core-tracing-opentelemetry}.
+ *
+ * <p>When a user passes a parent context to an Azure SDK call, the value is stored on the {@link
+ * com.azure.core.util.Context} under the {@code "trace-context"} key as the application's (unshaded)
+ * {@code io.opentelemetry.context.Context}. The bundled {@code OpenTelemetryTracer} reads that value
+ * back and expects it to be the agent's (shaded) context. Convert it here so the tracer does not
+ * need to reach into agent internals reflectively.
+ */
+class AzureContextInstrumentation implements TypeInstrumentation {
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return named("com.azure.core.util.Context");
+  }
+
+  @Override
+  public void transform(TypeTransformer transformer) {
+    transformer.applyAdviceToMethod(
+        named("getData").and(takesArguments(1)),
+        getClass().getName() + "$GetDataAdvice");
+  }
+
+  @SuppressWarnings("unused")
+  public static class GetDataAdvice {
+    @AssignReturned.ToReturned
+    @Advice.OnMethodExit(suppress = Throwable.class, inline = false)
+    public static Optional<Object> onExit(
+        @Advice.Argument(0) Object key, @Advice.Return Optional<Object> data) {
+      return AzureExplicitParentContextHelper.bridgeApplicationContext(key, data);
+    }
+  }
+}

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

@@ -0,0 +1,43 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.azurecore.v1_53;
+
+import io.opentelemetry.context.Context;
+import java.util.Optional;
+
+/**
+ * Converts an explicitly supplied application parent context into the agent context.
+ *
+ * <p>In the agent the application's {@code io.opentelemetry.context.Context} (referenced here as
+ * {@code application.io.opentelemetry.context.Context}) is bridged to the agent context by the
+ * opentelemetry-api instrumentation. Making the application context current and reading back the
+ * agent context performs that conversion using only public API, so this works in both inline and
+ * indy mode without reaching into agent-internal helper classes.
+ */
+public final class AzureExplicitParentContextHelper {
+
+  // azure-core-tracing-opentelemetry stores the explicit parent context under this key
+  private static final String PARENT_TRACE_CONTEXT_KEY = "trace-context";
+
+  public static Optional<Object> bridgeApplicationContext(Object key, Optional<Object> data) {
+    if (!PARENT_TRACE_CONTEXT_KEY.equals(key) || data == null || !data.isPresent()) {
+      return data;
+    }
+    Object value = data.get();
+    if (!(value instanceof application.io.opentelemetry.context.Context)) {
+      return data;
+    }
+    application.io.opentelemetry.context.Context applicationContext =
+        (application.io.opentelemetry.context.Context) value;
+    Context agentContext;
+    try (application.io.opentelemetry.context.Scope ignored = applicationContext.makeCurrent()) {
+      agentContext = Context.current();
+    }
+    return Optional.of(agentContext);
+  }
+
+  private AzureExplicitParentContextHelper() {}
+}

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

@@ -57,7 +57,10 @@ public ElementMatcher.Junction<ClassLoader> classLoaderMatcher() {
 
   @Override
   public List<TypeInstrumentation> typeInstrumentations() {
-    return asList(new EmptyTypeInstrumentation(), new AzureHttpClientInstrumentation());
+    return asList(
+        new EmptyTypeInstrumentation(),
+        new AzureContextInstrumentation(),
+        new AzureHttpClientInstrumentation());
   }
 
   private static class EmptyTypeInstrumentation implements TypeInstrumentation {

instrumentation/azure-core/azure-core-1.53/javaagent/src/testAzure/java/io/opentelemetry/javaagent/instrumentation/azurecore/v1_53/AzureSdkTest.java

@@ -32,6 +32,8 @@
 import com.azure.core.util.TracingOptions;
 import com.azure.core.util.tracing.Tracer;
 import com.azure.core.util.tracing.TracerProvider;
+import io.opentelemetry.api.GlobalOpenTelemetry;
+import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.api.trace.SpanKind;
 import io.opentelemetry.instrumentation.api.internal.SpanKey;
 import io.opentelemetry.instrumentation.testing.junit.AgentInstrumentationExtension;
@@ -78,6 +80,34 @@ void testSpan() {
                             equalTo(stringKey("az.namespace"), "otel.tests"))));
   }
 
+  @Test
+  void testExplicitParentContextBridge() {
+    // Azure's bundled OpenTelemetryTracer reflectively calls
+    // io.opentelemetry.javaagent.instrumentation.opentelemetryapi.context.AgentContextStorage
+    //   #getAgentContext(application io.opentelemetry.context.Context)
+    // to convert an explicitly-supplied application context into the agent (shaded) context.
+    // This test exercises that path: the parent span is supplied only via the Azure context bag
+    // under PARENT_TRACE_CONTEXT_KEY and is never made current, so the only way the child can be
+    // linked to it is through the reflective bridge.
+    Tracer azTracer = createAzTracer();
+
+    Span parentSpan = GlobalOpenTelemetry.getTracer("test").spanBuilder("parent").startSpan();
+    // application (unshaded) context carrying the parent span, NOT made current
+    io.opentelemetry.context.Context parentContext =
+        io.opentelemetry.context.Context.root().with(parentSpan);
+
+    Context azContext = new Context(Tracer.PARENT_TRACE_CONTEXT_KEY, parentContext);
+    Context child = azTracer.start("child", azContext);
+    azTracer.end(null, null, child);
+    parentSpan.end();
+
+    testing.waitAndAssertTracesWithoutScopeVersionVerification(
+        trace ->
+            trace.hasSpansSatisfyingExactly(
+                span -> span.hasName("parent").hasNoParent(),
+                span -> span.hasName("child").hasParent(trace.getSpan(0))));
+  }
+
   @Test
   void testPipelineAndSuppression() {
     AtomicBoolean hasClientAndHttpKeys = new AtomicBoolean(false);