Skip to content

Commit 66755dc

Browse files
SylvainJugeSylvainJuge
committed
aws-lambda-events
1 parent 2bd450b commit 66755dc

File tree

3 files changed

+126
-71
lines changed

3 files changed

+126
-71
lines changed

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

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,12 +10,14 @@
1010
import com.google.auto.service.AutoService;
1111
import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
1212
import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
13+
import io.opentelemetry.javaagent.extension.instrumentation.internal.ExperimentalInstrumentationModule;
1314
import java.util.Arrays;
1415
import java.util.List;
1516
import net.bytebuddy.matcher.ElementMatcher;
1617

1718
@AutoService(InstrumentationModule.class)
18-
public class AwsLambdaInstrumentationModule extends InstrumentationModule {
19+
public class AwsLambdaInstrumentationModule extends InstrumentationModule
20+
implements ExperimentalInstrumentationModule {
1921
public AwsLambdaInstrumentationModule() {
2022
super("aws-lambda-events", "aws-lambda-events-2.2", "aws-lambda");
2123
}
@@ -36,4 +38,9 @@ public List<TypeInstrumentation> typeInstrumentations() {
3638
new AwsLambdaRequestHandlerInstrumentation(),
3739
new AwsLambdaRequestStreamHandlerInstrumentation());
3840
}
41+
42+
@Override
43+
public boolean isIndyReady() {
44+
return true;
45+
}
3946
}

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

