Skip to content

Commit 0de7f49

Browse files
committed
Use context instead of request attributes for servlet async instrumentation
1 parent dd071d2 commit 0de7f49

File tree

2 files changed

+86
-16
lines changed

2 files changed

+86
-16
lines changed

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

+74-1
Original file line numberDiff line numberDiff line change
@@ -5,12 +5,15 @@
55

66
package io.opentelemetry.javaagent.instrumentation.servlet;
77

8+
import static io.opentelemetry.context.ContextKey.named;
89
import static io.opentelemetry.instrumentation.api.semconv.http.HttpServerRouteSource.SERVER;
910
import static io.opentelemetry.instrumentation.api.semconv.http.HttpServerRouteSource.SERVER_FILTER;
1011

1112
import io.opentelemetry.api.trace.Span;
1213
import io.opentelemetry.api.trace.SpanContext;
1314
import io.opentelemetry.context.Context;
15+
import io.opentelemetry.context.ContextKey;
16+
import io.opentelemetry.context.ImplicitContextKeyed;
1417
import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
1518
import io.opentelemetry.instrumentation.api.instrumenter.LocalRootSpan;
1619
import io.opentelemetry.instrumentation.api.semconv.http.HttpServerRoute;
@@ -21,6 +24,7 @@
2124
import io.opentelemetry.semconv.incubating.EnduserIncubatingAttributes;
2225
import java.security.Principal;
2326
import java.util.function.Function;
27+
import javax.annotation.Nullable;
2428

