Fix pekko route naming

Author: samwright

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

@@ -9,6 +9,8 @@
 import static io.opentelemetry.javaagent.instrumentation.pekkohttp.v1_0.server.PekkoHttpServerSingletons.instrumenter;
 
 import io.opentelemetry.context.Context;
+import io.opentelemetry.instrumentation.api.semconv.http.HttpServerRoute;
+import io.opentelemetry.instrumentation.api.semconv.http.HttpServerRouteSource;
 import io.opentelemetry.javaagent.bootstrap.http.HttpServerResponseCustomizerHolder;
 import io.opentelemetry.javaagent.instrumentation.pekkohttp.v1_0.server.route.PekkoRouteHolder;
 import java.util.ArrayDeque;
@@ -143,6 +145,14 @@ public void onPush() {
                 if (!headers.isEmpty()) {
                   response = (HttpResponse) response.addHeaders(headers);
                 }
+                PekkoRouteHolder routeHolder = tracingRequest.context.get(PekkoRouteHolder.KEY);
+                if (routeHolder != null) {
+                  routeHolder.pushIfNotCompletelyMatched("**");
+                  HttpServerRoute.update(
+                      tracingRequest.context,
+                      HttpServerRouteSource.CONTROLLER,
+                      routeHolder.route());
+                }
 
                 instrumenter().end(tracingRequest.context, tracingRequest.request, response, null);
               }

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

@@ -9,7 +9,9 @@
 import static net.bytebuddy.matcher.ElementMatchers.named;
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 
+import io.opentelemetry.context.Context;
 import io.opentelemetry.instrumentation.api.util.VirtualField;
+import io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
 import net.bytebuddy.asm.Advice;
@@ -18,7 +20,6 @@
 import org.apache.pekko.http.scaladsl.model.Uri;
 import org.apache.pekko.http.scaladsl.server.PathMatcher;
 import org.apache.pekko.http.scaladsl.server.PathMatchers;
-import org.apache.pekko.http.scaladsl.server.PathMatchers$;
 
 public class PathMatcherStaticInstrumentation implements TypeInstrumentation {
   @Override
@@ -43,11 +44,13 @@ public static void onExit(
         @Advice.Argument(0) Uri.Path path,
         @Advice.Return PathMatcher.Matching<?> result) {
       // result is either matched or unmatched, we only care about the matches
+      Context context = Java8BytecodeBridge.currentContext();
+      PekkoRouteHolder routeHolder = context.get(PekkoRouteHolder.KEY);
+      if (routeHolder == null) {
+        return;
+      }
       if (result.getClass() == PathMatcher.Matched.class) {
-        if (PathMatchers$.PathEnd$.class == pathMatcher.getClass()) {
-          PekkoRouteHolder.endMatched();
-          return;
-        }
+        PathMatcher.Matched<?> match = (PathMatcher.Matched<?>) result;
         // if present use the matched path that was remembered in PathMatcherInstrumentation,
         // otherwise just use a *
         String prefix = VirtualField.find(PathMatcher.class, String.class).get(pathMatcher);
@@ -58,9 +61,9 @@ public static void onExit(
             prefix = "*";
           }
         }
-        if (prefix != null) {
-          PekkoRouteHolder.push(prefix);
-        }
+        routeHolder.push(path, match.pathRest(), prefix);
+      } else {
+        routeHolder.didNotMatch();
       }
     }
   }

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

@@ -34,7 +34,6 @@ public List<TypeInstrumentation> typeInstrumentations() {
     return asList(
         new PathMatcherInstrumentation(),
         new PathMatcherStaticInstrumentation(),
-        new RouteConcatenationInstrumentation(),
-        new PathConcatenationInstrumentation());
+        new RouteConcatenationInstrumentation());
   }
 }

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

@@ -10,18 +10,17 @@
 import io.opentelemetry.context.Context;
 import io.opentelemetry.context.ContextKey;
 import io.opentelemetry.context.ImplicitContextKeyed;
