|
14 | 14 | import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer; |
15 | 15 | import java.util.concurrent.CompletableFuture; |
16 | 16 | import net.bytebuddy.asm.Advice; |
| 17 | +import net.bytebuddy.asm.Advice.AssignReturned; |
| 18 | +import net.bytebuddy.asm.Advice.AssignReturned.ToArguments.ToArgument; |
17 | 19 | import net.bytebuddy.description.type.TypeDescription; |
18 | 20 | import net.bytebuddy.matcher.ElementMatcher; |
19 | 21 | import org.reactivestreams.Subscriber; |
@@ -52,20 +54,21 @@ public void transform(TypeTransformer transformer) { |
52 | 54 | @SuppressWarnings("unused") |
53 | 55 | public static class WrapSubscriberAdvice { |
54 | 56 |
|
| 57 | + @AssignReturned.ToArguments(@ToArgument(1)) |
55 | 58 | @Advice.OnMethodEnter(suppress = Throwable.class) |
56 | | - public static void wrapSubscriber( |
57 | | - @Advice.Argument(value = 1, readOnly = false) Subscriber<?> subscriber) { |
58 | | - subscriber = SubscriberWrapper.wrap(subscriber); |
| 59 | + public static Subscriber<?> wrapSubscriber(@Advice.Argument(1) Subscriber<?> subscriber) { |
| 60 | + return SubscriberWrapper.wrap(subscriber); |
59 | 61 | } |
60 | 62 | } |
61 | 63 |
|
62 | 64 | @SuppressWarnings("unused") |
63 | 65 | public static class WrapCompletableFutureAdvice { |
64 | 66 |
|
| 67 | + @AssignReturned.ToArguments(@ToArgument(4)) |
65 | 68 | @Advice.OnMethodEnter(suppress = Throwable.class) |
66 | | - public static void wrapCompletableFuture( |
67 | | - @Advice.Argument(value = 4, readOnly = false) CompletableFuture<?> future) { |
68 | | - future = CompletableFutureWrapper.wrap(future); |
| 69 | + public static CompletableFuture<?> wrapCompletableFuture( |
| 70 | + @Advice.Argument(4) CompletableFuture<?> future) { |
| 71 | + return CompletableFutureWrapper.wrap(future); |
69 | 72 | } |
70 | 73 | } |
71 | 74 | } |
0 commit comments