Use context instead of request attributes for servlet async instrumentation

instrumentation/jetty/jetty-11.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/jetty/v11_0/Jetty11HandlerAdvice.java

@@ -45,7 +45,7 @@ public static void onEnter(
     scope = context.makeCurrent();
 
     // Must be set here since Jetty handlers can use startAsync outside of servlet scope.
-    helper().setAsyncListenerResponse(request, response);
+    helper().setAsyncListenerResponse(context, response);
 
     HttpServerResponseCustomizerHolder.getCustomizer()
         .customize(context, response, Jetty11ResponseMutator.INSTANCE);

instrumentation/jetty/jetty-12.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/jetty/v12_0/Jetty12Helper.java

@@ -8,7 +8,7 @@
 import io.opentelemetry.context.Context;
 import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
 import io.opentelemetry.javaagent.bootstrap.servlet.AppServerBridge;
-import io.opentelemetry.javaagent.instrumentation.servlet.ServletHelper;
+import io.opentelemetry.javaagent.bootstrap.servlet.ServletAsyncContext;
 import javax.annotation.Nullable;
 import org.eclipse.jetty.server.HttpStream;
 import org.eclipse.jetty.server.Request;
@@ -54,7 +54,7 @@ public void end(Context context, Request request, Response response, @Nullable T
       error = AppServerBridge.getException(context);
     }
     if (error == null) {
-      error = (Throwable) request.getAttribute(ServletHelper.ASYNC_EXCEPTION_ATTRIBUTE);
+      error = ServletAsyncContext.getAsyncException(context);
     }
 
     instrumenter.end(context, request, response, error);

instrumentation/jetty/jetty-8.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/jetty/v8_0/Jetty8HandlerAdvice.java

@@ -45,7 +45,7 @@ public static void onEnter(
     scope = context.makeCurrent();
 
     // Must be set here since Jetty handlers can use startAsync outside of servlet scope.
-    helper().setAsyncListenerResponse(request, response);
+    helper().setAsyncListenerResponse(context, response);
 
     HttpServerResponseCustomizerHolder.getCustomizer()
         .customize(context, response, Jetty8ResponseMutator.INSTANCE);

instrumentation/jetty/jetty-common/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/jetty/common/JettyHelper.java

@@ -41,7 +41,7 @@ public void end(
     }
 
     ServletResponseContext<RESPONSE> responseContext = new ServletResponseContext<>(response);
-    if (throwable != null || mustEndOnHandlerMethodExit(request)) {
+    if (throwable != null || mustEndOnHandlerMethodExit(context)) {
       instrumenter.end(context, requestContext, responseContext, throwable);
     }
   }

instrumentation/liberty/liberty-20.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/liberty/LibertyHelper.java

@@ -40,7 +40,7 @@ public void end(
     }
 
     ServletResponseContext<RESPONSE> responseContext = new ServletResponseContext<>(response);
-    if (throwable != null || mustEndOnHandlerMethodExit(request)) {
+    if (throwable != null || mustEndOnHandlerMethodExit(context)) {
       instrumenter.end(context, requestContext, responseContext, throwable);
     }
   }

instrumentation/liberty/liberty-20.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/liberty/LibertyWebAppInstrumentation.java

@@ -125,7 +125,7 @@ public static void onEnter() {
 
       // Must be set here since Liberty RequestProcessors can use startAsync outside of servlet
       // scope.
-      helper().setAsyncListenerResponse(requestInfo.getRequest(), requestInfo.getResponse());
+      helper().setAsyncListenerResponse(context, requestInfo.getResponse());
 
       HttpServerResponseCustomizerHolder.getCustomizer()
           .customize(context, requestInfo.getResponse(), Servlet3Accessor.INSTANCE);

instrumentation/servlet/servlet-3.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/servlet/v3_0/Servlet3Advice.java

@@ -60,7 +60,7 @@ public static void onEnter(
     requestContext = new ServletRequestContext<>(httpServletRequest, servletOrFilter);
     if (attachedContext == null && helper().shouldStart(currentContext, requestContext)) {
       context = helper().start(currentContext, requestContext);
-      helper().setAsyncListenerResponse(httpServletRequest, (HttpServletResponse) response);
+      helper().setAsyncListenerResponse(context, (HttpServletResponse) response);
 
       contextToUpdate = context;
     } else if (attachedContext != null

instrumentation/servlet/servlet-3.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/servlet/v3_0/Servlet3AsyncContextStartAdvice.java

@@ -7,21 +7,13 @@
 
 import static io.opentelemetry.javaagent.instrumentation.servlet.v3_0.Servlet3Singletons.helper;
 
-import javax.servlet.AsyncContext;
-import javax.servlet.ServletRequest;
-import javax.servlet.http.HttpServletRequest;
 import net.bytebuddy.asm.Advice;
 
 @SuppressWarnings("unused")
 public class Servlet3AsyncContextStartAdvice {
 
   @Advice.OnMethodEnter(suppress = Throwable.class)
-  public static void start(
-      @Advice.This AsyncContext asyncContext,
-      @Advice.Argument(value = 0, readOnly = false) Runnable runnable) {
-    ServletRequest request = asyncContext.getRequest();
-    if (request instanceof HttpServletRequest) {
-      runnable = helper().wrapAsyncRunnable((HttpServletRequest) request, runnable);
-    }
+  public static void start(@Advice.Argument(value = 0, readOnly = false) Runnable runnable) {
+    runnable = helper().wrapAsyncRunnable(runnable);
   }
 }

instrumentation/servlet/servlet-3.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/servlet/v3_0/Servlet3AsyncStartAdvice.java

@@ -8,6 +8,7 @@
 import static io.opentelemetry.javaagent.instrumentation.servlet.v3_0.Servlet3Singletons.helper;
 
 import io.opentelemetry.javaagent.bootstrap.CallDepth;
+import io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge;
 import javax.servlet.AsyncContext;
 import javax.servlet.ServletRequest;
 import javax.servlet.http.HttpServletRequest;
@@ -36,9 +37,7 @@ public static void startAsyncExit(
     if (servletRequest instanceof HttpServletRequest) {
       HttpServletRequest request = (HttpServletRequest) servletRequest;
 
-      if (!helper().isAsyncListenerAttached(request)) {
-        helper().attachAsyncListener(request);
-      }
+      helper().attachAsyncListener(request, Java8BytecodeBridge.currentContext());
     }
   }
 }

instrumentation/servlet/servlet-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/servlet/v5_0/async/AsyncContextStartAdvice.java

@@ -7,21 +7,13 @@
 
 import static io.opentelemetry.javaagent.instrumentation.servlet.v5_0.Servlet5Singletons.helper;
 
-import jakarta.servlet.AsyncContext;
-import jakarta.servlet.ServletRequest;
-import jakarta.servlet.http.HttpServletRequest;
 import net.bytebuddy.asm.Advice;
 
 @SuppressWarnings("unused")
 public class AsyncContextStartAdvice {
 
   @Advice.OnMethodEnter(suppress = Throwable.class)
-  public static void start(
-      @Advice.This AsyncContext asyncContext,
-      @Advice.Argument(value = 0, readOnly = false) Runnable runnable) {
-    ServletRequest request = asyncContext.getRequest();
-    if (request instanceof HttpServletRequest) {
-      runnable = helper().wrapAsyncRunnable((HttpServletRequest) request, runnable);
-    }
+  public static void start(@Advice.Argument(value = 0, readOnly = false) Runnable runnable) {
+    runnable = helper().wrapAsyncRunnable(runnable);
   }
 }

instrumentation/servlet/servlet-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/servlet/v5_0/async/AsyncStartAdvice.java

@@ -8,6 +8,7 @@
 import static io.opentelemetry.javaagent.instrumentation.servlet.v5_0.Servlet5Singletons.helper;
 
 import io.opentelemetry.javaagent.bootstrap.CallDepth;
+import io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge;
 import jakarta.servlet.AsyncContext;
 import jakarta.servlet.http.HttpServletRequest;
 import net.bytebuddy.asm.Advice;
@@ -34,9 +35,7 @@ public static void startAsyncExit(
     }
 
     if (request != null) {
-      if (!helper().isAsyncListenerAttached(request)) {
-        helper().attachAsyncListener(request);
-      }
+      helper().attachAsyncListener(request, Java8BytecodeBridge.currentContext());
     }
   }
 }

instrumentation/servlet/servlet-5.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/servlet/v5_0/service/JakartaServletServiceAdvice.java

@@ -63,7 +63,7 @@ public static void onEnter(
     requestContext = new ServletRequestContext<>(httpServletRequest, servletOrFilter);
     if (attachedContext == null && helper().shouldStart(currentContext, requestContext)) {
       context = helper().start(currentContext, requestContext);
-      helper().setAsyncListenerResponse(httpServletRequest, (HttpServletResponse) response);
+      helper().setAsyncListenerResponse(context, (HttpServletResponse) response);
 
       contextToUpdate = context;
     } else if (attachedContext != null

instrumentation/servlet/servlet-common/bootstrap/src/main/java/io/opentelemetry/javaagent/bootstrap/servlet/AppServerBridge.java

@@ -125,7 +125,9 @@ public Builder captureServletAttributes() {
      * @return new context with AppServerBridge attached.
      */
     public Context init(Context context) {
-      return context.with(AppServerBridge.CONTEXT_KEY, new AppServerBridge(this));
+      Context result = context.with(AppServerBridge.CONTEXT_KEY, new AppServerBridge(this));
+      result = ServletAsyncContext.init(result);
+      return result;
     }
   }
 }

instrumentation/servlet/servlet-common/bootstrap/src/main/java/io/opentelemetry/javaagent/bootstrap/servlet/ServletAsyncContext.java

@@ -0,0 +1,75 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.bootstrap.servlet;
+
+import static io.opentelemetry.context.ContextKey.named;
+
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.ContextKey;
+import io.opentelemetry.context.ImplicitContextKeyed;
+import javax.annotation.Nullable;
+
+public class ServletAsyncContext implements ImplicitContextKeyed {
+  private static final ContextKey<ServletAsyncContext> CONTEXT_KEY =
+      named("opentelemetry-servlet-async-context");
+
+  private boolean isAsyncListenerAttached;
+  private Throwable throwable;
+  private Object response;
+
+  public static Context init(Context context) {
+    if (context.get(CONTEXT_KEY) != null) {
+      return context;
+    }
+    return context.with(new ServletAsyncContext());
+  }
+
+  @Nullable
+  public static ServletAsyncContext get(@Nullable Context context) {
+    return context != null ? context.get(CONTEXT_KEY) : null;
+  }
+
+  public static boolean isAsyncListenerAttached(@Nullable Context context) {
+    ServletAsyncContext servletAsyncContext = get(context);
+    return servletAsyncContext != null && servletAsyncContext.isAsyncListenerAttached;
+  }
+
+  public static void setAsyncListenerAttached(@Nullable Context context, boolean value) {
+    ServletAsyncContext servletAsyncContext = get(context);
+    if (servletAsyncContext != null) {
+      servletAsyncContext.isAsyncListenerAttached = value;
+    }
+  }
+
+  public static Throwable getAsyncException(@Nullable Context context) {
+    ServletAsyncContext servletAsyncContext = get(context);
+    return servletAsyncContext != null ? servletAsyncContext.throwable : null;
+  }
+
+  public static void recordAsyncException(@Nullable Context context, Throwable throwable) {
+    ServletAsyncContext servletAsyncContext = get(context);
+    if (servletAsyncContext != null) {
+      servletAsyncContext.throwable = throwable;
+    }
+  }
+
+  public static Object getAsyncListenerResponse(@Nullable Context context) {
+    ServletAsyncContext servletAsyncContext = get(context);
+    return servletAsyncContext != null ? servletAsyncContext.response : null;
+  }
+
+  public static void setAsyncListenerResponse(@Nullable Context context, Object response) {
+    ServletAsyncContext servletAsyncContext = get(context);
+    if (servletAsyncContext != null) {
+      servletAsyncContext.response = response;
+    }
+  }
+
+  @Override
+  public Context storeInContext(Context context) {
+    return context.with(CONTEXT_KEY, this);
+  }
+}

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

@@ -33,18 +33,18 @@ public AsyncRequestCompletionListener(
   public void onComplete(RESPONSE response) {
     if (responseHandled.compareAndSet(false, true)) {
       ServletResponseContext<RESPONSE> responseContext = new ServletResponseContext<>(response);
-      Throwable throwable = servletHelper.getAsyncException(requestContext.request());
+      Throwable throwable = servletHelper.getAsyncException(context);
       instrumenter.end(context, requestContext, responseContext, throwable);
     }
   }
 
   @Override
   public void onTimeout(long timeout) {
     if (responseHandled.compareAndSet(false, true)) {
-      RESPONSE response = servletHelper.getAsyncListenerResponse(requestContext.request());
+      RESPONSE response = servletHelper.getAsyncListenerResponse(context);
       ServletResponseContext<RESPONSE> responseContext = new ServletResponseContext<>(response);
       responseContext.setTimeout(timeout);
-      Throwable throwable = servletHelper.getAsyncException(requestContext.request());
+      Throwable throwable = servletHelper.getAsyncException(context);
       instrumenter.end(context, requestContext, responseContext, throwable);
     }
   }

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

@@ -5,32 +5,32 @@
 
 package io.opentelemetry.javaagent.instrumentation.servlet;
 
+import io.opentelemetry.context.Context;
+
 public class AsyncRunnableWrapper<REQUEST> implements Runnable {
   private final ServletHelper<REQUEST, ?> helper;
-  private final REQUEST request;
   private final Runnable runnable;
+  private final Context context;
 
-  private AsyncRunnableWrapper(
-      ServletHelper<REQUEST, ?> helper, REQUEST request, Runnable runnable) {
+  private AsyncRunnableWrapper(ServletHelper<REQUEST, ?> helper, Runnable runnable) {
     this.helper = helper;
-    this.request = request;
     this.runnable = runnable;
+    this.context = Context.current();
   }
 
-  public static <REQUEST> Runnable wrap(
-      ServletHelper<REQUEST, ?> helper, REQUEST request, Runnable runnable) {
+  public static <REQUEST> Runnable wrap(ServletHelper<REQUEST, ?> helper, Runnable runnable) {
     if (runnable == null || runnable instanceof AsyncRunnableWrapper) {
       return runnable;
     }
-    return new AsyncRunnableWrapper<>(helper, request, runnable);
+    return new AsyncRunnableWrapper<>(helper, runnable);
   }
 
   @Override
   public void run() {
     try {
       runnable.run();
     } catch (Throwable throwable) {
-      helper.recordAsyncException(request, throwable);
+      helper.recordAsyncException(context, throwable);
       throw throwable;
     }
   }

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

@@ -17,6 +17,7 @@
 import io.opentelemetry.javaagent.bootstrap.internal.AgentCommonConfig;
 import io.opentelemetry.javaagent.bootstrap.servlet.AppServerBridge;
 import io.opentelemetry.javaagent.bootstrap.servlet.MappingResolver;
+import io.opentelemetry.javaagent.bootstrap.servlet.ServletAsyncContext;
 import io.opentelemetry.javaagent.bootstrap.servlet.ServletContextPath;
 import io.opentelemetry.semconv.incubating.EnduserIncubatingAttributes;
 import java.security.Principal;
@@ -61,6 +62,7 @@ public Context start(Context parentContext, ServletRequestContext<REQUEST> reque
     accessor.setRequestAttribute(request, "span_id", spanContext.getSpanId());
 
     context = addServletContextPath(context, request);
+    context = addAsyncContext(context);
 
     attachServerContext(context, request);
 
@@ -71,6 +73,10 @@ protected Context addServletContextPath(Context context, REQUEST request) {
     return ServletContextPath.init(context, contextPathExtractor, request);
   }
 
+  protected Context addAsyncContext(Context context) {
+    return ServletAsyncContext.init(context);
+  }
+
   public Context getServerContext(REQUEST request) {
     Object context = accessor.getRequestAttribute(request, ServletHelper.CONTEXT_ATTRIBUTE);
     return context instanceof Context ? (Context) context : null;
@@ -87,6 +93,8 @@ public void recordException(Context context, Throwable throwable) {
   public Context updateContext(
       Context context, REQUEST request, MappingResolver mappingResolver, boolean servlet) {
     Context result = addServletContextPath(context, request);
+    result = addAsyncContext(result);
+
     if (mappingResolver != null) {
       HttpServerRoute.update(
           result, servlet ? SERVER : SERVER_FILTER, spanNameProvider, mappingResolver, request);
@@ -125,7 +133,7 @@ private void captureRequestParameters(Span serverSpan, REQUEST request) {
       return;
     }
 
-    parameterExtractor.setAttributes(request, (key, value) -> serverSpan.setAttribute(key, value));
+    parameterExtractor.setAttributes(request, serverSpan::setAttribute);
   }
 
   /**