Lines changed: 74 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasClassesNamed;
99
import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.implementsInterface;
1010
import static io.opentelemetry.javaagent.instrumentation.awslambdaevents.v2_2.AwsLambdaSingletons.flushTimeout;
11+
import static io.opentelemetry.javaagent.instrumentation.awslambdaevents.v2_2.AwsLambdaSingletons.functionInstrumenter;
1112
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
1213
import static net.bytebuddy.matcher.ElementMatchers.isPublic;
1314
import static net.bytebuddy.matcher.ElementMatchers.named;
@@ -25,6 +26,7 @@
2526
import java.util.Collections;
2627
import java.util.Map;
2728
import java.util.concurrent.TimeUnit;
29+
import javax.annotation.Nullable;
2830
import net.bytebuddy.asm.Advice;
2931
import net.bytebuddy.description.type.TypeDescription;
3032
import net.bytebuddy.implementation.bytecode.assign.Assigner.Typing;
@@ -55,63 +57,90 @@ public void transform(TypeTransformer transformer) {
5557
@SuppressWarnings("unused")
5658
public static class HandleRequestAdvice {
5759

58-
@Advice.OnMethodEnter(suppress = Throwable.class)
59-
public static void onEnter(
60-
@Advice.Argument(value = 0, typing = Typing.DYNAMIC) Object arg,
61-
@Advice.Argument(1) Context context,
62-
@Advice.Local("otelInput") AwsLambdaRequest input,
63-
@Advice.Local("otelFunctionContext") io.opentelemetry.context.Context functionContext,
64-
@Advice.Local("otelFunctionScope") Scope functionScope,
65-
@Advice.Local("otelMessageContext") io.opentelemetry.context.Context messageContext,
66-
@Advice.Local("otelMessageScope") Scope messageScope) {
67-
Map<String, String> headers = Collections.emptyMap();
68-
if (arg instanceof APIGatewayProxyRequestEvent) {
69-
headers = MapUtils.lowercaseMap(((APIGatewayProxyRequestEvent) arg).getHeaders());
60+
public static class AdviceScope {
61+
private final AwsLambdaRequest lambdaRequest;
62+
private final Scope functionScope;
63+
private final io.opentelemetry.context.Context functionContext;
64+
private final Scope messageScope;
65+
private final io.opentelemetry.context.Context messageContext;
66+
67+
private AdviceScope(
68+
AwsLambdaRequest lambdaRequest,
69+
io.opentelemetry.context.Context functionContext,
70+
Scope functionScope,
71+
io.opentelemetry.context.Context messageContext,
72+
Scope messageScope) {
73+
this.lambdaRequest = lambdaRequest;
74+
this.functionContext = functionContext;
75+
this.functionScope = functionScope;
76+
this.messageContext = messageContext;
77+
this.messageScope = messageScope;
7078
}
71-
input = AwsLambdaRequest.create(context, arg, headers);
72-
io.opentelemetry.context.Context parentContext =
73-
AwsLambdaSingletons.functionInstrumenter().extract(input);
7479

75-
if (!AwsLambdaSingletons.functionInstrumenter().shouldStart(parentContext, input)) {
76-
return;
77-
}
80+
@Nullable
81+
public static AdviceScope start(Object arg, Context context) {
7882

79-
functionContext = AwsLambdaSingletons.functionInstrumenter().start(parentContext, input);
80-
functionScope = functionContext.makeCurrent();
83+
Map<String, String> headers = Collections.emptyMap();
84+
if (arg instanceof APIGatewayProxyRequestEvent) {
85+
headers = MapUtils.lowercaseMap(((APIGatewayProxyRequestEvent) arg).getHeaders());
86+
}
87+
AwsLambdaRequest lambdaRequest = AwsLambdaRequest.create(context, arg, headers);
88+
io.opentelemetry.context.Context parentContext =
89+
functionInstrumenter().extract(lambdaRequest);
8190

82-
if (arg instanceof SQSEvent) {
83-
if (AwsLambdaSingletons.messageInstrumenter()
84-
.shouldStart(functionContext, (SQSEvent) arg)) {
85-
messageContext =
86-
AwsLambdaSingletons.messageInstrumenter().start(functionContext, (SQSEvent) arg);
87-
messageScope = messageContext.makeCurrent();
91+
if (!functionInstrumenter().shouldStart(parentContext, lambdaRequest)) {
92+
return null;
8893
}
94+
95+
io.opentelemetry.context.Context functionContext =
96+
functionInstrumenter().start(parentContext, lambdaRequest);
97+
Scope functionScope = functionContext.makeCurrent();
98+
99+
io.opentelemetry.context.Context messageContext = null;
100+
Scope messageScope = null;
101+
if (arg instanceof SQSEvent) {
102+
if (AwsLambdaSingletons.messageInstrumenter()
103+
.shouldStart(functionContext, (SQSEvent) arg)) {
104+
messageContext =
105+
AwsLambdaSingletons.messageInstrumenter().start(functionContext, (SQSEvent) arg);
106+
messageScope = messageContext.makeCurrent();
107+
}
108+
}
109+
return new AdviceScope(
110+
lambdaRequest, functionContext, functionScope, messageContext, messageScope);
89111
}
112+
113+
public void end(Object arg, @Nullable Object result, @Nullable Throwable throwable) {
114+
if (messageScope != null) {
115+
messageScope.close();
116+
AwsLambdaSingletons.messageInstrumenter()
117+
.end(messageContext, (SQSEvent) arg, null, throwable);
118+
}
119+
if (functionScope != null) {
120+
functionScope.close();
121+
functionInstrumenter().end(functionContext, lambdaRequest, result, throwable);
122+
}
123+
OpenTelemetrySdkAccess.forceFlush(flushTimeout().toNanos(), TimeUnit.NANOSECONDS);
124+
}
125+
}
126+
127+
@Nullable
128+
@Advice.OnMethodEnter(suppress = Throwable.class)
129+
public static AdviceScope onEnter(
130+
@Advice.Argument(value = 0, typing = Typing.DYNAMIC) Object arg,
131+
@Advice.Argument(1) Context context) {
132+
return AdviceScope.start(arg, context);
90133
}
91134

92135
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
93136
public static void stopSpan(
94137
@Advice.Argument(value = 0, typing = Typing.DYNAMIC) Object arg,
95-
@Advice.Return Object result,
96-
@Advice.Thrown Throwable throwable,
97-
@Advice.Local("otelInput") AwsLambdaRequest input,
98-
@Advice.Local("otelFunctionContext") io.opentelemetry.context.Context functionContext,
99-
@Advice.Local("otelFunctionScope") Scope functionScope,
100-
@Advice.Local("otelMessageContext") io.opentelemetry.context.Context messageContext,
101-
@Advice.Local("otelMessageScope") Scope messageScope) {
102-
103-
if (messageScope != null) {
104-
messageScope.close();
105-
AwsLambdaSingletons.messageInstrumenter()
106-
.end(messageContext, (SQSEvent) arg, null, throwable);
107-
}
108-
109-
if (functionScope != null) {
110-
functionScope.close();
111-
AwsLambdaSingletons.functionInstrumenter().end(functionContext, input, result, throwable);
138+
@Advice.Return @Nullable Object result,
139+
@Advice.Thrown @Nullable Throwable throwable,
140+
@Advice.Enter @Nullable AdviceScope adviceScope) {
141+
if (adviceScope != null) {
142+
adviceScope.end(arg, result, throwable);
112143
}
113-
114-
OpenTelemetrySdkAccess.forceFlush(flushTimeout().toNanos(), TimeUnit.NANOSECONDS);
115144
}
116145
}
117146
}

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

Lines changed: 44 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasClassesNamed;
99
import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.implementsInterface;
1010
import static io.opentelemetry.javaagent.instrumentation.awslambdaevents.v2_2.AwsLambdaSingletons.flushTimeout;
11+
import static io.opentelemetry.javaagent.instrumentation.awslambdaevents.v2_2.AwsLambdaSingletons.functionInstrumenter;
1112
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
1213
import static net.bytebuddy.matcher.ElementMatchers.isPublic;
1314
import static net.bytebuddy.matcher.ElementMatchers.named;
@@ -21,8 +22,8 @@
2122
import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
2223
import java.io.InputStream;
2324
import java.util.Collections;
24-
import java.util.Map;
2525
import java.util.concurrent.TimeUnit;
26+
import javax.annotation.Nullable;
2627
import net.bytebuddy.asm.Advice;
2728
import net.bytebuddy.description.type.TypeDescription;
2829
import net.bytebuddy.matcher.ElementMatcher;
@@ -52,38 +53,56 @@ public void transform(TypeTransformer transformer) {
5253
@SuppressWarnings("unused")
5354
public static class HandleRequestAdvice {
5455

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-
AwsLambdaSingletons.functionInstrumenter().extract(otelInput);
56+
public static class AdviceScope {
57+
private final AwsLambdaRequest lambdaRequest;
58+
private final io.opentelemetry.context.Context context;
59+
private final Scope scope;
6660

67-
if (!AwsLambdaSingletons.functionInstrumenter().shouldStart(parentContext, otelInput)) {
68-
return;
61+
private AdviceScope(
62+
AwsLambdaRequest lambdaRequest,
63+
io.opentelemetry.context.Context otelContext,
64+
Scope scope) {
65+
this.lambdaRequest = lambdaRequest;
66+
this.context = otelContext;
67+
this.scope = scope;
6968
}
7069

71-
otelContext = AwsLambdaSingletons.functionInstrumenter().start(parentContext, otelInput);
72-
otelScope = otelContext.makeCurrent();
70+
@Nullable
71+
public static AdviceScope start(Object arg, Context context) {
72+
AwsLambdaRequest lambdaRequest =
73+
AwsLambdaRequest.create(context, arg, Collections.emptyMap());
74+
io.opentelemetry.context.Context parentContext =
75+
functionInstrumenter().extract(lambdaRequest);
76+
if (!functionInstrumenter().shouldStart(parentContext, lambdaRequest)) {
77+
return null;
78+
}
79+
80+
io.opentelemetry.context.Context otelContext =
81+
functionInstrumenter().start(parentContext, lambdaRequest);
82+
return new AdviceScope(lambdaRequest, otelContext, otelContext.makeCurrent());
83+
}
84+
85+
public void end(@Nullable Throwable throwable) {
86+
scope.close();
87+
functionInstrumenter().end(context, lambdaRequest, null, throwable);
88+
OpenTelemetrySdkAccess.forceFlush(flushTimeout().toNanos(), TimeUnit.NANOSECONDS);
89+
}
90+
}
91+
92+
@Nullable
93+
@Advice.OnMethodEnter(suppress = Throwable.class)
94+
public static AdviceScope onEnter(
95+
@Advice.Argument(0) InputStream input, @Advice.Argument(2) Context context) {
96+
return AdviceScope.start(input, context);
7397
}
7498

7599
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
76100
public static void stopSpan(
77-
@Advice.Thrown Throwable throwable,
78-
@Advice.Local("otelInput") AwsLambdaRequest input,
79-
@Advice.Local("otelContext") io.opentelemetry.context.Context functionContext,
80-
@Advice.Local("otelScope") Scope functionScope) {
81-
if (functionScope != null) {
82-
functionScope.close();
83-
AwsLambdaSingletons.functionInstrumenter().end(functionContext, input, null, throwable);
101+
@Advice.Thrown @Nullable Throwable throwable,
102+
@Advice.Enter @Nullable AdviceScope adviceScope) {
103+
if (adviceScope != null) {
104+
adviceScope.end(throwable);
84105
}
85-
86-
OpenTelemetrySdkAccess.forceFlush(flushTimeout().toNanos(), TimeUnit.NANOSECONDS);
87106
}
88107
}
89108
}

0 commit comments

Comments
 (0)