Skip to content

Commit 9028d23

Browse files
lauritlaurit
committed
don't realy on context being present in request attribute
1 parent aae9aa8 commit 9028d23

File tree

21 files changed

+143
-168
lines changed

21 files changed

+143
-168
lines changed

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

+1-1
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@ public static void onEnter(
4545
scope = context.makeCurrent();
4646

4747
// Must be set here since Jetty handlers can use startAsync outside of servlet scope.
48-
helper().setAsyncListenerResponse(request, response);
48+
helper().setAsyncListenerResponse(context, response);
4949

5050
HttpServerResponseCustomizerHolder.getCustomizer()
5151
.customize(context, response, Jetty11ResponseMutator.INSTANCE);

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

+1-1
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
import io.opentelemetry.context.Context;
99
import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
1010
import io.opentelemetry.javaagent.bootstrap.servlet.AppServerBridge;
11-
import io.opentelemetry.javaagent.instrumentation.servlet.BaseServletHelper.ServletAsyncContext;
11+
import io.opentelemetry.javaagent.bootstrap.servlet.ServletAsyncContext;
1212
import javax.annotation.Nullable;
1313
import org.eclipse.jetty.server.HttpStream;
1414
import org.eclipse.jetty.server.Request;

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

+1-1
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@ public static void onEnter(
4545
scope = context.makeCurrent();
4646

4747
// Must be set here since Jetty handlers can use startAsync outside of servlet scope.
48-
helper().setAsyncListenerResponse(request, response);
48+
helper().setAsyncListenerResponse(context, response);
4949

5050
HttpServerResponseCustomizerHolder.getCustomizer()
5151
.customize(context, response, Jetty8ResponseMutator.INSTANCE);

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

+1-1
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,7 @@ public void end(
4141
}
4242

4343
ServletResponseContext<RESPONSE> responseContext = new ServletResponseContext<>(response);
44-
if (throwable != null || mustEndOnHandlerMethodExit(request)) {
44+
if (throwable != null || mustEndOnHandlerMethodExit(context)) {
4545
instrumenter.end(context, requestContext, responseContext, throwable);
4646
}
4747
}

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

+1-1
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,7 @@ public void end(
4040
}
4141

4242
ServletResponseContext<RESPONSE> responseContext = new ServletResponseContext<>(response);
43-
if (throwable != null || mustEndOnHandlerMethodExit(request)) {
43+
if (throwable != null || mustEndOnHandlerMethodExit(context)) {
4444
instrumenter.end(context, requestContext, responseContext, throwable);
4545
}
4646
}

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

+1-1
Original file line numberDiff line numberDiff line change
@@ -125,7 +125,7 @@ public static void onEnter() {
125125

126126
// Must be set here since Liberty RequestProcessors can use startAsync outside of servlet
127127
// scope.
128-
helper().setAsyncListenerResponse(requestInfo.getRequest(), requestInfo.getResponse());
128+
helper().setAsyncListenerResponse(context, requestInfo.getResponse());
129129

130130
HttpServerResponseCustomizerHolder.getCustomizer()
131131
.customize(context, requestInfo.getResponse(), Servlet3Accessor.INSTANCE);

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

+1-1
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,7 @@ public static void onEnter(
6060
requestContext = new ServletRequestContext<>(httpServletRequest, servletOrFilter);
6161
if (attachedContext == null && helper().shouldStart(currentContext, requestContext)) {
6262
context = helper().start(currentContext, requestContext);
63-
helper().setAsyncListenerResponse(httpServletRequest, (HttpServletResponse) response);
63+
helper().setAsyncListenerResponse(context, (HttpServletResponse) response);
6464

6565
contextToUpdate = context;
6666
} else if (attachedContext != null

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

+2-10
Original file line numberDiff line numberDiff line change
@@ -7,21 +7,13 @@
77

88
import static io.opentelemetry.javaagent.instrumentation.servlet.v3_0.Servlet3Singletons.helper;
99

10-
import javax.servlet.AsyncContext;
11-
import javax.servlet.ServletRequest;
12-
import javax.servlet.http.HttpServletRequest;
1310
import net.bytebuddy.asm.Advice;
1411

1512
@SuppressWarnings("unused")
1613
public class Servlet3AsyncContextStartAdvice {
1714

1815
@Advice.OnMethodEnter(suppress = Throwable.class)
19-
public static void start(
20-
@Advice.This AsyncContext asyncContext,
21-
@Advice.Argument(value = 0, readOnly = false) Runnable runnable) {
22-
ServletRequest request = asyncContext.getRequest();
23-
if (request instanceof HttpServletRequest) {
24-
runnable = helper().wrapAsyncRunnable((HttpServletRequest) request, runnable);
25-
}
16+
public static void start(@Advice.Argument(value = 0, readOnly = false) Runnable runnable) {
17+
runnable = helper().wrapAsyncRunnable(runnable);
2618
}
2719
}

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

+2-3
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
import static io.opentelemetry.javaagent.instrumentation.servlet.v3_0.Servlet3Singletons.helper;
99

1010
import io.opentelemetry.javaagent.bootstrap.CallDepth;
11+
import io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge;
1112
import javax.servlet.AsyncContext;
1213
import javax.servlet.ServletRequest;
1314
import javax.servlet.http.HttpServletRequest;
@@ -36,9 +37,7 @@ public static void startAsyncExit(
3637
if (servletRequest instanceof HttpServletRequest) {
3738
HttpServletRequest request = (HttpServletRequest) servletRequest;
3839

39-
if (!helper().isAsyncListenerAttached(request)) {
40-
helper().attachAsyncListener(request);
41-
}
40+
helper().attachAsyncListener(request, Java8BytecodeBridge.currentContext());
4241
}
4342
}
4443
}

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

+2-10
Original file line numberDiff line numberDiff line change
@@ -7,21 +7,13 @@
77

88
import static io.opentelemetry.javaagent.instrumentation.servlet.v5_0.Servlet5Singletons.helper;
99

10-
import jakarta.servlet.AsyncContext;
11-
import jakarta.servlet.ServletRequest;
12-
import jakarta.servlet.http.HttpServletRequest;
1310
import net.bytebuddy.asm.Advice;
1411

1512
@SuppressWarnings("unused")
1613
public class AsyncContextStartAdvice {
1714

1815
@Advice.OnMethodEnter(suppress = Throwable.class)
19-
public static void start(
20-
@Advice.This AsyncContext asyncContext,
21-
@Advice.Argument(value = 0, readOnly = false) Runnable runnable) {
22-
ServletRequest request = asyncContext.getRequest();
23-
if (request instanceof HttpServletRequest) {
24-
runnable = helper().wrapAsyncRunnable((HttpServletRequest) request, runnable);
25-
}
16+
public static void start(@Advice.Argument(value = 0, readOnly = false) Runnable runnable) {
17+
runnable = helper().wrapAsyncRunnable(runnable);
2618
}
2719
}

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

+2-3
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
import static io.opentelemetry.javaagent.instrumentation.servlet.v5_0.Servlet5Singletons.helper;
99

1010
import io.opentelemetry.javaagent.bootstrap.CallDepth;
11+
import io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge;
1112
import jakarta.servlet.AsyncContext;
1213
import jakarta.servlet.http.HttpServletRequest;
1314
import net.bytebuddy.asm.Advice;
@@ -34,9 +35,7 @@ public static void startAsyncExit(
3435
}
3536

3637
if (request != null) {
37-
if (!helper().isAsyncListenerAttached(request)) {
38-
helper().attachAsyncListener(request);
39-
}
38+
helper().attachAsyncListener(request, Java8BytecodeBridge.currentContext());
4039
}
4140
}
4241
}

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

+1-1
Original file line numberDiff line numberDiff line change
@@ -63,7 +63,7 @@ public static void onEnter(
6363
requestContext = new ServletRequestContext<>(httpServletRequest, servletOrFilter);
6464
if (attachedContext == null && helper().shouldStart(currentContext, requestContext)) {
6565
context = helper().start(currentContext, requestContext);
66-
helper().setAsyncListenerResponse(httpServletRequest, (HttpServletResponse) response);
66+
helper().setAsyncListenerResponse(context, (HttpServletResponse) response);
6767

6868
contextToUpdate = context;
6969
} else if (attachedContext != null

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

+3-1
Original file line numberDiff line numberDiff line change
@@ -125,7 +125,9 @@ public Builder captureServletAttributes() {
125125
* @return new context with AppServerBridge attached.
126126
*/
127127
public Context init(Context context) {
128-
return context.with(AppServerBridge.CONTEXT_KEY, new AppServerBridge(this));
128+
Context result = context.with(AppServerBridge.CONTEXT_KEY, new AppServerBridge(this));
129+
result = ServletAsyncContext.init(result);
130+
return result;
129131
}
130132
}
131133
}

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

+75
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,75 @@
1+
/*
2+
* Copyright The OpenTelemetry Authors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package io.opentelemetry.javaagent.bootstrap.servlet;
7+
8+
import static io.opentelemetry.context.ContextKey.named;
9+
10+
import io.opentelemetry.context.Context;
11+
import io.opentelemetry.context.ContextKey;
12+
import io.opentelemetry.context.ImplicitContextKeyed;
13+
import javax.annotation.Nullable;
14+
15+
public class ServletAsyncContext implements ImplicitContextKeyed {
16+
private static final ContextKey<ServletAsyncContext> CONTEXT_KEY =
17+
named("opentelemetry-servlet-async-context");
18+
19+
private boolean isAsyncListenerAttached;
20+
private Throwable throwable;
21+
private Object response;
22+
23+
public static Context init(Context context) {
24+
if (context.get(CONTEXT_KEY) != null) {
25+
return context;
26+
}
27+
return context.with(new ServletAsyncContext());
28+
}
29+
30+
@Nullable
31+
public static ServletAsyncContext get(@Nullable Context context) {
32+
return context != null ? context.get(CONTEXT_KEY) : null;
33+
}
34+
35+
public static boolean isAsyncListenerAttached(@Nullable Context context) {
36+
ServletAsyncContext servletAsyncContext = get(context);
37+
return servletAsyncContext != null && servletAsyncContext.isAsyncListenerAttached;
38+
}
39+
40+
public static void setAsyncListenerAttached(@Nullable Context context, boolean value) {
41+
ServletAsyncContext servletAsyncContext = get(context);
42+
if (servletAsyncContext != null) {
43+
servletAsyncContext.isAsyncListenerAttached = value;
44+
}
45+
}
46+
47+
public static Throwable getAsyncException(@Nullable Context context) {
48+
ServletAsyncContext servletAsyncContext = get(context);
49+
return servletAsyncContext != null ? servletAsyncContext.throwable : null;
50+
}
51+
52+
public static void recordAsyncException(@Nullable Context context, Throwable throwable) {
53+
ServletAsyncContext servletAsyncContext = get(context);
54+
if (servletAsyncContext != null) {
55+
servletAsyncContext.throwable = throwable;
56+
}
57+
}
58+
59+
public static Object getAsyncListenerResponse(@Nullable Context context) {
60+
ServletAsyncContext servletAsyncContext = get(context);
61+
return servletAsyncContext != null ? servletAsyncContext.response : null;
62+
}
63+
64+
public static void setAsyncListenerResponse(@Nullable Context context, Object response) {
65+
ServletAsyncContext servletAsyncContext = get(context);
66+
if (servletAsyncContext != null) {
67+
servletAsyncContext.response = response;
68+
}
69+
}
70+
71+
@Override
72+
public Context storeInContext(Context context) {
73+
return context.with(CONTEXT_KEY, this);
74+
}
75+
}

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

+3-3
Original file line numberDiff line numberDiff line change
@@ -33,18 +33,18 @@ public AsyncRequestCompletionListener(
3333
public void onComplete(RESPONSE response) {
3434
if (responseHandled.compareAndSet(false, true)) {
3535
ServletResponseContext<RESPONSE> responseContext = new ServletResponseContext<>(response);
36-
Throwable throwable = servletHelper.getAsyncException(requestContext.request());
36+
Throwable throwable = servletHelper.getAsyncException(context);
3737
instrumenter.end(context, requestContext, responseContext, throwable);
3838
}
3939
}
4040

4141
@Override
4242
public void onTimeout(long timeout) {
4343
if (responseHandled.compareAndSet(false, true)) {
44-
RESPONSE response = servletHelper.getAsyncListenerResponse(requestContext.request());
44+
RESPONSE response = servletHelper.getAsyncListenerResponse(context);
4545
ServletResponseContext<RESPONSE> responseContext = new ServletResponseContext<>(response);
4646
responseContext.setTimeout(timeout);
47-
Throwable throwable = servletHelper.getAsyncException(requestContext.request());
47+
Throwable throwable = servletHelper.getAsyncException(context);
4848
instrumenter.end(context, requestContext, responseContext, throwable);
4949
}
5050
}

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

+8-8
Original file line numberDiff line numberDiff line change
@@ -5,32 +5,32 @@
55

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

8+
import io.opentelemetry.context.Context;
9+
810
public class AsyncRunnableWrapper<REQUEST> implements Runnable {
911
private final ServletHelper<REQUEST, ?> helper;
10-
private final REQUEST request;
1112
private final Runnable runnable;
13+
private final Context context;
1214

13-
private AsyncRunnableWrapper(
14-
ServletHelper<REQUEST, ?> helper, REQUEST request, Runnable runnable) {
15+
private AsyncRunnableWrapper(ServletHelper<REQUEST, ?> helper, Runnable runnable) {
1516
this.helper = helper;
16-
this.request = request;
1717
this.runnable = runnable;
18+
this.context = Context.current();
1819
}
1920

20-
public static <REQUEST> Runnable wrap(
21-
ServletHelper<REQUEST, ?> helper, REQUEST request, Runnable runnable) {
21+
public static <REQUEST> Runnable wrap(ServletHelper<REQUEST, ?> helper, Runnable runnable) {
2222
if (runnable == null || runnable instanceof AsyncRunnableWrapper) {
2323
return runnable;
2424
}
25-
return new AsyncRunnableWrapper<>(helper, request, runnable);
25+
return new AsyncRunnableWrapper<>(helper, runnable);
2626
}
2727

2828
@Override
2929
public void run() {
3030
try {
3131
runnable.run();
3232
} catch (Throwable throwable) {
33-
helper.recordAsyncException(request, throwable);
33+
helper.recordAsyncException(context, throwable);
3434
throw throwable;
3535
}
3636
}

0 commit comments

Comments
 (0)