Use context instead of request attributes for servlet async instrumentation (#13493)

This commit is contained in:
Lauri Tulmin 2025-04-10 02:38:36 +03:00 committed by GitHub
parent 33024dcc64
commit 940bce8784
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
21 changed files with 159 additions and 114 deletions

View File

@ -45,7 +45,7 @@ public class Jetty11HandlerAdvice {
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);

View File

@ -8,7 +8,7 @@ package io.opentelemetry.javaagent.instrumentation.jetty.v12_0;
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 class Jetty12Helper {
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);

View File

@ -45,7 +45,7 @@ public class Jetty8HandlerAdvice {
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);

View File

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

View File

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

View File

@ -125,7 +125,7 @@ public class LibertyWebAppInstrumentation implements TypeInstrumentation {
// 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);

View File

@ -60,7 +60,7 @@ public class Servlet3Advice {
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

View File

@ -7,21 +7,13 @@ package io.opentelemetry.javaagent.instrumentation.servlet.v3_0;
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);
}
}

View File

@ -8,6 +8,7 @@ package io.opentelemetry.javaagent.instrumentation.servlet.v3_0;
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 class Servlet3AsyncStartAdvice {
if (servletRequest instanceof HttpServletRequest) {
HttpServletRequest request = (HttpServletRequest) servletRequest;
if (!helper().isAsyncListenerAttached(request)) {
helper().attachAsyncListener(request);
}
helper().attachAsyncListener(request, Java8BytecodeBridge.currentContext());
}
}
}

View File

@ -7,21 +7,13 @@ package io.opentelemetry.javaagent.instrumentation.servlet.v5_0.async;
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);
}
}

View File

@ -8,6 +8,7 @@ package io.opentelemetry.javaagent.instrumentation.servlet.v5_0.async;
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 class AsyncStartAdvice {
}
if (request != null) {
if (!helper().isAsyncListenerAttached(request)) {
helper().attachAsyncListener(request);
}
helper().attachAsyncListener(request, Java8BytecodeBridge.currentContext());
}
}
}

View File

@ -63,7 +63,7 @@ public class JakartaServletServiceAdvice {
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

View File

@ -125,7 +125,9 @@ public class AppServerBridge {
* @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;
}
}
}

View File

@ -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);
}
}

View File

@ -33,7 +33,7 @@ public class AsyncRequestCompletionListener<REQUEST, RESPONSE>
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);
}
}
@ -41,10 +41,10 @@ public class AsyncRequestCompletionListener<REQUEST, RESPONSE>
@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);
}
}

View File

@ -5,24 +5,24 @@
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
@ -30,7 +30,7 @@ public class AsyncRunnableWrapper<REQUEST> implements Runnable {
try {
runnable.run();
} catch (Throwable throwable) {
helper.recordAsyncException(request, throwable);
helper.recordAsyncException(context, throwable);
throw throwable;
}
}

View File