-import io.opentelemetry.instrumentation.api.semconv.http.HttpServerRoute;
-import io.opentelemetry.instrumentation.api.semconv.http.HttpServerRouteSource;
-import java.util.ArrayDeque;
 import java.util.Deque;
+import java.util.LinkedList;
+import org.apache.pekko.http.scaladsl.model.Uri;
 
 public class PekkoRouteHolder implements ImplicitContextKeyed {
-  private static final ContextKey<PekkoRouteHolder> KEY = named("opentelemetry-pekko-route");
+  public static final ContextKey<PekkoRouteHolder> KEY = named("opentelemetry-pekko-route");
 
-  private String route = "";
-  private boolean newSegment = true;
-  private boolean endMatched;
-  private final Deque<String> stack = new ArrayDeque<>();
+  private StringBuilder route = new StringBuilder();
+  private Uri.Path lastUnmatchedPath = null;
+  private boolean lastWasMatched = false;
+  private final Deque<State> savedStates = new LinkedList<>();
 
   public static Context init(Context context) {
     if (context.get(KEY) != null) {
@@ -30,51 +29,47 @@ public static Context init(Context context) {
     return context.with(new PekkoRouteHolder());
   }
 
-  public static void push(String path) {
-    PekkoRouteHolder holder = Context.current().get(KEY);
-    if (holder != null && holder.newSegment && !holder.endMatched) {
-      holder.route += path;
-      holder.newSegment = false;
+  public void push(Uri.Path beforeMatch, Uri.Path afterMatch, String pathToPush) {
+    // Only accept the suggested 'pathToPush' if:
+    // - either this is the first match, or
+    // - the unmatched part of the path from the previous match is what the current match
+    //   acted upon. This avoids pushes from PathMatchers that compose other PathMatchers,
+    //   instead only accepting pushes from leaf-nodes in the PathMatcher hierarchy that actually
+    //   act on the path.
+    // AND:
+    // - some part of the path has now been matched by this matcher
+    if ((lastUnmatchedPath == null || lastUnmatchedPath.equals(beforeMatch))
+        && !afterMatch.equals(beforeMatch)) {
+      route.append(pathToPush);
+      lastUnmatchedPath = afterMatch;
     }
+    lastWasMatched = true;
   }
 
-  public static void startSegment() {
-    PekkoRouteHolder holder = Context.current().get(KEY);
-    if (holder != null) {
-      holder.newSegment = true;
-    }
+  public void didNotMatch() {
+    lastWasMatched = false;
   }
 
-  public static void endMatched() {
-    Context context = Context.current();
-    PekkoRouteHolder holder = context.get(KEY);
-    if (holder != null) {
-      holder.endMatched = true;
-      HttpServerRoute.update(context, HttpServerRouteSource.CONTROLLER, holder.route);
+  public void pushIfNotCompletelyMatched(String pathToPush) {
+    if (lastUnmatchedPath != null && !lastUnmatchedPath.isEmpty()) {
+      route.append(pathToPush);
     }
   }
 
-  public static void save() {
-    PekkoRouteHolder holder = Context.current().get(KEY);
-    if (holder != null) {
-      holder.stack.push(holder.route);
-      holder.newSegment = true;
-    }
+  public String route() {
+    return lastWasMatched ? route.toString() : null;
   }
 
-  public static void reset() {
-    PekkoRouteHolder holder = Context.current().get(KEY);
-    if (holder != null) {
-      holder.route = holder.stack.peek();
-      holder.newSegment = true;
-    }
+  public void save() {
+    savedStates.add(new State(lastUnmatchedPath, route));
+    route = new StringBuilder(route);
   }
 
-  public static void restore() {
-    PekkoRouteHolder holder = Context.current().get(KEY);
-    if (holder != null) {
-      holder.route = holder.stack.pop();
-      holder.newSegment = true;
+  public void restore() {
+    State popped = savedStates.pollLast();
+    if (popped != null) {
+      lastUnmatchedPath = popped.lastUnmatchedPath;
+      route = popped.route;
     }
   }
 
@@ -84,4 +79,14 @@ public Context storeInContext(Context context) {
   }
 
   private PekkoRouteHolder() {}
+
+  private static class State {
+    private final Uri.Path lastUnmatchedPath;
+    private final StringBuilder route;
+
+    private State(Uri.Path lastUnmatchedPath, StringBuilder route) {
+      this.lastUnmatchedPath = lastUnmatchedPath;
+      this.route = route;
+    }
+  }
 }

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

@@ -0,0 +1,42 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.pekkohttp.v1_0.server.route;
+
+import io.opentelemetry.context.Context;
+import io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge;
+import org.apache.pekko.http.scaladsl.server.RequestContext;
+import org.apache.pekko.http.scaladsl.server.RouteResult;
+import scala.Function1;
+import scala.concurrent.Future;
+
+public class PekkoRouteWrapper implements Function1<RequestContext, Future<RouteResult>> {
+  private final Function1<RequestContext, Future<RouteResult>> route;
+
+  public PekkoRouteWrapper(Function1<RequestContext, Future<RouteResult>> route) {
+    this.route = route;
+  }
+
+  @Override
+  public Future<RouteResult> apply(RequestContext ctx) {
+    Context context = Java8BytecodeBridge.currentContext();
+    PekkoRouteHolder routeHolder = context.get(PekkoRouteHolder.KEY);
+    if (routeHolder == null) {
+      return route.apply(ctx);
+    } else {
+      routeHolder.save();
+      return route
+          .apply(ctx)
+          .map(
+              result -> {
+                if (result.getClass() == RouteResult.Rejected.class) {
+                  routeHolder.restore();
+                }
+                return result;
+              },
+              ctx.executionContext());
+    }
+  }
+}

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

@@ -10,10 +10,12 @@
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
 import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 import org.apache.pekko.http.scaladsl.server.RequestContext;
 import org.apache.pekko.http.scaladsl.server.RouteResult;
+import scala.Function1;
 import scala.concurrent.Future;
 
 public class RouteConcatenationInstrumentation implements TypeInstrumentation {
@@ -25,41 +27,17 @@ public ElementMatcher<TypeDescription> typeMatcher() {
   @Override
   public void transform(TypeTransformer transformer) {
     transformer.applyAdviceToMethod(
-        named("$anonfun$$tilde$1"), this.getClass().getName() + "$ApplyAdvice");
-    transformer.applyAdviceToMethod(
-        named("$anonfun$$tilde$2"), this.getClass().getName() + "$Apply2Advice");
+        MethodDescription::isConstructor, this.getClass().getName() + "$ApplyAdvice");
+    transformer.applyAdviceToMethod(named("$tilde"), this.getClass().getName() + "$ApplyAdvice");
   }
 
-  @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
-      // in the route
-      PekkoRouteHolder.save();
-    }
-
-    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
-    public static void onExit(
-        @Advice.Argument(value = 2) RequestContext ctx,
-        @Advice.Return(readOnly = false) Future<RouteResult> future,
-        @Advice.Thrown Throwable throwable) {
-      if (throwable != null) {
-        PekkoRouteHolder.restore();
-      } else {
-        future = future.andThen(new RestoreOnExit(), ctx.executionContext());
-      }
-    }
-  }
-
-  @SuppressWarnings("unused")
-  public static class Apply2Advice {
-
-    @Advice.OnMethodEnter(suppress = Throwable.class)
-    public static void onEnter() {
-      PekkoRouteHolder.reset();
+    public static void onEnter(
+        @Advice.Argument(value = 0, readOnly = false)
+            Function1<RequestContext, Future<RouteResult>> route) {
+      route = new PekkoRouteWrapper(route);
     }
   }
 }