Skip to content

Commit 18280e1

Browse files
trasktrask
authored
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.
1 parent fcdf03d commit 18280e1

File tree

12 files changed

+219
-88
lines changed

12 files changed

+219
-88
lines changed

‎agent/agent-tooling/src/main/java/com/microsoft/applicationinsights/agent/internal/init/AzureFunctionsInitializer.java‎

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -152,9 +152,6 @@ static void setSelfDiagnosticsLevel(@Nullable String loggingLevel) {
152152
loggerList.forEach(configurator::updateLoggerLevel);
153153
}
154154

155-
// since the agent is already running at this point, this really just determines whether the
156-
// telemetry is sent to the ingestion service or not (essentially behaving to the user as if the
157-
// agent is not enabled)
158155
static boolean isAgentEnabled() {
159156
String enableAgent = System.getenv("APPLICATIONINSIGHTS_ENABLE_AGENT");
160157
boolean enableAgentDefault = Boolean.getBoolean("LazySetOptIn");

‎agent/instrumentation/azure-functions-worker-stub/build.gradle.kts‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
plugins {
2+
id("ai.java-conventions")
3+
}
4+
5+
// this module is needed since the azure functions worker artifact is not available in maven central

‎agent/instrumentation/azure-functions-worker-stub/src/main/java/com/microsoft/azure/functions/rpc/messages/InvocationRequest.java‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
// Copyright (c) Microsoft Corporation. All rights reserved.
2+
// Licensed under the MIT License.
3+
4+
package com.microsoft.azure.functions.rpc.messages;
5+
6+
public class InvocationRequest {
7+
8+
public RpcTraceContext getTraceContext() {
9+
throw new UnsupportedOperationException();
10+
}
11+
12+
public String getInvocationId() {
13+
throw new UnsupportedOperationException();
14+
}
15+
}

‎agent/instrumentation/azure-functions-worker-stub/src/main/java/com/microsoft/azure/functions/rpc/messages/RpcTraceContext.java‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
// Copyright (c) Microsoft Corporation. All rights reserved.
2+
// Licensed under the MIT License.
3+
4+
package com.microsoft.azure.functions.rpc.messages;
5+
6+
import java.util.Map;
7+
8+
public class RpcTraceContext {
9+
10+
public Map<String, String> getAttributesMap() {
11+
throw new UnsupportedOperationException();
12+
}
13+
14+
public String getTraceParent() {
15+
throw new UnsupportedOperationException();
16+
}
17+
18+
public String getTraceState() {
19+
throw new UnsupportedOperationException();
20+
}
21+
}

‎agent/instrumentation/azure-functions-worker-stub/src/main/java/com/microsoft/azure/functions/worker/handler/InvocationRequestHandler.java‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
// Copyright (c) Microsoft Corporation. All rights reserved.
2+
// Licensed under the MIT License.
3+
4+
package com.microsoft.azure.functions.worker.handler;
5+
6+
import com.microsoft.azure.functions.rpc.messages.InvocationRequest;
7+
8+
public class InvocationRequestHandler {
9+
10+
public void execute(InvocationRequest request) {
11+
verifyCurrentContext();
12+
}
13+
14+
// this doesn't exist in the real worker artifact
15+
// it only exists for testing the instrumentation
16+
protected void verifyCurrentContext() {}
17+
}

‎agent/instrumentation/azure-functions/build.gradle.kts‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,10 @@ muzzle {
1212
val otelInstrumentationAlphaVersion: String by project
1313

1414
dependencies {
15+
compileOnly(project(":agent:instrumentation:azure-functions-worker-stub"))
16+
17+
testImplementation(project(":agent:instrumentation:azure-functions-worker-stub"))
18+
1519
// TODO remove when start using io.opentelemetry.instrumentation.javaagent-instrumentation plugin
1620
add("codegen", "io.opentelemetry.javaagent:opentelemetry-javaagent-tooling:$otelInstrumentationAlphaVersion")
1721
add("muzzleBootstrap", "io.opentelemetry.instrumentation:opentelemetry-instrumentation-annotations-support:$otelInstrumentationAlphaVersion")

‎agent/instrumentation/azure-functions/src/main/java/io/opentelemetry/javaagent/instrumentation/azurefunctions/InvocationInstrumentation.java‎

Lines changed: 16 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,8 @@
1111
import com.microsoft.applicationinsights.agent.bootstrap.AzureFunctions;
1212
import com.microsoft.applicationinsights.agent.bootstrap.AzureFunctionsCustomDimensions;
1313
import com.microsoft.applicationinsights.agent.bootstrap.BytecodeUtil;
14+
import com.microsoft.azure.functions.rpc.messages.InvocationRequest;
15+
import com.microsoft.azure.functions.rpc.messages.RpcTraceContext;
1416
import io.opentelemetry.api.GlobalOpenTelemetry;
1517
import io.opentelemetry.api.trace.Span;
1618
import io.opentelemetry.api.trace.SpanContext;
@@ -43,23 +45,17 @@ public void transform(TypeTransformer transformer) {
4345
InvocationInstrumentation.class.getName() + "$ExecuteAdvice");
4446
}
4547

46-
@SuppressWarnings({
47-
"unused",
48-
"PrivateConstructorForUtilityClass",
49-
"MustBeClosedChecker",
50-
"unchecked"
51-
})
48+
@SuppressWarnings({"unused", "PrivateConstructorForUtilityClass", "MustBeClosedChecker"})
5249
public static class ExecuteAdvice {
5350
@Nullable
5451
@Advice.OnMethodEnter(suppress = Throwable.class)
55-
public static Scope methodEnter(@Advice.Argument(0) Object request)
56-
throws ReflectiveOperationException {
52+
public static Scope methodEnter(@Advice.Argument(0) InvocationRequest request) {
5753

5854
if (!AzureFunctions.hasConnectionString()) {
5955
return null;
6056
}
6157

62-
Object traceContext = InvocationRequestExtractAdapter.getTraceContextMethod.invoke(request);
58+
RpcTraceContext traceContext = request.getTraceContext();
6359
Context extractedContext =
6460
GlobalOpenTelemetry.getPropagators()
6561
.getTextMapPropagator()
@@ -78,10 +74,17 @@ public static Scope methodEnter(@Advice.Argument(0) Object request)
7874
traceFlags,
7975
spanContext.getTraceState());
8076

81-
return Context.current()
82-
.with(Span.wrap(spanContext))
83-
.with(generateCustomDimensions(request, traceContext))
84-
.makeCurrent();
77+
Map<String, String> attributesMap = traceContext.getAttributesMap();
78+
AzureFunctionsCustomDimensions customDimensions =
79+
new AzureFunctionsCustomDimensions(
80+
request.getInvocationId(),
81+
attributesMap.get("ProcessId"),
82+
attributesMap.get("LogLevel"),
83+
attributesMap.get("Category"),
84+
attributesMap.get("HostInstanceId"),
85+
attributesMap.get("#AzFuncLiveLogsSessionId"));
86+
87+
return Context.current().with(Span.wrap(spanContext)).with(customDimensions).makeCurrent();
8588
}
8689

8790
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
@@ -90,21 +93,5 @@ public static void methodExit(@Advice.Enter @Nullable Scope scope) {
9093
scope.close();
9194
}
9295
}
93-
94-
private static AzureFunctionsCustomDimensions generateCustomDimensions(
95-
Object request, Object traceContext) throws ReflectiveOperationException {
96-
String invocationId =
97-
(String) InvocationRequestExtractAdapter.getInvocationId.invoke(request);
98-
Map<String, String> attributesMap =
99-
(Map<String, String>)
100-
InvocationRequestExtractAdapter.getAttributesMap.invoke(traceContext);
101-
return new AzureFunctionsCustomDimensions(
102-
invocationId,
103-
attributesMap.get("ProcessId"),
104-
attributesMap.get("LogLevel"),
105-
attributesMap.get("Category"),
106-
attributesMap.get("HostInstanceId"),
107-
attributesMap.get("#AzFuncLiveLogsSessionId"));
108-
}
10996
}
11097
}

