Use context instead of request attributes for servlet async instrumentation

Author: laurit

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

@@ -5,12 +5,15 @@
 
 package io.opentelemetry.javaagent.instrumentation.servlet;
 
+import static io.opentelemetry.context.ContextKey.named;
 import static io.opentelemetry.instrumentation.api.semconv.http.HttpServerRouteSource.SERVER;
 import static io.opentelemetry.instrumentation.api.semconv.http.HttpServerRouteSource.SERVER_FILTER;
 
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.api.trace.SpanContext;
 import io.opentelemetry.context.Context;
+import io.opentelemetry.context.ContextKey;
+import io.opentelemetry.context.ImplicitContextKeyed;
 import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
 import io.opentelemetry.instrumentation.api.instrumenter.LocalRootSpan;
 import io.opentelemetry.instrumentation.api.semconv.http.HttpServerRoute;
@@ -21,6 +24,7 @@
 import io.opentelemetry.semconv.incubating.EnduserIncubatingAttributes;
 import java.security.Principal;
 import java.util.function.Function;
+import javax.annotation.Nullable;
 
 @SuppressWarnings("deprecation") // using deprecated semconv
 public abstract class BaseServletHelper<REQUEST, RESPONSE> {
@@ -61,6 +65,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 +76,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 +96,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 +136,7 @@ private void captureRequestParameters(Span serverSpan, REQUEST request) {
       return;
     }
 
-    parameterExtractor.setAttributes(request, (key, value) -> serverSpan.setAttribute(key, value));
+    parameterExtractor.setAttributes(request, serverSpan::setAttribute);
   }
 
   /**
@@ -169,4 +180,66 @@ public boolean needsRescoping(Context currentContext, Context attachedContext) {
   private static boolean sameTrace(Span oneSpan, Span otherSpan) {
     return oneSpan.getSpanContext().getTraceId().equals(otherSpan.getSpanContext().getTraceId());
   }
+
+  protected static 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/ServletHelper.java

@@ -10,12 +10,6 @@
 import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
 
 public class ServletHelper<REQUEST, RESPONSE> extends BaseServletHelper<REQUEST, RESPONSE> {
-  private static final String ASYNC_LISTENER_ATTRIBUTE =
-      ServletHelper.class.getName() + ".AsyncListener";
-  private static final String ASYNC_LISTENER_RESPONSE_ATTRIBUTE =
-      ServletHelper.class.getName() + ".AsyncListenerResponse";
-  public static final String ASYNC_EXCEPTION_ATTRIBUTE =
-      ServletHelper.class.getName() + ".AsyncException";
   public static final String CONTEXT_ATTRIBUTE = ServletHelper.class.getName() + ".Context";
 
   public ServletHelper(
@@ -87,14 +81,14 @@ public boolean mustEndOnHandlerMethodExit(REQUEST request) {
    * is not possible to access response from async event in listeners.
    */
   public void setAsyncListenerResponse(REQUEST request, RESPONSE response) {
-    accessor.setRequestAttribute(request, ASYNC_LISTENER_RESPONSE_ATTRIBUTE, response);
+    Context context = getServerContext(request);
+    ServletAsyncContext.setAsyncListenerResponse(context, response);
   }
 
+  @SuppressWarnings("unchecked")
   public RESPONSE getAsyncListenerResponse(REQUEST request) {
-    @SuppressWarnings("unchecked")
-    RESPONSE response =
-        (RESPONSE) accessor.getRequestAttribute(request, ASYNC_LISTENER_RESPONSE_ATTRIBUTE);
-    return response;
+    Context context = getServerContext(request);
+    return (RESPONSE) ServletAsyncContext.getAsyncListenerResponse(context);
   }
 
   public void attachAsyncListener(REQUEST request) {
@@ -113,23 +107,26 @@ private void attachAsyncListener(ServletRequestContext<REQUEST> requestContext)
           request,
           new AsyncRequestCompletionListener<>(this, instrumenter, requestContext, context),
           response);
-      accessor.setRequestAttribute(request, ASYNC_LISTENER_ATTRIBUTE, true);
+      ServletAsyncContext.setAsyncListenerAttached(context, true);
     }
   }
 
   public boolean isAsyncListenerAttached(REQUEST request) {
-    return accessor.getRequestAttribute(request, ASYNC_LISTENER_ATTRIBUTE) != null;
+    Context context = getServerContext(request);
+    return ServletAsyncContext.isAsyncListenerAttached(context);
   }
 
   public Runnable wrapAsyncRunnable(REQUEST request, Runnable runnable) {
     return AsyncRunnableWrapper.wrap(this, request, runnable);
   }
 
   public void recordAsyncException(REQUEST request, Throwable throwable) {
-    accessor.setRequestAttribute(request, ASYNC_EXCEPTION_ATTRIBUTE, throwable);
+    Context context = getServerContext(request);
+    ServletAsyncContext.recordAsyncException(context, throwable);
   }
 
   public Throwable getAsyncException(REQUEST request) {
-    return (Throwable) accessor.getRequestAttribute(request, ASYNC_EXCEPTION_ATTRIBUTE);
+    Context context = getServerContext(request);
+    return ServletAsyncContext.getAsyncException(context);
   }
 }