5
5
6
6
package io .opentelemetry .javaagent .instrumentation .servlet ;
7
7
8
+ import static io .opentelemetry .context .ContextKey .named ;
8
9
import static io .opentelemetry .instrumentation .api .semconv .http .HttpServerRouteSource .SERVER ;
9
10
import static io .opentelemetry .instrumentation .api .semconv .http .HttpServerRouteSource .SERVER_FILTER ;
10
11
11
12
import io .opentelemetry .api .trace .Span ;
12
13
import io .opentelemetry .api .trace .SpanContext ;
13
14
import io .opentelemetry .context .Context ;
15
+ import io .opentelemetry .context .ContextKey ;
16
+ import io .opentelemetry .context .ImplicitContextKeyed ;
14
17
import io .opentelemetry .instrumentation .api .instrumenter .Instrumenter ;
15
18
import io .opentelemetry .instrumentation .api .instrumenter .LocalRootSpan ;
16
19
import io .opentelemetry .instrumentation .api .semconv .http .HttpServerRoute ;
21
24
import io .opentelemetry .semconv .incubating .EnduserIncubatingAttributes ;
22
25
import java .security .Principal ;
23
26
import java .util .function .Function ;
27
+ import javax .annotation .Nullable ;
24
28
25
29
@ SuppressWarnings ("deprecation" ) // using deprecated semconv
26
30
public abstract class BaseServletHelper <REQUEST , RESPONSE > {
@@ -61,6 +65,7 @@ public Context start(Context parentContext, ServletRequestContext<REQUEST> reque
61
65
accessor .setRequestAttribute (request , "span_id" , spanContext .getSpanId ());
62
66
63
67
context = addServletContextPath (context , request );
68
+ context = addAsyncContext (context );
64
69
65
70
attachServerContext (context , request );
66
71
@@ -71,6 +76,10 @@ protected Context addServletContextPath(Context context, REQUEST request) {
71
76
return ServletContextPath .init (context , contextPathExtractor , request );
72
77
}
73
78
79
+ protected Context addAsyncContext (Context context ) {
80
+ return ServletAsyncContext .init (context );
81
+ }
82
+
74
83
public Context getServerContext (REQUEST request ) {
75
84
Object context = accessor .getRequestAttribute (request , ServletHelper .CONTEXT_ATTRIBUTE );
76
85
return context instanceof Context ? (Context ) context : null ;
@@ -87,6 +96,8 @@ public void recordException(Context context, Throwable throwable) {
87
96
public Context updateContext (
88
97
Context context , REQUEST request , MappingResolver mappingResolver , boolean servlet ) {
89
98
Context result = addServletContextPath (context , request );
99
+ result = addAsyncContext (result );
100
+
90
101
if (mappingResolver != null ) {
91
102
HttpServerRoute .update (
92
103
result , servlet ? SERVER : SERVER_FILTER , spanNameProvider , mappingResolver , request );
@@ -125,7 +136,7 @@ private void captureRequestParameters(Span serverSpan, REQUEST request) {
125
136
return ;
126
137
}
127
138
128
- parameterExtractor .setAttributes (request , ( key , value ) -> serverSpan . setAttribute ( key , value ) );
139
+ parameterExtractor .setAttributes (request , serverSpan :: setAttribute );
129
140
}
130
141
131
142
/**
@@ -169,4 +180,66 @@ public boolean needsRescoping(Context currentContext, Context attachedContext) {
169
180
private static boolean sameTrace (Span oneSpan , Span otherSpan ) {
170
181
return oneSpan .getSpanContext ().getTraceId ().equals (otherSpan .getSpanContext ().getTraceId ());
171
182
}
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
+ }
172
245
}
0 commit comments