Skip to content

Commit 7f9a4a0

Browse files
SylvainJugeSylvainJuge
authored
make aws-lambda indy-ready (#14829)
1 parent 8de9ca4 commit 7f9a4a0

File tree

6 files changed

+222
-119
lines changed

6 files changed

+222
-119
lines changed

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

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

1819
@AutoService(InstrumentationModule.class)
19-
public class AwsLambdaInstrumentationModule extends InstrumentationModule {
20+
public class AwsLambdaInstrumentationModule extends InstrumentationModule
21+
implements ExperimentalInstrumentationModule {
2022

2123
public AwsLambdaInstrumentationModule() {
2224
super("aws-lambda-core", "aws-lambda-core-1.0", "aws-lambda");
@@ -39,4 +41,9 @@ public List<TypeInstrumentation> typeInstrumentations() {
3941
new AwsLambdaRequestHandlerInstrumentation(),
4042
new AwsLambdaRequestStreamHandlerInstrumentation());
4143
}
44+
45+
@Override
46+
public boolean isIndyReady() {
47+
return true;
48+
}
4249
}

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

Lines changed: 42 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
2525
import java.util.Collections;
2626
import java.util.concurrent.TimeUnit;
27+
import javax.annotation.Nullable;
2728
import net.bytebuddy.asm.Advice;
2829
import net.bytebuddy.description.type.TypeDescription;
2930
import net.bytebuddy.implementation.bytecode.assign.Assigner.Typing;
@@ -60,38 +61,57 @@ public void transform(TypeTransformer transformer) {
6061
@SuppressWarnings("unused")
6162
public static class HandleRequestAdvice {
6263

63-
@Advice.OnMethodEnter(suppress = Throwable.class)
64-
public static void onEnter(
65-
@Advice.Argument(value = 0, typing = Typing.DYNAMIC) Object arg,
66-
@Advice.Argument(1) Context context,
67-
@Advice.Local("otelInput") AwsLambdaRequest input,
68-
@Advice.Local("otelContext") io.opentelemetry.context.Context otelContext,
69-
@Advice.Local("otelScope") Scope otelScope) {
70-
input = AwsLambdaRequest.create(context, arg, Collections.emptyMap());
71-
io.opentelemetry.context.Context parentContext = functionInstrumenter().extract(input);
64+
public static class AdviceScope {
65+
private final AwsLambdaRequest lambdaRequest;
66+
private final io.opentelemetry.context.Context context;
67+
private final Scope scope;
68+
69+
private AdviceScope(
70+
AwsLambdaRequest lambdaRequest,
71+
io.opentelemetry.context.Context otelContext,
72+
Scope scope) {
73+
this.lambdaRequest = lambdaRequest;
74+
this.context = otelContext;
75+
this.scope = scope;
76+
}
77+
78+
@Nullable
79+
public static AdviceScope start(Object arg, Context context) {
80+
AwsLambdaRequest lambdaRequest =
81+
AwsLambdaRequest.create(context, arg, Collections.emptyMap());
82+
io.opentelemetry.context.Context parentContext =
83+
functionInstrumenter().extract(lambdaRequest);
84+
if (!functionInstrumenter().shouldStart(parentContext, lambdaRequest)) {
85+
return null;
86+
}
87+
88+
io.opentelemetry.context.Context otelContext =
89+
functionInstrumenter().start(parentContext, lambdaRequest);
90+
return new AdviceScope(lambdaRequest, otelContext, otelContext.makeCurrent());
91+
}
7292

73-
if (!functionInstrumenter().shouldStart(parentContext, input)) {
74-
return;
93+
public void end(@Nullable Throwable throwable) {
94+
scope.close();
95+
functionInstrumenter().end(context, lambdaRequest, null, throwable);
96+
OpenTelemetrySdkAccess.forceFlush(flushTimeout().toNanos(), TimeUnit.NANOSECONDS);
7597
}
98+
}
7699

77-
otelContext = functionInstrumenter().start(parentContext, input);
78-
otelScope = otelContext.makeCurrent();
100+
@Advice.OnMethodEnter(suppress = Throwable.class)
101+
public static AdviceScope onEnter(
102+
@Advice.Argument(value = 0, typing = Typing.DYNAMIC) Object arg,
103+
@Advice.Argument(1) Context context) {
104+
return AdviceScope.start(arg, context);
79105
}
80106

81107
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
82108
public static void stopSpan(
83109
@Advice.Argument(value = 0, typing = Typing.DYNAMIC) Object arg,
84110
@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);
111+
@Advice.Enter @Nullable AdviceScope adviceScope) {
112+
if (adviceScope != null) {
113+
adviceScope.end(throwable);
92114
}
93-
94-
OpenTelemetrySdkAccess.forceFlush(flushTimeout().toNanos(), TimeUnit.NANOSECONDS);
95115
}
96116
}
97117
}

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

Lines changed: 46 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
import java.io.InputStream;
2626
import java.util.Collections;
2727
import java.util.concurrent.TimeUnit;
28+
import javax.annotation.Nullable;
2829
import net.bytebuddy.asm.Advice;
2930
import net.bytebuddy.description.type.TypeDescription;
3031
import net.bytebuddy.matcher.ElementMatcher;
@@ -60,38 +61,58 @@ public void transform(TypeTransformer transformer) {
6061
@SuppressWarnings("unused")
6162
public static class HandleRequestAdvice {
6263

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;
64+
public static class AdviceScope {
65+
private final AwsLambdaRequest lambdaRequest;
66+
private final io.opentelemetry.context.Context context;
67+
private final Scope scope;
68+
69+
private AdviceScope(
70+
AwsLambdaRequest lambdaRequest,
71+
io.opentelemetry.context.Context otelContext,
72+
Scope scope) {
73+
this.lambdaRequest = lambdaRequest;
74+
this.context = otelContext;
75+
this.scope = scope;
76+
}
77+
78+
@Nullable
79+
public static AdviceScope start(InputStream input, Context context) {
80+
AwsLambdaRequest lambdaRequest =
81+
AwsLambdaRequest.create(context, input, Collections.emptyMap());
82+
io.opentelemetry.context.Context parentContext =
83+
functionInstrumenter().extract(lambdaRequest);
84+
85+
if (!functionInstrumenter().shouldStart(parentContext, lambdaRequest)) {
86+
return null;
87+
}
88+
89+
io.opentelemetry.context.Context otelContext =
90+
functionInstrumenter().start(parentContext, lambdaRequest);
91+
return new AdviceScope(lambdaRequest, otelContext, otelContext.makeCurrent());
7692
}
7793

78-
otelContext = functionInstrumenter().start(parentContext, otelInput);
79-
otelScope = otelContext.makeCurrent();
94+
public void end(@Nullable Throwable throwable) {
95+
scope.close();
96+
functionInstrumenter().end(context, lambdaRequest, null, throwable);
97+
98+
OpenTelemetrySdkAccess.forceFlush(flushTimeout().toNanos(), TimeUnit.NANOSECONDS);
99+
}
100+
}
101+
102+
@Nullable
103+
@Advice.OnMethodEnter(suppress = Throwable.class)
104+
public static AdviceScope onEnter(
105+
@Advice.Argument(0) InputStream input, @Advice.Argument(2) Context context) {
106+
return AdviceScope.start(input, context);
80107
}
81108

82109
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
83110
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);
111+
@Advice.Thrown @Nullable Throwable throwable,
112+
@Advice.Enter @Nullable AdviceScope adviceScope) {
113+
if (adviceScope != null) {
114+
adviceScope.end(throwable);
92115
}
93-
94-
OpenTelemetrySdkAccess.forceFlush(flushTimeout().toNanos(), TimeUnit.NANOSECONDS);
95116
}
96117
}
97118
}

‎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
}

0 commit comments

Comments
 (0)