‎agent/instrumentation/azure-functions/src/main/java/io/opentelemetry/javaagent/instrumentation/azurefunctions/InvocationRequestExtractAdapter.java‎

Lines changed: 12 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -3,75 +3,31 @@
33

44
package io.opentelemetry.javaagent.instrumentation.azurefunctions;
55

6+
import com.microsoft.azure.functions.rpc.messages.RpcTraceContext;
67
import io.opentelemetry.context.propagation.TextMapGetter;
7-
import java.lang.reflect.Method;
8-
import java.util.logging.Level;
9-
import java.util.logging.Logger;
108
import javax.annotation.Nullable;
119

12-
// using reflection because these classes are not published to maven central that we can easily
13-
// compile against
14-
public class InvocationRequestExtractAdapter implements TextMapGetter<Object> {
15-
16-
private static final Logger logger =
17-
Logger.getLogger(InvocationRequestExtractAdapter.class.getName());
10+
public class InvocationRequestExtractAdapter implements TextMapGetter<RpcTraceContext> {
1811

1912
public static final InvocationRequestExtractAdapter GETTER =
2013
new InvocationRequestExtractAdapter();
2114

22-
public static final Method getTraceContextMethod;
23-
public static final Method getInvocationId;
24-
public static final Method getAttributesMap;
25-
private static final Method getTraceParentMethod;
26-
private static final Method getTraceStateMethod;
27-
28-
static {
29-
Method getTraceContextMethodLocal = null;
30-
Method getInvocationIdLocal = null;
31-
Method getAttributesMapLocal = null;
32-
Method getTraceParentMethodLocal = null;
33-
Method getTraceStateMethodLocal = null;
34-
try {
35-
Class<?> invocationRequestClass =
36-
Class.forName("com.microsoft.azure.functions.rpc.messages.InvocationRequest");
37-
Class<?> rpcTraceContextClass =
38-
Class.forName("com.microsoft.azure.functions.rpc.messages.RpcTraceContext");
39-
getTraceContextMethodLocal = invocationRequestClass.getMethod("getTraceContext");
40-
getInvocationIdLocal = invocationRequestClass.getMethod("getInvocationId");
41-
getAttributesMapLocal = rpcTraceContextClass.getMethod("getAttributesMap");
42-
getTraceParentMethodLocal = rpcTraceContextClass.getMethod("getTraceParent");
43-
getTraceStateMethodLocal = rpcTraceContextClass.getMethod("getTraceState");
44-
} catch (ReflectiveOperationException e) {
45-
logger.log(Level.SEVERE, e.getMessage(), e);
46-
}
47-
getTraceContextMethod = getTraceContextMethodLocal;
48-
getInvocationId = getInvocationIdLocal;
49-
getAttributesMap = getAttributesMapLocal;
50-
getTraceParentMethod = getTraceParentMethodLocal;
51-
getTraceStateMethod = getTraceStateMethodLocal;
52-
}
53-
5415
@Override
55-
public Iterable<String> keys(Object carrier) {
16+
public Iterable<String> keys(RpcTraceContext carrier) {
5617
return null;
5718
}
5819

5920
@Override
6021
@Nullable
61-
public String get(Object carrier, String key) {
62-
try {
63-
// only supports W3C propagator
64-
switch (key) {
65-
case "traceparent":
66-
return (String) getTraceParentMethod.invoke(carrier);
67-
case "tracestate":
68-
return (String) getTraceStateMethod.invoke(carrier);
69-
default:
70-
return null;
71-
}
72-
} catch (ReflectiveOperationException e) {
73-
logger.log(Level.FINE, e.getMessage(), e);
74-
return null;
22+
public String get(RpcTraceContext carrier, String key) {
23+
// only supports W3C propagator
24+
switch (key) {
25+
case "traceparent":
26+
return carrier.getTraceParent();
27+
case "tracestate":
28+
return carrier.getTraceState();
29+
default:
30+
return null;
7531
}
7632
}
7733
}

‎agent/instrumentation/azure-functions/src/test/java/AzureFunctionsTest.java‎

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
// Copyright (c) Microsoft Corporation. All rights reserved.
2+
// Licensed under the MIT License.
3+
4+
import static java.util.Collections.emptyMap;
5+
import static org.assertj.core.api.Assertions.assertThat;
6+
import static org.mockito.ArgumentMatchers.anyString;
7+
import static org.mockito.Mockito.mock;
8+
import static org.mockito.Mockito.when;
9+
10+
import com.microsoft.applicationinsights.agent.bootstrap.AzureFunctions;
11+
import com.microsoft.applicationinsights.agent.bootstrap.BytecodeUtil;
12+
import com.microsoft.applicationinsights.agent.bootstrap.BytecodeUtil.BytecodeUtilDelegate;
13+
import com.microsoft.azure.functions.rpc.messages.InvocationRequest;
14+
import com.microsoft.azure.functions.rpc.messages.RpcTraceContext;
15+
import com.microsoft.azure.functions.worker.handler.InvocationRequestHandler;
16+
import io.opentelemetry.api.trace.Span;
17+
import io.opentelemetry.api.trace.SpanContext;
18+
import io.opentelemetry.context.Context;
19+
import io.opentelemetry.instrumentation.testing.junit.AgentInstrumentationExtension;
20+
import io.opentelemetry.instrumentation.testing.junit.InstrumentationExtension;
21+
import java.util.Map;
22+
import java.util.concurrent.atomic.AtomicReference;
23+
import org.junit.jupiter.api.Test;
24+
import org.junit.jupiter.api.extension.RegisterExtension;
25+
26+
class AzureFunctionsTest {
27+
28+
@RegisterExtension
29+
static final InstrumentationExtension testing = AgentInstrumentationExtension.create();
30+
31+
static {
32+
// this is needed since currently tests are run against otel javaagent, and not ai javaagent
33+
AzureFunctions.setup(() -> true, () -> {});
34+
BytecodeUtilDelegate delegate = mock(BytecodeUtilDelegate.class);
35+
when(delegate.shouldSample(anyString())).thenReturn(true);
36+
BytecodeUtil.setDelegate(delegate);
37+
}
38+
39+
@Test
40+
void setRequestProperty() {
41+
// given
42+
String traceParent = "00-11111111111111111111111111111111-1111111111111111-00";
43+
String traceState = null;
44+
Map<String, String> attributesMap = emptyMap();
45+
RpcTraceContext traceContext = new MockRpcTraceContext(traceParent, traceState, attributesMap);
46+
47+
String invocationId = null;
48+
InvocationRequest request = new MockInvocationRequest(traceContext, invocationId);
49+
50+
AtomicReference<Context> contextRef = new AtomicReference<>();
51+
InvocationRequestHandler handler =
52+
new InvocationRequestHandler() {
53+
@Override
54+
protected void verifyCurrentContext() {
55+
contextRef.set(Context.current());
56+
}
57+
};
58+
59+
// when
60+
handler.execute(request);
61+
62+
// then
63+
Context context = contextRef.get();
64+
SpanContext spanContext = Span.fromContext(context).getSpanContext();
65+
assertThat(spanContext.getTraceId()).isEqualTo("11111111111111111111111111111111");
66+
assertThat(spanContext.getSpanId()).isEqualTo("1111111111111111");
67+
assertThat(spanContext.getTraceFlags().isSampled()).isTrue();
68+
}
69+
}

‎agent/instrumentation/azure-functions/src/test/java/MockInvocationRequest.java‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
// Copyright (c) Microsoft Corporation. All rights reserved.
2+
// Licensed under the MIT License.
3+
4+
import com.microsoft.azure.functions.rpc.messages.InvocationRequest;
5+
import com.microsoft.azure.functions.rpc.messages.RpcTraceContext;
6+
7+
class MockInvocationRequest extends InvocationRequest {
8+
9+
private final RpcTraceContext traceContext;
10+
private final String invocationId;
11+
12+
MockInvocationRequest(RpcTraceContext traceContext, String invocationId) {
13+
this.traceContext = traceContext;
14+
this.invocationId = invocationId;
15+
}
16+
17+
@Override
18+
public RpcTraceContext getTraceContext() {
19+
return traceContext;
20+
}
21+
22+
@Override
23+
public String getInvocationId() {
24+
return invocationId;
25+
}
26+
}

0 commit comments

Comments
 (0)