|
24 | 24 | import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer; |
25 | 25 | import java.util.Collections; |
26 | 26 | import java.util.concurrent.TimeUnit; |
| 27 | +import javax.annotation.Nullable; |
27 | 28 | import net.bytebuddy.asm.Advice; |
28 | 29 | import net.bytebuddy.description.type.TypeDescription; |
29 | 30 | import net.bytebuddy.implementation.bytecode.assign.Assigner.Typing; |
@@ -60,38 +61,57 @@ public void transform(TypeTransformer transformer) { |
60 | 61 | @SuppressWarnings("unused") |
61 | 62 | public static class HandleRequestAdvice { |
62 | 63 |
|
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 | + } |
72 | 92 |
|
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); |
75 | 97 | } |
| 98 | + } |
76 | 99 |
|
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); |
79 | 105 | } |
80 | 106 |
|
81 | 107 | @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) |
82 | 108 | public static void stopSpan( |
83 | 109 | @Advice.Argument(value = 0, typing = Typing.DYNAMIC) Object arg, |
84 | 110 | @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); |
92 | 114 | } |
93 | | - |
94 | | - OpenTelemetrySdkAccess.forceFlush(flushTimeout().toNanos(), TimeUnit.NANOSECONDS); |
95 | 115 | } |
96 | 116 | } |
97 | 117 | } |
0 commit comments