fix(akka): improve handling of http.route (#15504)

Co-authored-by: Abhishek Rao <[email protected]>
Co-authored-by: Lauri Tulmin <[email protected]>

Author: 3 people

instrumentation/akka/akka-http-10.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/akkahttp/client/HttpExtClientInstrumentation.java

@@ -5,7 +5,6 @@
 
 package io.opentelemetry.javaagent.instrumentation.akkahttp.client;
 
-import static io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge.currentContext;
 import static io.opentelemetry.javaagent.instrumentation.akkahttp.client.AkkaHttpClientSingletons.instrumenter;
 import static io.opentelemetry.javaagent.instrumentation.akkahttp.client.AkkaHttpClientSingletons.setter;
 import static net.bytebuddy.matcher.ElementMatchers.named;
@@ -56,7 +55,7 @@ private AdviceScope(Context context, Scope scope) {
       }
 
       public static AdviceScope start(HttpRequest request) {
-        Context parentContext = currentContext();
+        Context parentContext = Context.current();
         if (!instrumenter().shouldStart(parentContext, request)) {
           return null;
         }
@@ -83,7 +82,7 @@ public Future<HttpResponse> end(
           if (throwable == null) {
             responseFuture.onComplete(
                 new OnCompleteHandler(context, request), actorSystem.dispatcher());
-            return FutureWrapper.wrap(responseFuture, actorSystem.dispatcher(), currentContext());
+            return FutureWrapper.wrap(responseFuture, actorSystem.dispatcher(), Context.current());
           } else {
             instrumenter().end(context, request, null, throwable);
           }

instrumentation/akka/akka-http-10.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/akkahttp/server/AkkaFlowWrapper.java

@@ -5,7 +5,6 @@
 
 package io.opentelemetry.javaagent.instrumentation.akkahttp.server;
 
-import static io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge.currentContext;
 import static io.opentelemetry.javaagent.instrumentation.akkahttp.server.AkkaHttpServerSingletons.errorResponse;
 import static io.opentelemetry.javaagent.instrumentation.akkahttp.server.AkkaHttpServerSingletons.instrumenter;
 
@@ -23,6 +22,8 @@
 import akka.stream.stage.GraphStageLogic;
 import akka.stream.stage.OutHandler;
 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.akkahttp.server.route.AkkaRouteHolder;
 import java.util.ArrayDeque;
@@ -116,7 +117,7 @@ public void onPush() {
               HttpRequest request = grab(requestIn);
 
               TracingRequest tracingRequest = TracingRequest.EMPTY;
-              Context parentContext = currentContext();
+              Context parentContext = Context.current();
               if (instrumenter().shouldStart(parentContext, request)) {
                 Context context = instrumenter().start(parentContext, request);
                 context = AkkaRouteHolder.init(context);
@@ -160,6 +161,15 @@ public void onPush() {
                   response = (HttpResponse) response.addHeaders(headers);
                 }
 
+                AkkaRouteHolder routeHolder = AkkaRouteHolder.get(tracingRequest.context);
+                if (routeHolder != null) {
+                  routeHolder.pushIfNotCompletelyMatched("*");
+                  HttpServerRoute.update(
+                      tracingRequest.context,
+                      HttpServerRouteSource.CONTROLLER,
+                      routeHolder.route());
+                }
+
                 instrumenter().end(tracingRequest.context, tracingRequest.request, response, null);
               }
               push(responseOut, response);

instrumentation/akka/akka-http-10.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/akkahttp/server/route/AkkaHttpServerRouteInstrumentationModule.java

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

instrumentation/akka/akka-http-10.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/akkahttp/server/route/AkkaRouteHolder.java

@@ -7,21 +7,20 @@
 
 import static io.opentelemetry.context.ContextKey.named;
 
+import akka.http.scaladsl.model.Uri;
 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;
 
 public class AkkaRouteHolder implements ImplicitContextKeyed {
   private static final ContextKey<AkkaRouteHolder> KEY = named("opentelemetry-akka-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 ArrayDeque<>();
 
   public static Context init(Context context) {
     if (context.get(KEY) != null) {
@@ -30,52 +29,51 @@ public static Context init(Context context) {
     return context.with(new AkkaRouteHolder());
   }
 
-  public static void push(String path) {
-    AkkaRouteHolder holder = Context.current().get(KEY);
-    if (holder != null && holder.newSegment && !holder.endMatched) {
-      holder.route += path;
-      holder.newSegment = false;
-    }
+  public static AkkaRouteHolder get(Context context) {
+    return context.get(KEY);
   }
 
-  public static void startSegment() {
-    AkkaRouteHolder holder = Context.current().get(KEY);
-    if (holder != null) {
-      holder.newSegment = true;
+  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 endMatched() {
-    Context context = Context.current();
-    AkkaRouteHolder holder = context.get(KEY);
-    if (holder != null) {
-      holder.endMatched = true;
-      HttpServerRoute.update(context, HttpServerRouteSource.CONTROLLER, holder.route);
-    }
+  public void didNotMatch() {
+    lastWasMatched = false;
   }
 
-  public static void save() {
-    AkkaRouteHolder holder = Context.current().get(KEY);
-    if (holder != null) {
-      holder.stack.push(holder.route);
-      holder.newSegment = true;
+  public void pushIfNotCompletelyMatched(String pathToPush) {
+    if (lastUnmatchedPath != null && !lastUnmatchedPath.isEmpty()) {
+      route.append(pathToPush);
     }
   }
 
-  public static void restore() {
-    AkkaRouteHolder holder = Context.current().get(KEY);
-    if (holder != null) {
-      holder.route = holder.stack.pop();
-      holder.newSegment = true;
-    }
+  public String route() {
+    return lastWasMatched ? route.toString() : null;
+  }
+
+  public void save() {
+    savedStates.add(new State(lastUnmatchedPath, route));
+    route = new StringBuilder(route);
   }
 
-  // reset the state to when save was called
-  public static void reset() {
-    AkkaRouteHolder holder = Context.current().get(KEY);
-    if (holder != null) {
-      holder.route = holder.stack.peek();
-      holder.newSegment = true;
+  public void restore() {
+    State popped = savedStates.pollLast();
+    if (popped != null) {
+      lastUnmatchedPath = popped.lastUnmatchedPath;
+      route = popped.route;
     }
   }
 
@@ -85,4 +83,14 @@ public Context storeInContext(Context context) {
   }
 
   private AkkaRouteHolder() {}
+
+  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/akka/akka-http-10.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/akkahttp/server/route/AkkaRouteUtil.java

@@ -0,0 +1,17 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.akkahttp.server.route;
+
+import akka.http.scaladsl.server.PathMatcher;
+import io.opentelemetry.instrumentation.api.util.VirtualField;
+
+public final class AkkaRouteUtil {
+
+  public static final VirtualField<PathMatcher<?>, String> PREFIX =
+      VirtualField.find(PathMatcher.class, String.class);
+
+  private AkkaRouteUtil() {}
+}

instrumentation/akka/akka-http-10.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/akkahttp/server/route/AkkaRouteWrapper.java

@@ -0,0 +1,45 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.akkahttp.server.route;
+
+import akka.http.scaladsl.server.RequestContext;
+import akka.http.scaladsl.server.RouteResult;
+import io.opentelemetry.context.Context;
+import scala.Function1;
+import scala.concurrent.Future;
+import scala.runtime.AbstractFunction1;
+
+public class AkkaRouteWrapper extends AbstractFunction1<RequestContext, Future<RouteResult>> {
+  private final Function1<RequestContext, Future<RouteResult>> route;
+
+  public AkkaRouteWrapper(Function1<RequestContext, Future<RouteResult>> route) {
+    this.route = route;
+  }
+
+  @Override
+  public Future<RouteResult> apply(RequestContext ctx) {
+    Context context = Context.current();
+    AkkaRouteHolder routeHolder = AkkaRouteHolder.get(context);
+    if (routeHolder == null) {
+      return route.apply(ctx);
+    } else {
+      routeHolder.save();
+      return route
+          .apply(ctx)
+          .map(
+              new AbstractFunction1<RouteResult, RouteResult>() {
+                @Override
+                public RouteResult apply(RouteResult result) {
+                  if (result.getClass() == RouteResult.Rejected.class) {
+                    routeHolder.restore();
+                  }
+                  return result;
+                }
+              },
+              ctx.executionContext());
+    }
+  }
+}

instrumentation/akka/akka-http-10.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/akkahttp/server/route/PathConcatenationInstrumentation.java

@@ -5,6 +5,7 @@
 
 package io.opentelemetry.javaagent.instrumentation.akkahttp.server.route;
 
+import static io.opentelemetry.javaagent.instrumentation.akkahttp.server.route.AkkaRouteUtil.PREFIX;
 import static net.bytebuddy.matcher.ElementMatchers.named;
 import static net.bytebuddy.matcher.ElementMatchers.returns;
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
@@ -40,7 +41,7 @@ public static void onEnter(
         @Advice.Argument(0) Uri.Path prefix, @Advice.Return PathMatcher<?> result) {
       // store the path being matched inside a VirtualField on the given matcher, so it can be used
       // for constructing the route
-      PathMatcherUtil.setMatched(result, prefix.toString());
+      PREFIX.set(result, prefix.toString());
     }
   }
 }

instrumentation/akka/akka-http-10.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/akkahttp/server/route/PathMatcherInstrumentation.java

@@ -6,12 +6,15 @@
 package io.opentelemetry.javaagent.instrumentation.akkahttp.server.route;
 
 import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.extendsClass;
+import static io.opentelemetry.javaagent.instrumentation.akkahttp.server.route.AkkaRouteUtil.PREFIX;
 import static net.bytebuddy.matcher.ElementMatchers.named;
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 
 import akka.http.scaladsl.model.Uri;
 import akka.http.scaladsl.server.PathMatcher;
 import akka.http.scaladsl.server.PathMatchers;
+import io.opentelemetry.context.Context;
+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;
@@ -40,24 +43,26 @@ 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();
+      AkkaRouteHolder routeHolder = AkkaRouteHolder.get(context);
+      if (routeHolder == null) {
+        return;
+      }
       if (result.getClass() == PathMatcher.Matched.class) {
-        if (PathMatchers.PathEnd$.class == pathMatcher.getClass()) {
-          AkkaRouteHolder.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 = PathMatcherUtil.getMatched(pathMatcher);
+        String prefix = PREFIX.get(pathMatcher);
         if (prefix == null) {
           if (PathMatchers.Slash$.class == pathMatcher.getClass()) {
             prefix = "/";
           } else {
             prefix = "*";
           }
         }
-        if (prefix != null) {
-          AkkaRouteHolder.push(prefix);
-        }
+        routeHolder.push(path, match.pathRest(), prefix);
+      } else {
+        routeHolder.didNotMatch();
       }
     }
   }