Hopefully fix flaky ratpack test
@@ -10,8 +10,10 @@ import static net.bytebuddy.matcher.ElementMatchers.nameStartsWith;
import static net.bytebuddy.matcher.ElementMatchers.named;
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+import io.opentelemetry.context.Scope;
import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import io.opentelemetry.javaagent.instrumentation.api.Java8BytecodeBridge;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.matcher.ElementMatcher;
@@ -30,14 +32,35 @@ public class DefaultExecStarterInstrumentation implements TypeInstrumentation {
transformer.applyAdviceToMethod(
named("onComplete")
.or(named("onError"))
+ .or(named("onStart"))
.and(takesArgument(0, named("ratpack.func.Action"))),
DefaultExecStarterInstrumentation.class.getName() + "$WrapActionAdvice");
+ transformer.applyAdviceToMethod(
+ named("start").and(takesArgument(0, named("ratpack.func.Action"))),
+ DefaultExecStarterInstrumentation.class.getName() + "$StartAdvice");
}
+ @SuppressWarnings("unused")
public static class WrapActionAdvice {
@Advice.OnMethodEnter(suppress = Throwable.class)
public static void wrapAction(@Advice.Argument(value = 0, readOnly = false) Action<?> action) {
action = ActionWrapper.wrapIfNeeded(action);
}
}
+
+ @SuppressWarnings("unused")
+ public static class StartAdvice {
+ @Advice.OnMethodEnter(suppress = Throwable.class)
+ public static Scope enter(@Advice.Argument(value = 0, readOnly = false) Action<?> action) {
+ action = ActionWrapper.wrapIfNeeded(action);
+ return Java8BytecodeBridge.rootContext().makeCurrent();
+ }
+
+ @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+ public static void exit(@Advice.Enter Scope scope) {
+ if (scope != null) {
+ scope.close();
+ }
+ }
+ }
}
