Skip to content

Commit 3c14532

Browse files
lukeina2zlukeina2z
authored
feat: Add instrumentation for Lambda Java interface HandleStreamRequest (#13466)
1 parent dd05ebb commit 3c14532

File tree

6 files changed

+424
-4
lines changed

6 files changed

+424
-4
lines changed

‎instrumentation/aws-lambda/aws-lambda-core-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/awslambdacore/v1_0/AwsLambdaInstrumentationModule.java

+5-2
Original file line numberDiff line numberDiff line change
@@ -6,17 +6,18 @@
66
package io.opentelemetry.javaagent.instrumentation.awslambdacore.v1_0;
77

88
import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasClassesNamed;
9-
import static java.util.Collections.singletonList;
109
import static net.bytebuddy.matcher.ElementMatchers.not;
1110

1211
import com.google.auto.service.AutoService;
1312
import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
1413
import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
14+
import java.util.Arrays;
1515
import java.util.List;
1616
import net.bytebuddy.matcher.ElementMatcher;
1717

1818
@AutoService(InstrumentationModule.class)
1919
public class AwsLambdaInstrumentationModule extends InstrumentationModule {
20+
2021
public AwsLambdaInstrumentationModule() {
2122
super("aws-lambda-core", "aws-lambda-core-1.0", "aws-lambda");
2223
}
@@ -34,6 +35,8 @@ public boolean isHelperClass(String className) {
3435

3536
@Override
3637
public List<TypeInstrumentation> typeInstrumentations() {
37-
return singletonList(new AwsLambdaRequestHandlerInstrumentation());
38+
return Arrays.asList(
39+
new AwsLambdaRequestHandlerInstrumentation(),
40+
new AwsLambdaRequestStreamHandlerInstrumentation());
3841
}
3942
}

‎instrumentation/aws-lambda/aws-lambda-core-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/awslambdacore/v1_0/AwsLambdaRequestStreamHandlerInstrumentation.java

+97
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,97 @@
1+
/*
2+
* Copyright The OpenTelemetry Authors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package io.opentelemetry.javaagent.instrumentation.awslambdacore.v1_0;
7+
8+
import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasClassesNamed;
9+
import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.implementsInterface;
10+
import static io.opentelemetry.javaagent.instrumentation.awslambdacore.v1_0.AwsLambdaInstrumentationHelper.flushTimeout;
11+
import static io.opentelemetry.javaagent.instrumentation.awslambdacore.v1_0.AwsLambdaInstrumentationHelper.functionInstrumenter;
12+
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
13+
import static net.bytebuddy.matcher.ElementMatchers.isPublic;
14+
import static net.bytebuddy.matcher.ElementMatchers.nameStartsWith;
15+
import static net.bytebuddy.matcher.ElementMatchers.named;
16+
import static net.bytebuddy.matcher.ElementMatchers.not;
17+
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
18+
19+
import com.amazonaws.services.lambda.runtime.Context;
20+
import io.opentelemetry.context.Scope;
21+
import io.opentelemetry.instrumentation.awslambdacore.v1_0.AwsLambdaRequest;
22+
import io.opentelemetry.javaagent.bootstrap.OpenTelemetrySdkAccess;
23+
import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
24+
import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
25+
import java.io.InputStream;
26+
import java.util.Collections;
27+
import java.util.concurrent.TimeUnit;
28+
import net.bytebuddy.asm.Advice;
29+
import net.bytebuddy.description.type.TypeDescription;
30+
import net.bytebuddy.matcher.ElementMatcher;
31+
32+
public class AwsLambdaRequestStreamHandlerInstrumentation implements TypeInstrumentation {
33+
34+
@Override
35+
public ElementMatcher<ClassLoader> classLoaderOptimization() {
36+
return hasClassesNamed("com.amazonaws.services.lambda.runtime.RequestStreamHandler");
37+
}
38+
39+
@Override
40+
public ElementMatcher<TypeDescription> typeMatcher() {
41+
return implementsInterface(named("com.amazonaws.services.lambda.runtime.RequestStreamHandler"))
42+
.and(not(nameStartsWith("com.amazonaws.services.lambda.runtime.api.client")))
43+
// In Java 8 and Java 11 runtimes,
44+
// AWS Lambda runtime is packaged under `lambdainternal` package.
45+
// But it is `com.amazonaws.services.lambda.runtime.api.client`
46+
// for new runtime likes Java 17 and Java 21.
47+
.and(not(nameStartsWith("lambdainternal")));
48+
}
49+
50+
@Override
51+
public void transform(TypeTransformer transformer) {
52+
transformer.applyAdviceToMethod(
53+
isMethod()
54+
.and(isPublic())
55+
.and(named("handleRequest"))
56+
.and(takesArgument(2, named("com.amazonaws.services.lambda.runtime.Context"))),
57+
AwsLambdaRequestStreamHandlerInstrumentation.class.getName() + "$HandleRequestAdvice");
58+
}
59+
60+
@SuppressWarnings("unused")
61+
public static class HandleRequestAdvice {
62+
63+
@Advice.OnMethodEnter(suppress = Throwable.class)
64+
public static void onEnter(
65+
@Advice.Argument(0) InputStream input,
66+
@Advice.Argument(2) Context context,
67+
@Advice.Local("otelInput") AwsLambdaRequest otelInput,
68+
@Advice.Local("otelContext") io.opentelemetry.context.Context otelContext,
69+
@Advice.Local("otelScope") Scope otelScope) {
70+
71+
otelInput = AwsLambdaRequest.create(context, input, Collections.emptyMap());
72+
io.opentelemetry.context.Context parentContext = functionInstrumenter().extract(otelInput);
73+
74+
if (!functionInstrumenter().shouldStart(parentContext, otelInput)) {
75+
return;
76+
}
77+
78+
otelContext = functionInstrumenter().start(parentContext, otelInput);
79+
otelScope = otelContext.makeCurrent();
80+
}
81+
82+
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
83+
public static void stopSpan(
84+
@Advice.Thrown Throwable throwable,
85+
@Advice.Local("otelInput") AwsLambdaRequest input,
86+
@Advice.Local("otelContext") io.opentelemetry.context.Context functionContext,
87+
@Advice.Local("otelScope") Scope functionScope) {
88+
89+
if (functionScope != null) {
90+
functionScope.close();
91+
functionInstrumenter().end(functionContext, input, null, throwable);
92+
}
93+
94+
OpenTelemetrySdkAccess.forceFlush(flushTimeout().toNanos(), TimeUnit.NANOSECONDS);
95+
}
96+
}
97+
}

‎instrumentation/aws-lambda/aws-lambda-core-1.0/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/awslambdacore/v1_0/AwsLambdaStreamHandlerTest.java

+113
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,113 @@
1+
/*
2+
* Copyright The OpenTelemetry Authors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package io.opentelemetry.javaagent.instrumentation.awslambdacore.v1_0;
7+
8+
import static io.opentelemetry.sdk.testing.assertj.OpenTelemetryAssertions.equalTo;
9+
import static org.assertj.core.api.Assertions.assertThat;
10+
import static org.assertj.core.api.Assertions.catchThrowable;
11+
import static org.mockito.Mockito.when;
12+
13+
import com.amazonaws.services.lambda.runtime.Context;
14+
import com.amazonaws.services.lambda.runtime.RequestStreamHandler;
15+
import io.opentelemetry.api.trace.SpanKind;
16+
import io.opentelemetry.instrumentation.testing.junit.AgentInstrumentationExtension;
17+
import io.opentelemetry.instrumentation.testing.junit.InstrumentationExtension;
18+
import io.opentelemetry.sdk.trace.data.StatusData;
19+
import io.opentelemetry.semconv.incubating.FaasIncubatingAttributes;
20+
import java.io.BufferedReader;
21+
import java.io.BufferedWriter;
22+
import java.io.ByteArrayInputStream;
23+
import java.io.ByteArrayOutputStream;
24+
import java.io.IOException;
25+
import java.io.InputStream;
26+
import java.io.InputStreamReader;
27+
import java.io.OutputStream;
28+
import java.io.OutputStreamWriter;
29+
import java.nio.charset.StandardCharsets;
30+
import org.junit.jupiter.api.AfterEach;
31+
import org.junit.jupiter.api.BeforeEach;
32+
import org.junit.jupiter.api.Test;
33+
import org.junit.jupiter.api.extension.ExtendWith;
34+
import org.junit.jupiter.api.extension.RegisterExtension;
35+
import org.mockito.Mock;
36+
import org.mockito.junit.jupiter.MockitoExtension;
37+
38+
@ExtendWith(MockitoExtension.class)
39+
public class AwsLambdaStreamHandlerTest {
40+
41+
@RegisterExtension
42+
public static final InstrumentationExtension testing = AgentInstrumentationExtension.create();
43+
44+
@Mock private Context context;
45+
46+
@BeforeEach
47+
void setUp() {
48+
when(context.getFunctionName()).thenReturn("my_function");
49+
when(context.getAwsRequestId()).thenReturn("1-22-333");
50+
}
51+
52+
@AfterEach
53+
void tearDown() {
54+
assertThat(testing.forceFlushCalled()).isTrue();
55+
}
56+
57+
@Test
58+
void handlerTraced() throws Exception {
59+
InputStream input = new ByteArrayInputStream("hello\n".getBytes(StandardCharsets.UTF_8));
60+
OutputStream output = new ByteArrayOutputStream();
61+
RequestStreamHandlerTestImpl handler = new RequestStreamHandlerTestImpl();
62+
handler.handleRequest(input, output, context);
63+
64+
testing.waitAndAssertTraces(
65+
trace ->
66+
trace.hasSpansSatisfyingExactly(
67+
span ->
68+
span.hasName("my_function")
69+
.hasKind(SpanKind.SERVER)
70+
.hasAttributesSatisfyingExactly(
71+
equalTo(FaasIncubatingAttributes.FAAS_INVOCATION_ID, "1-22-333"))));
72+
}
73+
74+
@Test
75+
void handlerTracedWithException() {
76+
InputStream input = new ByteArrayInputStream("bye\n".getBytes(StandardCharsets.UTF_8));
77+
OutputStream output = new ByteArrayOutputStream();
78+
RequestStreamHandlerTestImpl handler = new RequestStreamHandlerTestImpl();
79+
80+
Throwable thrown = catchThrowable(() -> handler.handleRequest(input, output, context));
81+
assertThat(thrown).isInstanceOf(IllegalArgumentException.class);
82+
83+
testing.waitAndAssertTraces(
84+
trace ->
85+
trace.hasSpansSatisfyingExactly(
86+
span ->
87+
span.hasName("my_function")
88+
.hasKind(SpanKind.SERVER)
89+
.hasStatus(StatusData.error())
90+
.hasException(thrown)
91+
.hasAttributesSatisfyingExactly(
92+
equalTo(FaasIncubatingAttributes.FAAS_INVOCATION_ID, "1-22-333"))));
93+
}
94+
95+
static final class RequestStreamHandlerTestImpl implements RequestStreamHandler {
96+
@Override
97+
public void handleRequest(InputStream input, OutputStream output, Context context)
98+
throws IOException {
99+
BufferedReader reader =
100+
new BufferedReader(new InputStreamReader(input, StandardCharsets.UTF_8));
101+
BufferedWriter writer =
102+
new BufferedWriter(new OutputStreamWriter(output, StandardCharsets.UTF_8));
103+
String line = reader.readLine();
104+
if (line.equals("hello")) {
105+
writer.write("world");
106+
writer.flush();
107+
writer.close();
108+
} else {
109+
throw new IllegalArgumentException("bad argument");
110+
}
111+
}
112+
}
113+
}

‎instrumentation/aws-lambda/aws-lambda-events-2.2/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/awslambdaevents/v2_2/AwsLambdaInstrumentationModule.java

+4-2
Original file line numberDiff line numberDiff line change
@@ -6,11 +6,11 @@
66
package io.opentelemetry.javaagent.instrumentation.awslambdaevents.v2_2;
77

88
import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasClassesNamed;
9-
import static java.util.Collections.singletonList;
109

1110
import com.google.auto.service.AutoService;
1211
import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
1312
import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
13+
import java.util.Arrays;
1414
import java.util.List;
1515
import net.bytebuddy.matcher.ElementMatcher;
1616

@@ -32,6 +32,8 @@ public boolean isHelperClass(String className) {
3232

3333
@Override
3434
public List<TypeInstrumentation> typeInstrumentations() {
35-
return singletonList(new AwsLambdaRequestHandlerInstrumentation());
35+
return Arrays.asList(
36+
new AwsLambdaRequestHandlerInstrumentation(),
37+
new AwsLambdaRequestStreamHandlerInstrumentation());
3638
}
3739
}

‎instrumentation/aws-lambda/aws-lambda-events-2.2/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/awslambdaevents/v2_2/AwsLambdaRequestStreamHandlerInstrumentation.java

+92
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,92 @@
1+
/*
2+
* Copyright The OpenTelemetry Authors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package io.opentelemetry.javaagent.instrumentation.awslambdaevents.v2_2;
7+
8+
import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasClassesNamed;
9+
import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.implementsInterface;
10+
import static io.opentelemetry.javaagent.instrumentation.awslambdaevents.v2_2.AwsLambdaInstrumentationHelper.flushTimeout;
11+
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
12+
import static net.bytebuddy.matcher.ElementMatchers.isPublic;
13+
import static net.bytebuddy.matcher.ElementMatchers.named;
14+
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
15+
16+
import com.amazonaws.services.lambda.runtime.Context;
17+
import io.opentelemetry.context.Scope;
18+
import io.opentelemetry.instrumentation.awslambdacore.v1_0.AwsLambdaRequest;
19+
import io.opentelemetry.javaagent.bootstrap.OpenTelemetrySdkAccess;
20+
import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
21+
import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
22+
import java.io.InputStream;
23+
import java.util.Collections;
24+
import java.util.Map;
25+
import java.util.concurrent.TimeUnit;
26+
import net.bytebuddy.asm.Advice;
27+
import net.bytebuddy.description.type.TypeDescription;
28+
import net.bytebuddy.matcher.ElementMatcher;
29+
30+
public class AwsLambdaRequestStreamHandlerInstrumentation implements TypeInstrumentation {
31+
32+
@Override
33+
public ElementMatcher<ClassLoader> classLoaderOptimization() {
34+
return hasClassesNamed("com.amazonaws.services.lambda.runtime.RequestStreamHandler");
35+
}
36+
37+
@Override
38+
public ElementMatcher<TypeDescription> typeMatcher() {
39+
return implementsInterface(named("com.amazonaws.services.lambda.runtime.RequestStreamHandler"));
40+
}
41+
42+
@Override
43+
public void transform(TypeTransformer transformer) {
44+
transformer.applyAdviceToMethod(
45+
isMethod()
46+
.and(isPublic())
47+
.and(named("handleRequest"))
48+
.and(takesArgument(2, named("com.amazonaws.services.lambda.runtime.Context"))),
49+
AwsLambdaRequestStreamHandlerInstrumentation.class.getName() + "$HandleRequestAdvice");
50+
}
51+
52+
@SuppressWarnings("unused")
53+
public static class HandleRequestAdvice {
54+
55+
@Advice.OnMethodEnter(suppress = Throwable.class)
56+
public static void onEnter(
57+
@Advice.Argument(0) InputStream input,
58+
@Advice.Argument(2) Context context,
59+
@Advice.Local("otelInput") AwsLambdaRequest otelInput,
60+
@Advice.Local("otelContext") io.opentelemetry.context.Context otelContext,
61+
@Advice.Local("otelScope") Scope otelScope) {
62+
Map<String, String> headers = Collections.emptyMap();
63+
otelInput = AwsLambdaRequest.create(context, input, headers);
64+
io.opentelemetry.context.Context parentContext =
65+
AwsLambdaInstrumentationHelper.functionInstrumenter().extract(otelInput);
66+
67+
if (!AwsLambdaInstrumentationHelper.functionInstrumenter()
68+
.shouldStart(parentContext, otelInput)) {
69+
return;
70+
}
71+
72+
otelContext =
73+
AwsLambdaInstrumentationHelper.functionInstrumenter().start(parentContext, otelInput);
74+
otelScope = otelContext.makeCurrent();
75+
}
76+
77+
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
78+
public static void stopSpan(
79+
@Advice.Thrown Throwable throwable,
80+
@Advice.Local("otelInput") AwsLambdaRequest input,
81+
@Advice.Local("otelContext") io.opentelemetry.context.Context functionContext,
82+
@Advice.Local("otelScope") Scope functionScope) {
83+
if (functionScope != null) {
84+
functionScope.close();
85+
AwsLambdaInstrumentationHelper.functionInstrumenter()
86+
.end(functionContext, input, null, throwable);
87+
}
88+
89+
OpenTelemetrySdkAccess.forceFlush(flushTimeout().toNanos(), TimeUnit.NANOSECONDS);
90+
}
91+
}
92+
}

0 commit comments

Comments
 (0)