Fix Azure Functions instrumentation (#2684)

The fix was just un-inlining `generateCustomDimensions()`. But this pointed out our lack of basic test for this instrumentation, so added some stub classes to test the instrumentation (since azure functions worker classes are not on maven central). And using these stub classes also allows us to remove the reflection.
microsoft · Nov 11, 2022 · 18280e1 · 18280e1
1 parent fcdf03d
commit 18280e1
Show file tree

Hide file tree

Showing 12 changed files with 219 additions and 88 deletions.
diff --git a/...java/com/microsoft/applicationinsights/agent/internal/init/AzureFunctionsInitializer.java b/...java/com/microsoft/applicationinsights/agent/internal/init/AzureFunctionsInitializer.java
@@ -152,9 +152,6 @@ static void setSelfDiagnosticsLevel(@Nullable String loggingLevel) {
     loggerList.forEach(configurator::updateLoggerLevel);
   }
 
-  // since the agent is already running at this point, this really just determines whether the
-  // telemetry is sent to the ingestion service or not (essentially behaving to the user as if the
-  // agent is not enabled)
   static boolean isAgentEnabled() {
     String enableAgent = System.getenv("APPLICATIONINSIGHTS_ENABLE_AGENT");
     boolean enableAgentDefault = Boolean.getBoolean("LazySetOptIn");

diff --git a/agent/instrumentation/azure-functions-worker-stub/build.gradle.kts b/agent/instrumentation/azure-functions-worker-stub/build.gradle.kts
@@ -0,0 +1,5 @@
+plugins {
+  id("ai.java-conventions")
+}
+
+// this module is needed since the azure functions worker artifact is not available in maven central
diff --git a/...rker-stub/src/main/java/com/microsoft/azure/functions/rpc/messages/InvocationRequest.java b/...rker-stub/src/main/java/com/microsoft/azure/functions/rpc/messages/InvocationRequest.java
@@ -0,0 +1,15 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.microsoft.azure.functions.rpc.messages;
+
+public class InvocationRequest {
+
+  public RpcTraceContext getTraceContext() {
+    throw new UnsupportedOperationException();
+  }
+
+  public String getInvocationId() {
+    throw new UnsupportedOperationException();
+  }
+}
diff --git a/...worker-stub/src/main/java/com/microsoft/azure/functions/rpc/messages/RpcTraceContext.java b/...worker-stub/src/main/java/com/microsoft/azure/functions/rpc/messages/RpcTraceContext.java
@@ -0,0 +1,21 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.microsoft.azure.functions.rpc.messages;
+
+import java.util.Map;
+
+public class RpcTraceContext {
+
+  public Map<String, String> getAttributesMap() {
+    throw new UnsupportedOperationException();
+  }
+
+  public String getTraceParent() {
+    throw new UnsupportedOperationException();
+  }
+
+  public String getTraceState() {
+    throw new UnsupportedOperationException();
+  }
+}
diff --git a/.../src/main/java/com/microsoft/azure/functions/worker/handler/InvocationRequestHandler.java b/.../src/main/java/com/microsoft/azure/functions/worker/handler/InvocationRequestHandler.java
@@ -0,0 +1,17 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.microsoft.azure.functions.worker.handler;
+
+import com.microsoft.azure.functions.rpc.messages.InvocationRequest;
+
+public class InvocationRequestHandler {
+
+  public void execute(InvocationRequest request) {
+    verifyCurrentContext();
+  }
+
+  // this doesn't exist in the real worker artifact
+  // it only exists for testing the instrumentation
+  protected void verifyCurrentContext() {}
+}
diff --git a/agent/instrumentation/azure-functions/build.gradle.kts b/agent/instrumentation/azure-functions/build.gradle.kts
@@ -12,6 +12,10 @@ muzzle {
 val otelInstrumentationAlphaVersion: String by project
 
 dependencies {
+  compileOnly(project(":agent:instrumentation:azure-functions-worker-stub"))
+
+  testImplementation(project(":agent:instrumentation:azure-functions-worker-stub"))
+
   // TODO remove when start using io.opentelemetry.instrumentation.javaagent-instrumentation plugin
   add("codegen", "io.opentelemetry.javaagent:opentelemetry-javaagent-tooling:$otelInstrumentationAlphaVersion")
   add("muzzleBootstrap", "io.opentelemetry.instrumentation:opentelemetry-instrumentation-annotations-support:$otelInstrumentationAlphaVersion")

diff --git a/.../io/opentelemetry/javaagent/instrumentation/azurefunctions/InvocationInstrumentation.java b/.../io/opentelemetry/javaagent/instrumentation/azurefunctions/InvocationInstrumentation.java
@@ -11,6 +11,8 @@
 import com.microsoft.applicationinsights.agent.bootstrap.AzureFunctions;
 import com.microsoft.applicationinsights.agent.bootstrap.AzureFunctionsCustomDimensions;
 import com.microsoft.applicationinsights.agent.bootstrap.BytecodeUtil;
+import com.microsoft.azure.functions.rpc.messages.InvocationRequest;
+import com.microsoft.azure.functions.rpc.messages.RpcTraceContext;
 import io.opentelemetry.api.GlobalOpenTelemetry;
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.api.trace.SpanContext;
@@ -43,23 +45,17 @@ public void transform(TypeTransformer transformer) {
         InvocationInstrumentation.class.getName() + "$ExecuteAdvice");
   }
 
-  @SuppressWarnings({
-    "unused",
-    "PrivateConstructorForUtilityClass",
-    "MustBeClosedChecker",
-    "unchecked"
-  })
+  @SuppressWarnings({"unused", "PrivateConstructorForUtilityClass", "MustBeClosedChecker"})
   public static class ExecuteAdvice {
     @Nullable
     @Advice.OnMethodEnter(suppress = Throwable.class)
-    public static Scope methodEnter(@Advice.Argument(0) Object request)
-        throws ReflectiveOperationException {
+    public static Scope methodEnter(@Advice.Argument(0) InvocationRequest request) {
 
       if (!AzureFunctions.hasConnectionString()) {
         return null;
       }
 
-      Object traceContext = InvocationRequestExtractAdapter.getTraceContextMethod.invoke(request);
+      RpcTraceContext traceContext = request.getTraceContext();
       Context extractedContext =
           GlobalOpenTelemetry.getPropagators()
               .getTextMapPropagator()
@@ -78,10 +74,17 @@ public static Scope methodEnter(@Advice.Argument(0) Object request)
               traceFlags,
               spanContext.getTraceState());
 
-      return Context.current()
-          .with(Span.wrap(spanContext))
-          .with(generateCustomDimensions(request, traceContext))
-          .makeCurrent();
+      Map<String, String> attributesMap = traceContext.getAttributesMap();
+      AzureFunctionsCustomDimensions customDimensions =
+          new AzureFunctionsCustomDimensions(
+              request.getInvocationId(),
+              attributesMap.get("ProcessId"),
+              attributesMap.get("LogLevel"),
+              attributesMap.get("Category"),
+              attributesMap.get("HostInstanceId"),
+              attributesMap.get("#AzFuncLiveLogsSessionId"));
+
+      return Context.current().with(Span.wrap(spanContext)).with(customDimensions).makeCurrent();
     }
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
@@ -90,21 +93,5 @@ public static void methodExit(@Advice.Enter @Nullable Scope scope) {
         scope.close();
       }
     }
-
-    private static AzureFunctionsCustomDimensions generateCustomDimensions(
-        Object request, Object traceContext) throws ReflectiveOperationException {
-      String invocationId =
-          (String) InvocationRequestExtractAdapter.getInvocationId.invoke(request);
-      Map<String, String> attributesMap =
-          (Map<String, String>)
-              InvocationRequestExtractAdapter.getAttributesMap.invoke(traceContext);
-      return new AzureFunctionsCustomDimensions(
-          invocationId,
-          attributesMap.get("ProcessId"),
-          attributesMap.get("LogLevel"),
-          attributesMap.get("Category"),
-          attributesMap.get("HostInstanceId"),
-          attributesMap.get("#AzFuncLiveLogsSessionId"));
-    }
   }
 }
diff --git a/...entelemetry/javaagent/instrumentation/azurefunctions/InvocationRequestExtractAdapter.java b/...entelemetry/javaagent/instrumentation/azurefunctions/InvocationRequestExtractAdapter.java
@@ -3,75 +3,31 @@
 
 package io.opentelemetry.javaagent.instrumentation.azurefunctions;
 
+import com.microsoft.azure.functions.rpc.messages.RpcTraceContext;
 import io.opentelemetry.context.propagation.TextMapGetter;
-import java.lang.reflect.Method;
-import java.util.logging.Level;
-import java.util.logging.Logger;
 import javax.annotation.Nullable;
 
-// using reflection because these classes are not published to maven central that we can easily
-// compile against
-public class InvocationRequestExtractAdapter implements TextMapGetter<Object> {
-
-  private static final Logger logger =
-      Logger.getLogger(InvocationRequestExtractAdapter.class.getName());
+public class InvocationRequestExtractAdapter implements TextMapGetter<RpcTraceContext> {
 
   public static final InvocationRequestExtractAdapter GETTER =
       new InvocationRequestExtractAdapter();
 
-  public static final Method getTraceContextMethod;
-  public static final Method getInvocationId;
-  public static final Method getAttributesMap;
-  private static final Method getTraceParentMethod;
-  private static final Method getTraceStateMethod;
-
-  static {
-    Method getTraceContextMethodLocal = null;
-    Method getInvocationIdLocal = null;
-    Method getAttributesMapLocal = null;
-    Method getTraceParentMethodLocal = null;
-    Method getTraceStateMethodLocal = null;
-    try {
-      Class<?> invocationRequestClass =
-          Class.forName("com.microsoft.azure.functions.rpc.messages.InvocationRequest");
-      Class<?> rpcTraceContextClass =
-          Class.forName("com.microsoft.azure.functions.rpc.messages.RpcTraceContext");
-      getTraceContextMethodLocal = invocationRequestClass.getMethod("getTraceContext");
-      getInvocationIdLocal = invocationRequestClass.getMethod("getInvocationId");
-      getAttributesMapLocal = rpcTraceContextClass.getMethod("getAttributesMap");
-      getTraceParentMethodLocal = rpcTraceContextClass.getMethod("getTraceParent");
-      getTraceStateMethodLocal = rpcTraceContextClass.getMethod("getTraceState");
-    } catch (ReflectiveOperationException e) {
-      logger.log(Level.SEVERE, e.getMessage(), e);
-    }
-    getTraceContextMethod = getTraceContextMethodLocal;
-    getInvocationId = getInvocationIdLocal;
-    getAttributesMap = getAttributesMapLocal;
-    getTraceParentMethod = getTraceParentMethodLocal;
-    getTraceStateMethod = getTraceStateMethodLocal;
-  }
-
   @Override
-  public Iterable<String> keys(Object carrier) {
+  public Iterable<String> keys(RpcTraceContext carrier) {
     return null;
   }
 
   @Override
   @Nullable
-  public String get(Object carrier, String key) {
-    try {
-      // only supports W3C propagator
-      switch (key) {
-        case "traceparent":
-          return (String) getTraceParentMethod.invoke(carrier);
-        case "tracestate":
-          return (String) getTraceStateMethod.invoke(carrier);
-        default:
-          return null;
-      }
-    } catch (ReflectiveOperationException e) {
-      logger.log(Level.FINE, e.getMessage(), e);
-      return null;
+  public String get(RpcTraceContext carrier, String key) {
+    // only supports W3C propagator
+    switch (key) {
+      case "traceparent":
+        return carrier.getTraceParent();
+      case "tracestate":
+        return carrier.getTraceState();
+      default:
+        return null;
     }
   }
 }
diff --git a/agent/instrumentation/azure-functions/src/test/java/AzureFunctionsTest.java b/agent/instrumentation/azure-functions/src/test/java/AzureFunctionsTest.java
@@ -0,0 +1,69 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+import static java.util.Collections.emptyMap;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import com.microsoft.applicationinsights.agent.bootstrap.AzureFunctions;
+import com.microsoft.applicationinsights.agent.bootstrap.BytecodeUtil;
+import com.microsoft.applicationinsights.agent.bootstrap.BytecodeUtil.BytecodeUtilDelegate;
+import com.microsoft.azure.functions.rpc.messages.InvocationRequest;
+import com.microsoft.azure.functions.rpc.messages.RpcTraceContext;
+import com.microsoft.azure.functions.worker.handler.InvocationRequestHandler;
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.SpanContext;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.instrumentation.testing.junit.AgentInstrumentationExtension;
+import io.opentelemetry.instrumentation.testing.junit.InstrumentationExtension;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicReference;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+class AzureFunctionsTest {
+
+  @RegisterExtension
+  static final InstrumentationExtension testing = AgentInstrumentationExtension.create();
+
+  static {
+    // this is needed since currently tests are run against otel javaagent, and not ai javaagent
+    AzureFunctions.setup(() -> true, () -> {});
+    BytecodeUtilDelegate delegate = mock(BytecodeUtilDelegate.class);
+    when(delegate.shouldSample(anyString())).thenReturn(true);
+    BytecodeUtil.setDelegate(delegate);
+  }
+
+  @Test
+  void setRequestProperty() {
+    // given
+    String traceParent = "00-11111111111111111111111111111111-1111111111111111-00";
+    String traceState = null;
+    Map<String, String> attributesMap = emptyMap();
+    RpcTraceContext traceContext = new MockRpcTraceContext(traceParent, traceState, attributesMap);
+
+    String invocationId = null;
+    InvocationRequest request = new MockInvocationRequest(traceContext, invocationId);
+
+    AtomicReference<Context> contextRef = new AtomicReference<>();
+    InvocationRequestHandler handler =
+        new InvocationRequestHandler() {
+          @Override
+          protected void verifyCurrentContext() {
+            contextRef.set(Context.current());
+          }
+        };
+
+    // when
+    handler.execute(request);
+
+    // then
+    Context context = contextRef.get();
+    SpanContext spanContext = Span.fromContext(context).getSpanContext();
+    assertThat(spanContext.getTraceId()).isEqualTo("11111111111111111111111111111111");
+    assertThat(spanContext.getSpanId()).isEqualTo("1111111111111111");
+    assertThat(spanContext.getTraceFlags().isSampled()).isTrue();
+  }
+}
diff --git a/agent/instrumentation/azure-functions/src/test/java/MockInvocationRequest.java b/agent/instrumentation/azure-functions/src/test/java/MockInvocationRequest.java
@@ -0,0 +1,26 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+import com.microsoft.azure.functions.rpc.messages.InvocationRequest;
+import com.microsoft.azure.functions.rpc.messages.RpcTraceContext;
+
+class MockInvocationRequest extends InvocationRequest {
+
+  private final RpcTraceContext traceContext;
+  private final String invocationId;
+
+  MockInvocationRequest(RpcTraceContext traceContext, String invocationId) {
+    this.traceContext = traceContext;
+    this.invocationId = invocationId;
+  }
+
+  @Override
+  public RpcTraceContext getTraceContext() {
+    return traceContext;
+  }
+
+  @Override
+  public String getInvocationId() {
+    return invocationId;
+  }
+}
diff --git a/agent/instrumentation/azure-functions/src/test/java/MockRpcTraceContext.java b/agent/instrumentation/azure-functions/src/test/java/MockRpcTraceContext.java
@@ -0,0 +1,33 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+import com.microsoft.azure.functions.rpc.messages.RpcTraceContext;
+import java.util.Map;
+
+class MockRpcTraceContext extends RpcTraceContext {
+
+  private final String traceParent;
+  private final String traceState;
+  private final Map<String, String> attributesMap;
+
+  MockRpcTraceContext(String traceParent, String traceState, Map<String, String> attributesMap) {
+    this.traceParent = traceParent;
+    this.traceState = traceState;
+    this.attributesMap = attributesMap;
+  }
+
+  @Override
+  public String getTraceParent() {
+    return traceParent;
+  }
+
+  @Override
+  public String getTraceState() {
+    return traceState;
+  }
+
+  @Override
+  public Map<String, String> getAttributesMap() {
+    return attributesMap;
+  }
+}
diff --git a/settings.gradle.kts b/settings.gradle.kts
@@ -59,6 +59,7 @@ include(":agent:azure-monitor-exporter")
 include(":agent:agent-for-testing")
 include(":agent:instrumentation:applicationinsights-web-2.3")
 include(":agent:instrumentation:azure-functions")
+include(":agent:instrumentation:azure-functions-worker-stub")
 include(":agent:instrumentation:methods")
 include(":agent:instrumentation:micrometer-1.0")
 include(":agent:agent")