|
5 | 5 |
|
6 | 6 | package io.opentelemetry.javaagent.instrumentation.powerjob.v4_0; |
7 | 7 |
|
8 | | -import static io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge.currentContext; |
9 | 8 | import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.implementsInterface; |
10 | 9 | import static io.opentelemetry.javaagent.instrumentation.powerjob.v4_0.PowerJobSingletons.instrumenter; |
11 | 10 | import static net.bytebuddy.matcher.ElementMatchers.isPublic; |
|
17 | 16 | import io.opentelemetry.context.Scope; |
18 | 17 | import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation; |
19 | 18 | import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer; |
| 19 | +import javax.annotation.Nullable; |
20 | 20 | import net.bytebuddy.asm.Advice; |
21 | 21 | import net.bytebuddy.description.type.TypeDescription; |
22 | 22 | import net.bytebuddy.matcher.ElementMatcher; |
@@ -45,43 +45,58 @@ public void transform(TypeTransformer transformer) { |
45 | 45 |
|
46 | 46 | public static class ProcessAdvice { |
47 | 47 |
|
48 | | - @SuppressWarnings("unused") |
49 | | - @Advice.OnMethodEnter(suppress = Throwable.class) |
50 | | - public static void onSchedule( |
51 | | - @Advice.This BasicProcessor handler, |
52 | | - @Advice.Argument(0) TaskContext taskContext, |
53 | | - @Advice.Local("otelRequest") PowerJobProcessRequest request, |
54 | | - @Advice.Local("otelContext") Context context, |
55 | | - @Advice.Local("otelScope") Scope scope) { |
56 | | - Context parentContext = currentContext(); |
57 | | - request = |
58 | | - PowerJobProcessRequest.createRequest( |
59 | | - taskContext.getJobId(), |
60 | | - handler, |
61 | | - "process", |
62 | | - taskContext.getJobParams(), |
63 | | - taskContext.getInstanceParams()); |
| 48 | + public static class AdviceScope { |
| 49 | + private final PowerJobProcessRequest request; |
| 50 | + private final Context context; |
| 51 | + private final Scope scope; |
| 52 | + |
| 53 | + private AdviceScope(PowerJobProcessRequest request, Context context, Scope scope) { |
| 54 | + this.request = request; |
| 55 | + this.context = context; |
| 56 | + this.scope = scope; |
| 57 | + } |
| 58 | + |
| 59 | + @Nullable |
| 60 | + public static AdviceScope start(BasicProcessor handler, TaskContext taskContext) { |
| 61 | + Context parentContext = Context.current(); |
| 62 | + PowerJobProcessRequest request = |
| 63 | + PowerJobProcessRequest.createRequest( |
| 64 | + taskContext.getJobId(), |
| 65 | + handler, |
| 66 | + "process", |
| 67 | + taskContext.getJobParams(), |
| 68 | + taskContext.getInstanceParams()); |
| 69 | + |
| 70 | + if (!instrumenter().shouldStart(parentContext, request)) { |
| 71 | + return null; |
| 72 | + } |
64 | 73 |
|
65 | | - if (!instrumenter().shouldStart(parentContext, request)) { |
66 | | - return; |
| 74 | + Context context = instrumenter().start(parentContext, request); |
| 75 | + return new AdviceScope(request, context, context.makeCurrent()); |
67 | 76 | } |
68 | | - context = instrumenter().start(parentContext, request); |
69 | | - scope = context.makeCurrent(); |
| 77 | + |
| 78 | + public void end(ProcessResult result, Throwable throwable) { |
| 79 | + scope.close(); |
| 80 | + instrumenter().end(context, request, result, throwable); |
| 81 | + } |
| 82 | + } |
| 83 | + |
| 84 | + @SuppressWarnings("unused") |
| 85 | + @Advice.OnMethodEnter(suppress = Throwable.class) |
| 86 | + public static AdviceScope onSchedule( |
| 87 | + @Advice.This BasicProcessor handler, @Advice.Argument(0) TaskContext taskContext) { |
| 88 | + return AdviceScope.start(handler, taskContext); |
70 | 89 | } |
71 | 90 |
|
72 | 91 | @SuppressWarnings("unused") |
73 | 92 | @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) |
74 | 93 | public static void stopSpan( |
75 | 94 | @Advice.Return ProcessResult result, |
76 | | - @Advice.Thrown Throwable throwable, |
77 | | - @Advice.Local("otelRequest") PowerJobProcessRequest request, |
78 | | - @Advice.Local("otelContext") Context context, |
79 | | - @Advice.Local("otelScope") Scope scope) { |
80 | | - if (scope == null) { |
81 | | - return; |
| 95 | + @Advice.Thrown @Nullable Throwable throwable, |
| 96 | + @Advice.Enter @Nullable AdviceScope adviceScope) { |
| 97 | + if (adviceScope != null) { |
| 98 | + adviceScope.end(result, throwable); |
82 | 99 | } |
83 | | - scope.close(); |
84 | | - instrumenter().end(context, request, result, throwable); |
85 | 100 | } |
86 | 101 | } |
87 | 102 | } |
0 commit comments