Clean up request dispatcher instrumentation (#2724)

instrumentation/servlet/servlet-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/servlet/v5_0/dispatcher/RequestDispatcherAdvice.java

@@ -10,6 +10,7 @@ import static io.opentelemetry.javaagent.instrumentation.servlet.v5_0.dispatcher
 
 import io.opentelemetry.context.Context;
 import io.opentelemetry.context.Scope;
+import io.opentelemetry.javaagent.instrumentation.api.Java8BytecodeBridge;
 import io.opentelemetry.javaagent.instrumentation.servlet.common.dispatcher.RequestDispatcherAdviceHelper;
 import jakarta.servlet.RequestDispatcher;
 import jakarta.servlet.ServletRequest;
@@ -22,48 +23,44 @@ public class RequestDispatcherAdvice {
   public static void start(
       @Advice.Origin Method method,
       @Advice.This RequestDispatcher dispatcher,
-      @Advice.Local("_originalContext") Object originalContext,
+      @Advice.Local("otelRequestContext") Context requestContext,
       @Advice.Local("otelContext") Context context,
       @Advice.Local("otelScope") Scope scope,
       @Advice.Argument(0) ServletRequest request) {
 
-    Context parent =
+    Context currentContext = Java8BytecodeBridge.currentContext();
-        RequestDispatcherAdviceHelper.getStartParentContext(
-            request.getAttribute(CONTEXT_ATTRIBUTE));
 
-    if (parent == null) {
+    Object requestContextAttr = request.getAttribute(CONTEXT_ATTRIBUTE);
+    requestContext = requestContextAttr instanceof Context ? (Context) requestContextAttr : null;
+
+    Context parentContext =
+        RequestDispatcherAdviceHelper.getStartParentContext(currentContext, requestContext);
+    if (parentContext == null) {
       return;
     }
 
-    try (Scope ignored = parent.makeCurrent()) {
+    context = tracer().startSpan(parentContext, method);
-      context = tracer().startSpan(method);
 
-      // save the original servlet span before overwriting the request attribute, so that it can
+    // this tells the dispatched servlet to use the current span as the parent for its work
-      // be restored on method exit
+    request.setAttribute(CONTEXT_ATTRIBUTE, context);
-      originalContext = request.getAttribute(CONTEXT_ATTRIBUTE);
 
-      // this tells the dispatched servlet to use the current span as the parent for its work
-      request.setAttribute(CONTEXT_ATTRIBUTE, context);
-    }
     scope = context.makeCurrent();
   }
 
   @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
   public static void stop(
-      @Advice.Local("_originalContext") Object originalContext,
       @Advice.Argument(0) ServletRequest request,
+      @Advice.Local("otelRequestContext") Context requestContext,
       @Advice.Local("otelContext") Context context,
       @Advice.Local("otelScope") Scope scope,
       @Advice.Thrown Throwable throwable) {
 
     scope.close();
 
-    // restore the original servlet span
+    if (requestContext != null) {
-    // since spanWithScope is non-null here, originalContext must have been set with the
+      // restore the original request context
-    // prior
+      request.setAttribute(CONTEXT_ATTRIBUTE, requestContext);
-    // servlet span (as opposed to remaining unset)
+    }
-    // TODO review this logic. Seems like manual context management
-    request.setAttribute(CONTEXT_ATTRIBUTE, originalContext);
 
     if (throwable != null) {
       tracer().endExceptionally(context, throwable);

instrumentation/servlet/servlet-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/servlet/v5_0/dispatcher/RequestDispatcherTracer.java

@@ -5,6 +5,7 @@
 
 package io.opentelemetry.javaagent.instrumentation.servlet.v5_0.dispatcher;
 
+import io.opentelemetry.api.trace.SpanKind;
 import io.opentelemetry.context.Context;
 import io.opentelemetry.instrumentation.api.tracer.BaseTracer;
 import java.lang.reflect.Method;
@@ -21,7 +22,7 @@ public class RequestDispatcherTracer extends BaseTracer {
     return "io.opentelemetry.javaagent.servlet-5.0";
   }
 
-  public Context startSpan(Method method) {
+  public Context startSpan(Context parentContext, Method method) {
-    return startSpan(spanNameForMethod(method));
+    return startSpan(parentContext, spanNameForMethod(method), SpanKind.INTERNAL);
   }
 }

instrumentation/servlet/servlet-common/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/servlet/common/dispatcher/RequestDispatcherAdviceHelper.java

@@ -9,43 +9,44 @@ import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.api.trace.SpanContext;
 import io.opentelemetry.context.Context;
 import io.opentelemetry.instrumentation.api.tracer.HttpServerTracer;
+import org.checkerframework.checker.nullness.qual.Nullable;
 
 public class RequestDispatcherAdviceHelper {
   /**
    * Determines if the advice for {@link RequestDispatcherInstrumentation} should create a new span
    * and provides the context in which that span should be created.
    *
-   * @param servletContextObject Value of the {@link HttpServerTracer#CONTEXT_ATTRIBUTE} attribute
+   * @param requestContext Value of the {@link HttpServerTracer#CONTEXT_ATTRIBUTE} attribute of the
-   *     of the servlet request.
+   *     servlet request.
    * @return The context in which the advice should create the dispatcher span in. Returns <code>
    *     null</code> in case a new span should not be created.
    */
-  public static Context getStartParentContext(Object servletContextObject) {
+  // TODO (trask) do we need to guard against context leak here?
-    Context parentContext = Context.current();
+  //  this could be simplified by always using currentContext, only falling back to requestContext
+  //  if currentContext does not have a valid span
+  public static @Nullable Context getStartParentContext(
+      Context currentContext, @Nullable Context requestContext) {
+    Span currentSpan = Span.fromContext(currentContext);
+    SpanContext currentSpanContext = currentSpan.getSpanContext();
 
-    Context servletContext =
+    if (!currentSpanContext.isValid()) {
-        servletContextObject instanceof Context ? (Context) servletContextObject : null;
+      return requestContext;
-
-    Span parentSpan = Span.fromContext(parentContext);
-    SpanContext parentSpanContext = parentSpan.getSpanContext();
-    if (!parentSpanContext.isValid() && servletContext == null) {
-      // Don't want to generate a new top-level span
-      return null;
     }
 
-    Span servletSpan = servletContext != null ? Span.fromContext(servletContext) : null;
+    if (requestContext == null) {
-    Context parent;
+      return currentContext;
-    if (servletContext == null
-        || (parentSpanContext.isValid()
-            && servletSpan.getSpanContext().getTraceId().equals(parentSpanContext.getTraceId()))) {
-      // Use the parentSpan if the servletSpan is null or part of the same trace.
-      parent = parentContext;
-    } else {
-      // parentSpan is part of a different trace, so lets ignore it.
-      // This can happen with the way Tomcat does error handling.
-      parent = servletContext;
     }
 
-    return parent;
+    // at this point: currentContext has a valid span and requestContext is not null
+
+    Span requestSpan = Span.fromContext(requestContext);
+    if (requestSpan.getSpanContext().getTraceId().equals(currentSpanContext.getTraceId())) {
+      // currentContext is part of the same trace, so return it
+      return currentContext;
+    }
+
+    // currentContext is part of a different trace, so lets ignore it.
+    // This can happen with the way Tomcat does error handling.
+    return requestContext;
   }
 }

instrumentation/servlet/servlet-javax-common/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/servlet/javax/dispatcher/RequestDispatcherAdvice.java

@@ -10,9 +10,9 @@ import static io.opentelemetry.javaagent.instrumentation.servlet.javax.dispatche
 
 import io.opentelemetry.context.Context;
 import io.opentelemetry.context.Scope;
+import io.opentelemetry.javaagent.instrumentation.api.Java8BytecodeBridge;
 import io.opentelemetry.javaagent.instrumentation.servlet.common.dispatcher.RequestDispatcherAdviceHelper;
 import java.lang.reflect.Method;
-import javax.servlet.RequestDispatcher;
 import javax.servlet.ServletRequest;
 import net.bytebuddy.asm.Advice;
 
@@ -21,49 +21,45 @@ public class RequestDispatcherAdvice {
   @Advice.OnMethodEnter(suppress = Throwable.class)
   public static void start(
       @Advice.Origin Method method,
-      @Advice.This RequestDispatcher dispatcher,
+      @Advice.Local("otelRequestContext") Context requestContext,
-      @Advice.Local("_originalContext") Object originalContext,
       @Advice.Local("otelContext") Context context,
       @Advice.Local("otelScope") Scope scope,
       @Advice.Argument(0) ServletRequest request) {
 
-    Context parent =
+    Context currentContext = Java8BytecodeBridge.currentContext();
-        RequestDispatcherAdviceHelper.getStartParentContext(
-            request.getAttribute(CONTEXT_ATTRIBUTE));
 
-    if (parent == null) {
+    Object requestContextAttr = request.getAttribute(CONTEXT_ATTRIBUTE);
+    requestContext = requestContextAttr instanceof Context ? (Context) requestContextAttr : null;
+
+    Context parentContext =
+        RequestDispatcherAdviceHelper.getStartParentContext(currentContext, requestContext);
+    if (parentContext == null) {
       return;
     }
 
-    try (Scope ignored = parent.makeCurrent()) {
+    context = tracer().startSpan(parentContext, method);
-      context = tracer().startSpan(method);
 
-      // save the original servlet span before overwriting the request attribute, so that it can
+    // this tells the dispatched servlet to use the current span as the parent for its work
-      // be restored on method exit
+    request.setAttribute(CONTEXT_ATTRIBUTE, context);
-      originalContext = request.getAttribute(CONTEXT_ATTRIBUTE);
 
-      // this tells the dispatched servlet to use the current span as the parent for its work
+    // TODO (trask) do we need this, since doing manual propagation above?
-      request.setAttribute(CONTEXT_ATTRIBUTE, context);
-    }
     scope = context.makeCurrent();
   }
 
   @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
   public static void stop(
-      @Advice.Local("_originalContext") Object originalContext,
       @Advice.Argument(0) ServletRequest request,
+      @Advice.Local("otelRequestContext") Context requestContext,
       @Advice.Local("otelContext") Context context,
       @Advice.Local("otelScope") Scope scope,
       @Advice.Thrown Throwable throwable) {
 
     scope.close();
 
-    // restore the original servlet span
+    if (requestContext != null) {
-    // since spanWithScope is non-null here, originalContext must have been set with the
+      // restore the original request context
-    // prior
+      request.setAttribute(CONTEXT_ATTRIBUTE, requestContext);
-    // servlet span (as opposed to remaining unset)
+    }
-    // TODO review this logic. Seems like manual context management
-    request.setAttribute(CONTEXT_ATTRIBUTE, originalContext);
 
     if (throwable != null) {
       tracer().endExceptionally(context, throwable);

instrumentation/servlet/servlet-javax-common/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/servlet/javax/dispatcher/RequestDispatcherTracer.java

@@ -5,6 +5,7 @@
 
 package io.opentelemetry.javaagent.instrumentation.servlet.javax.dispatcher;
 
+import io.opentelemetry.api.trace.SpanKind;
 import io.opentelemetry.context.Context;
 import io.opentelemetry.instrumentation.api.tracer.BaseTracer;
 import java.lang.reflect.Method;
@@ -21,7 +22,7 @@ public class RequestDispatcherTracer extends BaseTracer {
     return "io.opentelemetry.javaagent.servlet-javax-common";
   }
 
-  public Context startSpan(Method method) {
+  public Context startSpan(Context parentContext, Method method) {
-    return startSpan(spanNameForMethod(method));
+    return startSpan(parentContext, spanNameForMethod(method), SpanKind.INTERNAL);
   }
 }