2529
@SuppressWarnings("deprecation") // using deprecated semconv
2630
public abstract class BaseServletHelper<REQUEST, RESPONSE> {
@@ -61,6 +65,7 @@ public Context start(Context parentContext, ServletRequestContext<REQUEST> reque
6165
accessor.setRequestAttribute(request, "span_id", spanContext.getSpanId());
6266

6367
context = addServletContextPath(context, request);
68+
context = addAsyncContext(context);
6469

6570
attachServerContext(context, request);
6671

@@ -71,6 +76,10 @@ protected Context addServletContextPath(Context context, REQUEST request) {
7176
return ServletContextPath.init(context, contextPathExtractor, request);
7277
}
7378

79+
protected Context addAsyncContext(Context context) {
80+
return ServletAsyncContext.init(context);
81+
}
82+
7483
public Context getServerContext(REQUEST request) {
7584
Object context = accessor.getRequestAttribute(request, ServletHelper.CONTEXT_ATTRIBUTE);
7685
return context instanceof Context ? (Context) context : null;
@@ -87,6 +96,8 @@ public void recordException(Context context, Throwable throwable) {
8796
public Context updateContext(
8897
Context context, REQUEST request, MappingResolver mappingResolver, boolean servlet) {
8998
Context result = addServletContextPath(context, request);
99+
result = addAsyncContext(result);
100+
90101
if (mappingResolver != null) {
91102
HttpServerRoute.update(
92103
result, servlet ? SERVER : SERVER_FILTER, spanNameProvider, mappingResolver, request);
@@ -125,7 +136,7 @@ private void captureRequestParameters(Span serverSpan, REQUEST request) {
125136
return;
126137
}
127138

128-
parameterExtractor.setAttributes(request, (key, value) -> serverSpan.setAttribute(key, value));
139+
parameterExtractor.setAttributes(request, serverSpan::setAttribute);
129140
}
130141

131142
/**
@@ -169,4 +180,66 @@ public boolean needsRescoping(Context currentContext, Context attachedContext) {
169180
private static boolean sameTrace(Span oneSpan, Span otherSpan) {
170181
return oneSpan.getSpanContext().getTraceId().equals(otherSpan.getSpanContext().getTraceId());
171182
}
183+
184+
protected static class ServletAsyncContext implements ImplicitContextKeyed {
185+
private static final ContextKey<ServletAsyncContext> CONTEXT_KEY =
186+
named("opentelemetry-servlet-async-context");
187+
188+
private boolean isAsyncListenerAttached;
189+
private Throwable throwable;
190+
private Object response;
191+
192+
public static Context init(Context context) {
193+
if (context.get(CONTEXT_KEY) != null) {
194+
return context;
195+
}
196+
return context.with(new ServletAsyncContext());
197+
}
198+
199+
@Nullable
200+
public static ServletAsyncContext get(@Nullable Context context) {
201+
return context != null ? context.get(CONTEXT_KEY) : null;
202+
}
203+
204+
public static boolean isAsyncListenerAttached(@Nullable Context context) {
205+
ServletAsyncContext servletAsyncContext = get(context);
206+
return servletAsyncContext != null && servletAsyncContext.isAsyncListenerAttached;
207+
}
208+
209+
public static void setAsyncListenerAttached(@Nullable Context context, boolean value) {
210+
ServletAsyncContext servletAsyncContext = get(context);
211+
if (servletAsyncContext != null) {
212+
servletAsyncContext.isAsyncListenerAttached = value;
213+
}
214+
}
215+
216+
public static Throwable getAsyncException(@Nullable Context context) {
217+
ServletAsyncContext servletAsyncContext = get(context);
218+
return servletAsyncContext != null ? servletAsyncContext.throwable : null;
219+
}
220+
221+
public static void recordAsyncException(@Nullable Context context, Throwable throwable) {
222+
ServletAsyncContext servletAsyncContext = get(context);
223+
if (servletAsyncContext != null) {
224+
servletAsyncContext.throwable = throwable;
225+
}
226+
}
227+
228+
public static Object getAsyncListenerResponse(@Nullable Context context) {
229+
ServletAsyncContext servletAsyncContext = get(context);
230+
return servletAsyncContext != null ? servletAsyncContext.response : null;
231+
}
232+
233+
public static void setAsyncListenerResponse(@Nullable Context context, Object response) {
234+
ServletAsyncContext servletAsyncContext = get(context);
235+
if (servletAsyncContext != null) {
236+
servletAsyncContext.response = response;
237+
}
238+
}
239+
240+
@Override
241+
public Context storeInContext(Context context) {
242+
return context.with(CONTEXT_KEY, this);
243+
}
244+
}
172245
}

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

+12-15
Original file line numberDiff line numberDiff line change
@@ -10,12 +10,6 @@
1010
import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
1111

1212
public class ServletHelper<REQUEST, RESPONSE> extends BaseServletHelper<REQUEST, RESPONSE> {
13-
private static final String ASYNC_LISTENER_ATTRIBUTE =
14-
ServletHelper.class.getName() + ".AsyncListener";
15-
private static final String ASYNC_LISTENER_RESPONSE_ATTRIBUTE =
16-
ServletHelper.class.getName() + ".AsyncListenerResponse";
17-
public static final String ASYNC_EXCEPTION_ATTRIBUTE =
18-
ServletHelper.class.getName() + ".AsyncException";
1913
public static final String CONTEXT_ATTRIBUTE = ServletHelper.class.getName() + ".Context";
2014

2115
public ServletHelper(
@@ -87,14 +81,14 @@ public boolean mustEndOnHandlerMethodExit(REQUEST request) {
8781
* is not possible to access response from async event in listeners.
8882
*/
8983
public void setAsyncListenerResponse(REQUEST request, RESPONSE response) {
90-
accessor.setRequestAttribute(request, ASYNC_LISTENER_RESPONSE_ATTRIBUTE, response);
84+
Context context = getServerContext(request);
85+
ServletAsyncContext.setAsyncListenerResponse(context, response);
9186
}
9287

88+
@SuppressWarnings("unchecked")
9389
public RESPONSE getAsyncListenerResponse(REQUEST request) {
94-
@SuppressWarnings("unchecked")
95-
RESPONSE response =
96-
(RESPONSE) accessor.getRequestAttribute(request, ASYNC_LISTENER_RESPONSE_ATTRIBUTE);
97-
return response;
90+
Context context = getServerContext(request);
91+
return (RESPONSE) ServletAsyncContext.getAsyncListenerResponse(context);
9892
}
9993

10094
public void attachAsyncListener(REQUEST request) {
@@ -113,23 +107,26 @@ private void attachAsyncListener(ServletRequestContext<REQUEST> requestContext)
113107
request,
114108
new AsyncRequestCompletionListener<>(this, instrumenter, requestContext, context),
115109
response);
116-
accessor.setRequestAttribute(request, ASYNC_LISTENER_ATTRIBUTE, true);
110+
ServletAsyncContext.setAsyncListenerAttached(context, true);
117111
}
118112
}
119113

120114
public boolean isAsyncListenerAttached(REQUEST request) {
121-
return accessor.getRequestAttribute(request, ASYNC_LISTENER_ATTRIBUTE) != null;
115+
Context context = getServerContext(request);
116+
return ServletAsyncContext.isAsyncListenerAttached(context);
122117
}
123118

124119
public Runnable wrapAsyncRunnable(REQUEST request, Runnable runnable) {
125120
return AsyncRunnableWrapper.wrap(this, request, runnable);
126121
}
127122

128123
public void recordAsyncException(REQUEST request, Throwable throwable) {
129-
accessor.setRequestAttribute(request, ASYNC_EXCEPTION_ATTRIBUTE, throwable);
124+
Context context = getServerContext(request);
125+
ServletAsyncContext.recordAsyncException(context, throwable);
130126
}
131127

132128
public Throwable getAsyncException(REQUEST request) {
133-
return (Throwable) accessor.getRequestAttribute(request, ASYNC_EXCEPTION_ATTRIBUTE);
129+
Context context = getServerContext(request);
130+
return ServletAsyncContext.getAsyncException(context);
134131
}
135132
}

0 commit comments

Comments
 (0)