Ensure tilde$1 onExit is run in correct order

Author: masonedmison

instrumentation/pekko/pekko-http-1.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/pekkohttp/v1_0/server/route/RouteConcatenationInstrumentation.java

@@ -12,6 +12,12 @@
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
+import org.apache.pekko.http.javadsl.server.RouteResult;
+import org.apache.pekko.http.scaladsl.server.RequestContext;
+import scala.PartialFunction;
+import scala.Unit;
+import scala.concurrent.Future;
+import scala.util.Try;
 
 public class RouteConcatenationInstrumentation implements TypeInstrumentation {
   @Override
@@ -27,20 +33,39 @@ public void transform(TypeTransformer transformer) {
         named("$anonfun$$tilde$2"), this.getClass().getName() + "$Apply2Advice");
   }
 
+  public static class OnExitFinalizer implements PartialFunction<Try<RouteResult>, Unit> {
+    @Override
+    public boolean isDefinedAt(Try<RouteResult> x) {
+      return true;
+    }
+
+    @Override
+    public Unit apply(Try<RouteResult> v1) {
+      System.err.println("===tilde$1 onExit===");
+      PekkoRouteHolder.restore();
+      return null;
+    }
+  }
+
   @SuppressWarnings("unused")
   public static class ApplyAdvice {
 
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static void onEnter() {
-      // when routing dsl uses concat(path(...) {...}, path(...) {...}) we'll restore the currently
-      // matched route after each matcher so that match attempts that failed wouldn't get recorded
+      // when routing dsl uses concat(path(...) {...}, path(...) {...}) we'll restore
+      // the currently
+      // matched route after each matcher so that match attempts that failed wouldn't
+      // get recorded
       // in the route
+      System.err.println("===tilde$1 onEnter===");
       PekkoRouteHolder.save();
     }
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
-    public static void onExit() {
-      PekkoRouteHolder.restore();
+    public static void onExit(
+        @Advice.Argument(value = 2) RequestContext ctx,
+        @Advice.Return(readOnly = false) Future<RouteResult> fut) {
+      fut = fut.andThen(new OnExitFinalizer(), ctx.executionContext());
     }
   }
 
@@ -49,6 +74,7 @@ public static class Apply2Advice {
 
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static void onEnter() {
+      System.err.println("===tilde$2 onEnter===");
       PekkoRouteHolder.reset();
     }
   }