Update webflux to Instrumenter API (and improvements/simplifications) (…

…#3798) * Update Webflux to Instrumenter API * Change webflux handler to match other handlers * Further simplification * Fix Mono failure tests * Use extractors! * Renames * Fix comment * Update instrumentation/spring/spring-webflux-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spring/webflux/server/AdviceUtils.java Co-authored-by: Mateusz Rzeszutek <[email protected]>
open-telemetry · Aug 13, 2021 · 84e9d18 · 84e9d18
1 parent 44b71b7
commit 84e9d18
Show file tree

Hide file tree

Showing 10 changed files with 159 additions and 147 deletions.
diff --git a/...in/java/io/opentelemetry/javaagent/instrumentation/spring/webflux/server/AdviceUtils.java b/...in/java/io/opentelemetry/javaagent/instrumentation/spring/webflux/server/AdviceUtils.java
@@ -5,18 +5,17 @@
 
 package io.opentelemetry.javaagent.instrumentation.spring.webflux.server;
 
-import io.opentelemetry.api.trace.Span;
-import io.opentelemetry.api.trace.StatusCode;
+import static io.opentelemetry.javaagent.instrumentation.spring.webflux.server.WebfluxSingletons.instrumenter;
+
 import io.opentelemetry.context.Context;
 import io.opentelemetry.instrumentation.api.tracer.ClassNames;
-import java.util.Map;
-import org.springframework.web.reactive.function.server.ServerRequest;
+import org.checkerframework.checker.nullness.qual.Nullable;
 import org.springframework.web.server.ServerWebExchange;
 import reactor.core.publisher.Mono;
 
 public class AdviceUtils {
 
-  public static final String CONTEXT_ATTRIBUTE = AdviceUtils.class.getName() + ".Context";
+  public static final String ON_SPAN_END = AdviceUtils.class.getName() + ".Context";
 
   public static String spanNameForHandler(Object handler) {
     String className = ClassNames.simpleName(handler.getClass());
@@ -28,38 +27,30 @@ public static String spanNameForHandler(Object handler) {
     return className + ".handle";
   }
 
-  public static <T> Mono<T> setPublisherSpan(Mono<T> mono, Context context) {
-    return mono.doOnError(t -> finishSpanIfPresent(context, t))
-        .doOnSuccess(x -> finishSpanIfPresent(context, null))
-        .doOnCancel(() -> finishSpanIfPresent(context, null));
+  public static void registerOnSpanEnd(
+      ServerWebExchange exchange, Context context, Object handler) {
+    exchange
+        .getAttributes()
+        .put(
+            AdviceUtils.ON_SPAN_END,
+            (AdviceUtils.OnSpanEnd) t -> instrumenter().end(context, handler, null, t));
   }
 
-  public static void finishSpanIfPresent(ServerWebExchange exchange, Throwable throwable) {
-    if (exchange != null) {
-      finishSpanIfPresentInAttributes(exchange.getAttributes(), throwable);
-    }
-  }
-
-  public static void finishSpanIfPresent(ServerRequest serverRequest, Throwable throwable) {
-    if (serverRequest != null) {
-      finishSpanIfPresentInAttributes(serverRequest.attributes(), throwable);
-    }
+  public static <T> Mono<T> end(Mono<T> mono, ServerWebExchange exchange) {
+    return mono.doOnError(throwable -> end(exchange, throwable))
+        .doOnSuccess(t -> end(exchange, null))
+        .doOnCancel(() -> end(exchange, null));
   }
 
-  static void finishSpanIfPresent(Context context, Throwable throwable) {
-    if (context != null) {
-      Span span = Span.fromContext(context);
-      if (throwable != null) {
-        span.setStatus(StatusCode.ERROR);
-        span.recordException(throwable);
-      }
-      span.end();
+  private static void end(ServerWebExchange exchange, @Nullable Throwable throwable) {
+    OnSpanEnd onSpanEnd = (OnSpanEnd) exchange.getAttributes().get(AdviceUtils.ON_SPAN_END);
+    if (onSpanEnd != null) {
+      onSpanEnd.end(throwable);
     }
   }
 
-  private static void finishSpanIfPresentInAttributes(
-      Map<String, Object> attributes, Throwable throwable) {
-    Context context = (Context) attributes.remove(CONTEXT_ATTRIBUTE);
-    finishSpanIfPresent(context, throwable);
+  @FunctionalInterface
+  interface OnSpanEnd {
+    void end(Throwable throwable);
   }
 }
diff --git a/...try/javaagent/instrumentation/spring/webflux/server/DispatcherHandlerInstrumentation.java b/...try/javaagent/instrumentation/spring/webflux/server/DispatcherHandlerInstrumentation.java
@@ -5,16 +5,12 @@
 
 package io.opentelemetry.javaagent.instrumentation.spring.webflux.server;
 
-import static io.opentelemetry.javaagent.instrumentation.spring.webflux.server.SpringWebfluxHttpServerTracer.tracer;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.isPublic;
 import static net.bytebuddy.matcher.ElementMatchers.named;
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
 
-import io.opentelemetry.api.trace.SpanKind;
-import io.opentelemetry.context.Context;
-import io.opentelemetry.context.Scope;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
 import net.bytebuddy.asm.Advice;
@@ -41,41 +37,20 @@ public void transform(TypeTransformer transformer) {
         this.getClass().getName() + "$HandleAdvice");
   }
 
-  /**
-   * This is 'top level' advice for Webflux instrumentation. This handles creating and finishing
-   * Webflux span.
-   */
   @SuppressWarnings("unused")
   public static class HandleAdvice {
 
-    @Advice.OnMethodEnter(suppress = Throwable.class)
-    public static void methodEnter(
-        @Advice.Argument(0) ServerWebExchange exchange,
-        @Advice.Local("otelScope") Scope otelScope,
-        @Advice.Local("otelContext") Context otelContext) {
-
-      otelContext = tracer().startSpan("DispatcherHandler.handle", SpanKind.INTERNAL);
-      // Unfortunately Netty EventLoop is not instrumented well enough to attribute all work to the
-      // right things so we have to store the context in request itself.
-      exchange.getAttributes().put(AdviceUtils.CONTEXT_ATTRIBUTE, otelContext);
-
-      otelScope = otelContext.makeCurrent();
-    }
-
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void methodExit(
         @Advice.Thrown Throwable throwable,
         @Advice.Argument(0) ServerWebExchange exchange,
-        @Advice.Return(readOnly = false) Mono<Void> mono,
-        @Advice.Local("otelScope") Scope otelScope,
-        @Advice.Local("otelContext") Context otelContext) {
-      if (throwable == null && mono != null) {
-        mono = AdviceUtils.setPublisherSpan(mono, otelContext);
-      } else if (throwable != null) {
-        AdviceUtils.finishSpanIfPresent(exchange, throwable);
+        @Advice.Return(readOnly = false) Mono<Void> mono) {
+      if (mono != null) {
+        // note: it seems like this code should go in @OnMethodExit of
+        // HandlerAdapterInstrumentation.HandleAdvice instead, but for some reason "GET to bad
+        // endpoint annotation API fail Mono" test fails with that placement
+        mono = AdviceUtils.end(mono, exchange);
       }
-      otelScope.close();
-      // span finished in SpanFinishingSubscriber
     }
   }
 }
diff --git a/...etry/javaagent/instrumentation/spring/webflux/server/ExperimentalAttributesExtractor.java b/...etry/javaagent/instrumentation/spring/webflux/server/ExperimentalAttributesExtractor.java
@@ -0,0 +1,36 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.spring.webflux.server;
+
+import io.opentelemetry.api.common.AttributeKey;
+import io.opentelemetry.api.common.AttributesBuilder;
+import io.opentelemetry.instrumentation.api.instrumenter.AttributesExtractor;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.springframework.web.method.HandlerMethod;
+
+public class ExperimentalAttributesExtractor extends AttributesExtractor<Object, Void> {
+
+  private static final AttributeKey<String> HANDLER_TYPE =
+      AttributeKey.stringKey("spring-webflux.handler.type");
+
+  @Override
+  protected void onStart(AttributesBuilder attributes, Object handler) {
+    attributes.put(HANDLER_TYPE, getHandlerType(handler));
+  }
+
+  @Override
+  protected void onEnd(AttributesBuilder attributes, Object handler, @Nullable Void unused) {}
+
+  private static String getHandlerType(Object handler) {
+    if (handler instanceof HandlerMethod) {
+      // Special case for requests mapped with annotations
+      HandlerMethod handlerMethod = (HandlerMethod) handler;
+      return handlerMethod.getMethod().getDeclaringClass().getName();
+    } else {
+      return handler.getClass().getName();
+    }
+  }
+}
diff --git a/...emetry/javaagent/instrumentation/spring/webflux/server/HandlerAdapterInstrumentation.java b/...emetry/javaagent/instrumentation/spring/webflux/server/HandlerAdapterInstrumentation.java
@@ -7,6 +7,7 @@
 
 import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasClassesNamed;
 import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.implementsInterface;
+import static io.opentelemetry.javaagent.instrumentation.spring.webflux.server.WebfluxSingletons.instrumenter;
 import static net.bytebuddy.matcher.ElementMatchers.isAbstract;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.isPublic;
@@ -20,14 +21,11 @@
 import io.opentelemetry.context.Scope;
 import io.opentelemetry.instrumentation.api.servlet.ServletContextPath;
 import io.opentelemetry.instrumentation.api.tracer.ServerSpan;
-import io.opentelemetry.instrumentation.api.tracer.SpanNames;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
-import io.opentelemetry.javaagent.instrumentation.spring.webflux.SpringWebfluxConfig;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
-import org.springframework.web.method.HandlerMethod;
 import org.springframework.web.reactive.HandlerMapping;
 import org.springframework.web.server.ServerWebExchange;
 import org.springframework.web.util.pattern.PathPattern;
@@ -64,55 +62,51 @@ public static class HandleAdvice {
     public static void methodEnter(
         @Advice.Argument(0) ServerWebExchange exchange,
         @Advice.Argument(1) Object handler,
+        @Advice.Local("otelContext") Context context,
         @Advice.Local("otelScope") Scope scope) {
 
-      Context context = exchange.getAttribute(AdviceUtils.CONTEXT_ATTRIBUTE);
-      if (handler != null && context != null) {
-        Span span = Span.fromContext(context);
-        String handlerType;
-        String spanName;
-
-        if (handler instanceof HandlerMethod) {
-          // Special case for requests mapped with annotations
-          HandlerMethod handlerMethod = (HandlerMethod) handler;
-          spanName = SpanNames.fromMethod(handlerMethod.getMethod());
-          handlerType = handlerMethod.getMethod().getDeclaringClass().getName();
-        } else {
-          spanName = AdviceUtils.spanNameForHandler(handler);
-          handlerType = handler.getClass().getName();
-        }
-
-        span.updateName(spanName);
-        if (SpringWebfluxConfig.captureExperimentalSpanAttributes()) {
-          span.setAttribute("spring-webflux.handler.type", handlerType);
-        }
-
-        scope = context.makeCurrent();
+      Context parentContext = Context.current();
+
+      Span serverSpan = ServerSpan.fromContextOrNull(parentContext);
+
+      PathPattern bestPattern =
+          exchange.getAttribute(HandlerMapping.BEST_MATCHING_PATTERN_ATTRIBUTE);
+      if (serverSpan != null && bestPattern != null) {
+        serverSpan.updateName(
+            ServletContextPath.prepend(Context.current(), bestPattern.toString()));
       }
 
-      if (context != null) {
-        Span serverSpan = ServerSpan.fromContextOrNull(context);
+      if (handler == null) {
+        return;
+      }
 
-        PathPattern bestPattern =
-            exchange.getAttribute(HandlerMapping.BEST_MATCHING_PATTERN_ATTRIBUTE);
-        if (serverSpan != null && bestPattern != null) {
-          serverSpan.updateName(
-              ServletContextPath.prepend(Context.current(), bestPattern.toString()));
-        }
+      if (!instrumenter().shouldStart(parentContext, handler)) {
+        return;
       }
+
+      context = instrumenter().start(parentContext, handler);
+      scope = context.makeCurrent();
     }
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void methodExit(
         @Advice.Argument(0) ServerWebExchange exchange,
+        @Advice.Argument(1) Object handler,
         @Advice.Thrown Throwable throwable,
+        @Advice.Local("otelContext") Context context,
         @Advice.Local("otelScope") Scope scope) {
-      if (throwable != null) {
-        AdviceUtils.finishSpanIfPresent(exchange, throwable);
+      if (scope == null) {
+        return;
       }
-      if (scope != null) {
-        scope.close();
-        // span finished in SpanFinishingSubscriber
+      scope.close();
+
+      if (throwable != null) {
+        instrumenter().end(context, handler, null, throwable);
+      } else {
+        AdviceUtils.registerOnSpanEnd(exchange, context, handler);
+        // span finished by wrapped Mono in DispatcherHandlerInstrumentation
+        // the Mono is already wrapped at this point, but doesn't read the ON_SPAN_END until
+        // the Mono is resolved, which is after this point
       }
     }
   }

diff --git a/.../opentelemetry/javaagent/instrumentation/spring/webflux/server/RouteOnSuccessOrError.java b/.../opentelemetry/javaagent/instrumentation/spring/webflux/server/RouteOnSuccessOrError.java
@@ -9,12 +9,10 @@
 import io.opentelemetry.context.Context;
 import io.opentelemetry.instrumentation.api.servlet.ServletContextPath;
 import io.opentelemetry.instrumentation.api.tracer.ServerSpan;
-import io.opentelemetry.javaagent.instrumentation.spring.webflux.SpringWebfluxConfig;
 import java.util.function.BiConsumer;
 import java.util.regex.Pattern;
 import org.springframework.web.reactive.function.server.HandlerFunction;
 import org.springframework.web.reactive.function.server.RouterFunction;
-import org.springframework.web.reactive.function.server.ServerRequest;
 
 public class RouteOnSuccessOrError implements BiConsumer<HandlerFunction<?>, Throwable> {
 
@@ -23,26 +21,19 @@ public class RouteOnSuccessOrError implements BiConsumer<HandlerFunction<?>, Thr
   private static final Pattern METHOD_REGEX =
       Pattern.compile("^(GET|HEAD|POST|PUT|DELETE|CONNECT|OPTIONS|TRACE|PATCH) ");
 
-  private final RouterFunction routerFunction;
-  private final ServerRequest serverRequest;
+  private final RouterFunction<?> routerFunction;
 
-  public RouteOnSuccessOrError(RouterFunction routerFunction, ServerRequest serverRequest) {
+  public RouteOnSuccessOrError(RouterFunction<?> routerFunction) {
     this.routerFunction = routerFunction;
-    this.serverRequest = serverRequest;
   }
 
   @Override
   public void accept(HandlerFunction<?> handler, Throwable throwable) {
     if (handler != null) {
       String predicateString = parsePredicateString();
       if (predicateString != null) {
-        Context context = (Context) serverRequest.attributes().get(AdviceUtils.CONTEXT_ATTRIBUTE);
+        Context context = Context.current();
         if (context != null) {
-          if (SpringWebfluxConfig.captureExperimentalSpanAttributes()) {
-            Span span = Span.fromContext(context);
-            span.setAttribute("spring-webflux.request.predicate", predicateString);
-          }
-
           Span serverSpan = ServerSpan.fromContextOrNull(context);
           if (serverSpan != null) {
             serverSpan.updateName(ServletContextPath.prepend(context, parseRoute(predicateString)));

diff --git a/...emetry/javaagent/instrumentation/spring/webflux/server/RouterFunctionInstrumentation.java b/...emetry/javaagent/instrumentation/spring/webflux/server/RouterFunctionInstrumentation.java
@@ -22,7 +22,6 @@
 import net.bytebuddy.matcher.ElementMatcher;
 import org.springframework.web.reactive.function.server.HandlerFunction;
 import org.springframework.web.reactive.function.server.RouterFunction;
-import org.springframework.web.reactive.function.server.ServerRequest;
 import reactor.core.publisher.Mono;
 
 public class RouterFunctionInstrumentation implements TypeInstrumentation {
@@ -64,14 +63,11 @@ public static class RouteAdvice {
 
     @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
     public static void methodExit(
-        @Advice.This RouterFunction thiz,
-        @Advice.Argument(0) ServerRequest serverRequest,
+        @Advice.This RouterFunction<?> thiz,
         @Advice.Return(readOnly = false) Mono<HandlerFunction<?>> result,
         @Advice.Thrown Throwable throwable) {
       if (throwable == null) {
-        result = result.doOnSuccessOrError(new RouteOnSuccessOrError(thiz, serverRequest));
-      } else {
-        AdviceUtils.finishSpanIfPresent(serverRequest, throwable);
+        result = result.doOnSuccessOrError(new RouteOnSuccessOrError(thiz));
       }
     }
   }

diff --git a/...emetry/javaagent/instrumentation/spring/webflux/server/SpringWebfluxHttpServerTracer.java b/...emetry/javaagent/instrumentation/spring/webflux/server/SpringWebfluxHttpServerTracer.java