@ -17,6 +17,7 @@ import io.opentelemetry.instrumentation.api.semconv.http.HttpServerRoute;
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 abstract class BaseServletHelper<REQUEST, RESPONSE> {
accessor.setRequestAttribute(request, "span_id", spanContext.getSpanId());
context = addServletContextPath(context, request);
context = addAsyncContext(context);
attachServerContext(context, request);
@ -71,6 +73,10 @@ public abstract class BaseServletHelper<REQUEST, RESPONSE> {
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 abstract class BaseServletHelper<REQUEST, RESPONSE> {
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 @@ public abstract class BaseServletHelper<REQUEST, RESPONSE> {
return;
}
parameterExtractor.setAttributes(request, (key, value) -> serverSpan.setAttribute(key, value));
parameterExtractor.setAttributes(request, serverSpan::setAttribute);
}
/**

View File

@ -8,14 +8,9 @@ package io.opentelemetry.javaagent.instrumentation.servlet;
import io.opentelemetry.context.Context;
import io.opentelemetry.context.Scope;
import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
import io.opentelemetry.javaagent.bootstrap.servlet.ServletAsyncContext;
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(
@ -43,10 +38,10 @@ public class ServletHelper<REQUEST, RESPONSE> extends BaseServletHelper<REQUEST,
// instrumentation and we have an uncaught throwable. Let's add it to the current span.
if (throwable != null) {
recordException(currentContext, throwable);
if (!mustEndOnHandlerMethodExit(request)) {
if (!mustEndOnHandlerMethodExit(currentContext)) {
// We could be inside async dispatch. Unlike tomcat jetty does not call
// ServletAsyncListener.onError when exception is thrown inside async dispatch.
recordAsyncException(request, throwable);
recordAsyncException(currentContext, throwable);
}
}
// also capture request parameters as servlet attributes
@ -58,7 +53,7 @@ public class ServletHelper<REQUEST, RESPONSE> extends BaseServletHelper<REQUEST,
}
ServletResponseContext<RESPONSE> responseContext = new ServletResponseContext<>(response);
if (throwable != null || mustEndOnHandlerMethodExit(request)) {
if (throwable != null || mustEndOnHandlerMethodExit(context)) {
instrumenter.end(context, requestContext, responseContext, throwable);
}
}
@ -68,8 +63,8 @@ public class ServletHelper<REQUEST, RESPONSE> extends BaseServletHelper<REQUEST,
* that started a span must also end it, even if no error was detected. Extracted as a separate
* method to avoid duplicating the comments on the logic behind this choice.
*/
public boolean mustEndOnHandlerMethodExit(REQUEST request) {
if (isAsyncListenerAttached(request)) {
public boolean mustEndOnHandlerMethodExit(Context context) {
if (isAsyncListenerAttached(context)) {
// This request is handled asynchronously and startAsync instrumentation has already attached
// the listener.
return false;
@ -82,54 +77,47 @@ public class ServletHelper<REQUEST, RESPONSE> extends BaseServletHelper<REQUEST,
}
/**
* Response object must be attached to a request prior to {@link
* #attachAsyncListener(ServletRequestContext)} being called, as otherwise in some environments it
* is not possible to access response from async event in listeners.
* Response object must be attached to a request prior to {@link #attachAsyncListener(REQUEST,
* Context)} being called, as otherwise in some environments it 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);
public void setAsyncListenerResponse(Context context, RESPONSE response) {
ServletAsyncContext.setAsyncListenerResponse(context, response);
}
public RESPONSE getAsyncListenerResponse(REQUEST request) {
@SuppressWarnings("unchecked")
RESPONSE response =
(RESPONSE) accessor.getRequestAttribute(request, ASYNC_LISTENER_RESPONSE_ATTRIBUTE);
return response;
public RESPONSE getAsyncListenerResponse(Context context) {
return (RESPONSE) ServletAsyncContext.getAsyncListenerResponse(context);
}
public void attachAsyncListener(REQUEST request) {
public void attachAsyncListener(REQUEST request, Context context) {
if (isAsyncListenerAttached(context)) {
return;
}
Object response = getAsyncListenerResponse(context);
ServletRequestContext<REQUEST> requestContext = new ServletRequestContext<>(request, null);
attachAsyncListener(requestContext);
}
private void attachAsyncListener(ServletRequestContext<REQUEST> requestContext) {
REQUEST request = requestContext.request();
Context context = getServerContext(request);
if (context != null) {
Object response = getAsyncListenerResponse(request);
accessor.addRequestAsyncListener(
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;
private static boolean isAsyncListenerAttached(Context context) {
return ServletAsyncContext.isAsyncListenerAttached(context);
}
public Runnable wrapAsyncRunnable(REQUEST request, Runnable runnable) {
return AsyncRunnableWrapper.wrap(this, request, runnable);
public Runnable wrapAsyncRunnable(Runnable runnable) {
return AsyncRunnableWrapper.wrap(this, runnable);
}
public void recordAsyncException(REQUEST request, Throwable throwable) {
accessor.setRequestAttribute(request, ASYNC_EXCEPTION_ATTRIBUTE, throwable);
public void recordAsyncException(Context context, Throwable throwable) {
ServletAsyncContext.recordAsyncException(context, throwable);
}
public Throwable getAsyncException(REQUEST request) {
return (Throwable) accessor.getRequestAttribute(request, ASYNC_EXCEPTION_ATTRIBUTE);
public Throwable getAsyncException(Context context) {
return ServletAsyncContext.getAsyncException(context);
}
}

View File

@ -7,8 +7,8 @@ package io.opentelemetry.javaagent.instrumentation.tomcat.v10_0;
import static io.opentelemetry.javaagent.instrumentation.tomcat.v10_0.Tomcat10Singletons.helper;
import io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge;
import net.bytebuddy.asm.Advice;
import org.apache.coyote.Request;
import org.apache.coyote.Response;
@SuppressWarnings("unused")
@ -16,12 +16,10 @@ public class Tomcat10AttachResponseAdvice {
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
public static void attachResponse(
@Advice.Argument(0) Request request,
@Advice.Argument(2) Response response,
@Advice.Return boolean success) {
@Advice.Argument(2) Response response, @Advice.Return boolean success) {
if (success) {
helper().attachResponseToRequest(request, response);
helper().attachResponseToRequest(Java8BytecodeBridge.currentContext(), response);
}
}
}

View File

@ -7,8 +7,8 @@ package io.opentelemetry.javaagent.instrumentation.tomcat.v7_0;
import static io.opentelemetry.javaagent.instrumentation.tomcat.v7_0.Tomcat7Singletons.helper;
import io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge;
import net.bytebuddy.asm.Advice;
import org.apache.coyote.Request;
import org.apache.coyote.Response;
@SuppressWarnings("unused")
@ -16,12 +16,10 @@ public class Tomcat7AttachResponseAdvice {
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
public static void attachResponse(
@Advice.Argument(0) Request request,
@Advice.Argument(2) Response response,
@Advice.Return boolean success) {
@Advice.Argument(2) Response response, @Advice.Return boolean success) {
if (success) {
helper().attachResponseToRequest(request, response);
helper().attachResponseToRequest(Java8BytecodeBridge.currentContext(), response);
}
}
}

View File

@ -49,22 +49,16 @@ public class TomcatHelper<REQUEST, RESPONSE> {
throwable = AppServerBridge.getException(context);
}
if (throwable != null || mustEndOnHandlerMethodExit(request)) {
if (throwable != null || servletHelper.mustEndOnHandlerMethodExit(context)) {
instrumenter.end(context, request, response, throwable);
}
}
private boolean mustEndOnHandlerMethodExit(Request request) {
REQUEST servletRequest = servletEntityProvider.getServletRequest(request);
return servletRequest != null && servletHelper.mustEndOnHandlerMethodExit(servletRequest);
}
public void attachResponseToRequest(Request request, Response response) {
REQUEST servletRequest = servletEntityProvider.getServletRequest(request);
public void attachResponseToRequest(Context context, Response response) {
RESPONSE servletResponse = servletEntityProvider.getServletResponse(response);
if (servletRequest != null && servletResponse != null) {
servletHelper.setAsyncListenerResponse(servletRequest, servletResponse);
if (servletResponse != null) {
servletHelper.setAsyncListenerResponse(context, servletResponse);
}
}