Skip to content

Commit a0754fd

Browse files
eyalkoreneyalkoren
authored
Wrap Runnable in AsyncContext.start (#388)
1 parent 915b929 commit a0754fd

File tree

15 files changed

+348
-99
lines changed

15 files changed

+348
-99
lines changed

‎apm-agent-api/src/main/java/co/elastic/apm/api/AbstractSpanImpl.java

+2-2
Original file line numberDiff line numberDiff line change
@@ -7,9 +7,9 @@
77
* Licensed under the Apache License, Version 2.0 (the "License");
88
* you may not use this file except in compliance with the License.
99
* You may obtain a copy of the License at
10-
*
10+
*
1111
* http://www.apache.org/licenses/LICENSE-2.0
12-
*
12+
*
1313
* Unless required by applicable law or agreed to in writing, software
1414
* distributed under the License is distributed on an "AS IS" BASIS,
1515
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

‎apm-agent-core/src/main/java/co/elastic/apm/agent/impl/ElasticApmTracer.java

+16
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,7 @@ public class ElasticApmTracer {
6363
private final ObjectPool<Transaction> transactionPool;
6464
private final ObjectPool<Span> spanPool;
6565
private final ObjectPool<ErrorCapture> errorPool;
66+
private final ObjectPool<InScopeRunnableWrapper> runnableWrapperObjectPool;
6667
private final Reporter reporter;
6768
// Maintains a stack of all the activated spans
6869
// This way its easy to retrieve the bottom of the stack (the transaction)
@@ -107,6 +108,13 @@ public ErrorCapture createInstance() {
107108
return new ErrorCapture(ElasticApmTracer.this);
108109
}
109110
});
111+
runnableWrapperObjectPool = QueueBasedObjectPool.ofRecyclable(AtomicQueueFactory.<InScopeRunnableWrapper>newQueue(createBoundedMpmc(maxPooledElements)), false,
112+
new Allocator<InScopeRunnableWrapper>() {
113+
@Override
114+
public InScopeRunnableWrapper createInstance() {
115+
return new InScopeRunnableWrapper(ElasticApmTracer.this);
116+
}
117+
});
110118
sampler = ProbabilitySampler.of(coreConfiguration.getSampleRate().get());
111119
coreConfiguration.getSampleRate().addChangeListener(new ConfigurationOption.ChangeListener<Double>() {
112120
@Override
@@ -296,6 +304,14 @@ public void recycle(ErrorCapture error) {
296304
errorPool.recycle(error);
297305
}
298306

307+
public Runnable wrapRunnable(Runnable delegate, AbstractSpan<?> span) {
308+
return runnableWrapperObjectPool.createInstance().wrap(delegate, span);
309+
}
310+
311+
public void recycle(InScopeRunnableWrapper wrapper) {
312+
runnableWrapperObjectPool.recycle(wrapper);
313+
}
314+
299315
/**
300316
* Called when the container shuts down.
301317
* Cleans up thread pools and other resources.

‎apm-agent-core/src/main/java/co/elastic/apm/agent/impl/InScopeRunnableWrapper.java

+90
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,90 @@
1+
/*-
2+
* #%L
3+
* Elastic APM Java agent
4+
* %%
5+
* Copyright (C) 2018 Elastic and contributors
6+
* %%
7+
* Licensed under the Apache License, Version 2.0 (the "License");
8+
* you may not use this file except in compliance with the License.
9+
* You may obtain a copy of the License at
10+
*
11+
* http://www.apache.org/licenses/LICENSE-2.0
12+
*
13+
* Unless required by applicable law or agreed to in writing, software
14+
* distributed under the License is distributed on an "AS IS" BASIS,
15+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
16+
* See the License for the specific language governing permissions and
17+
* limitations under the License.
18+
* #L%
19+
*/
20+
package co.elastic.apm.agent.impl;
21+
22+
23+
import co.elastic.apm.agent.bci.VisibleForAdvice;
24+
import co.elastic.apm.agent.impl.transaction.AbstractSpan;
25+
import co.elastic.apm.agent.objectpool.Recyclable;
26+
import org.slf4j.Logger;
27+
import org.slf4j.LoggerFactory;
28+
29+
import javax.annotation.Nullable;
30+
31+
@VisibleForAdvice
32+
public class InScopeRunnableWrapper implements Runnable, Recyclable {
33+
private final Logger logger = LoggerFactory.getLogger(InScopeRunnableWrapper.class);
34+
35+
@Nullable
36+
private Runnable delegate;
37+
@Nullable
38+
private AbstractSpan<?> span;
39+
40+
private final ElasticApmTracer tracer;
41+
42+
InScopeRunnableWrapper(ElasticApmTracer tracer) {
43+
this.tracer = tracer;
44+
}
45+
46+
public InScopeRunnableWrapper wrap(Runnable delegate, AbstractSpan<?> span) {
47+
this.delegate = delegate;
48+
this.span = span;
49+
return this;
50+
}
51+
52+
@Override
53+
public void run() {
54+
if (span != null) {
55+
try {
56+
span.activate();
57+
} catch (Throwable throwable) {
58+
try {
59+
logger.warn("Failed to activate span");
60+
} catch (Throwable t) {
61+
// do nothing, just never fail
62+
}
63+
}
64+
}
65+
66+
try {
67+
//noinspection ConstantConditions
68+
delegate.run();
69+
} finally {
70+
try {
71+
if (span != null) {
72+
span.deactivate();
73+
}
74+
tracer.recycle(this);
75+
} catch (Throwable throwable) {
76+
try {
77+
logger.warn("Failed to deactivate span or recycle");
78+
} catch (Throwable t) {
79+
// do nothing, just never fail
80+
}
81+
}
82+
}
83+
}
84+
85+
@Override
86+
public void resetState() {
87+
delegate = null;
88+
span = null;
89+
}
90+
}

‎apm-agent-plugins/apm-servlet-plugin/src/main/java/co/elastic/apm/agent/servlet/AsyncInstrumentation.java

+92-50
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
import co.elastic.apm.agent.bci.HelperClassManager;
2424
import co.elastic.apm.agent.bci.VisibleForAdvice;
2525
import co.elastic.apm.agent.impl.ElasticApmTracer;
26+
import co.elastic.apm.agent.impl.transaction.Transaction;
2627
import net.bytebuddy.asm.Advice;
2728
import net.bytebuddy.description.NamedElement;
2829
import net.bytebuddy.description.method.MethodDescription;
@@ -52,83 +53,124 @@
5253
* {@code javax.servlet}, as these are inlined into the matching methods.
5354
* The agent itself does not have access to the Servlet API classes, as they are loaded by a child class loader.
5455
* See https://github.com/raphw/byte-buddy/issues/465 for more information.
55-
* However, the helper class {@link StartAsyncAdviceHelper} has access to the Servlet API,
56+
* However, the helper class {@link AsyncContextAdviceHelper} has access to the Servlet API,
5657
* as it is loaded by the child classloader of {@link AsyncContext}
57-
* (see {@link StartAsyncAdvice#onExitStartAsync(HttpServletRequest, AsyncContext)}).
58+
* (see {@link StartAsyncInstrumentation.StartAsyncAdvice#onExitStartAsync(HttpServletRequest, AsyncContext)}
59+
* and {@link AsyncContextInstrumentation.AsyncContextStartAdvice#onEnterAsyncContextStart(Runnable)}).
5860
*/
59-
public class AsyncInstrumentation extends ElasticApmInstrumentation {
61+
public abstract class AsyncInstrumentation extends ElasticApmInstrumentation {
6062

6163
private static final String SERVLET_API_ASYNC_GROUP_NAME = "servlet-api-async";
6264
@Nullable
6365
@VisibleForAdvice
6466
// referring to AsyncContext is legal because of type erasure
65-
public static HelperClassManager<StartAsyncAdviceHelper<AsyncContext>> asyncHelper;
67+
public static HelperClassManager<AsyncContextAdviceHelper<AsyncContext>> asyncHelper;
6668

6769
@Override
6870
public void init(ElasticApmTracer tracer) {
6971
asyncHelper = HelperClassManager.ForSingleClassLoader.of(tracer,
70-
"co.elastic.apm.agent.servlet.helper.StartAsyncAdviceHelperImpl",
72+
"co.elastic.apm.agent.servlet.helper.AsyncContextAdviceHelperImpl",
7173
"co.elastic.apm.agent.servlet.helper.ApmAsyncListener");
7274
}
7375

7476
@Override
75-
public ElementMatcher<? super NamedElement> getTypeMatcherPreFilter() {
76-
return nameContains("Request");
77+
public Collection<String> getInstrumentationGroupNames() {
78+
return Arrays.asList(ServletInstrumentation.SERVLET_API, SERVLET_API_ASYNC_GROUP_NAME);
7779
}
7880

79-
@Override
80-
public ElementMatcher<? super TypeDescription> getTypeMatcher() {
81-
return not(isInterface())
82-
.and(hasSuperType(named("javax.servlet.http.HttpServletRequest")));
81+
public interface AsyncContextAdviceHelper<T> {
82+
void onExitStartAsync(T asyncContext);
8383
}
8484

85-
/**
86-
* Matches
87-
* <ul>
88-
* <li>{@link HttpServletRequest#startAsync()}</li>
89-
* <li>{@link HttpServletRequest#startAsync(ServletRequest, ServletResponse)}</li>
90-
* </ul>
91-
*
92-
* @return
93-
*/
94-
@Override
95-
public ElementMatcher<? super MethodDescription> getMethodMatcher() {
96-
return named("startAsync")
97-
.and(returns(named("javax.servlet.AsyncContext")))
98-
.and(takesArguments(0)
99-
.or(
100-
takesArgument(0, named("javax.servlet.ServletRequest"))
101-
.and(takesArgument(1, named("javax.servlet.ServletResponse")))
102-
)
103-
)
104-
.and(isPublic());
105-
}
85+
public static class StartAsyncInstrumentation extends AsyncInstrumentation {
86+
@Override
87+
public ElementMatcher<? super NamedElement> getTypeMatcherPreFilter() {
88+
return nameContains("Request");
89+
}
10690

107-
@Override
108-
public Collection<String> getInstrumentationGroupNames() {
109-
return Arrays.asList(ServletInstrumentation.SERVLET_API, SERVLET_API_ASYNC_GROUP_NAME);
110-
}
91+
@Override
92+
public ElementMatcher<? super TypeDescription> getTypeMatcher() {
93+
return not(isInterface())
94+
.and(hasSuperType(named("javax.servlet.http.HttpServletRequest")));
95+
}
11196

112-
@Override
113-
public Class<?> getAdviceClass() {
114-
return StartAsyncAdvice.class;
115-
}
97+
/**
98+
* Matches
99+
* <ul>
100+
* <li>{@link HttpServletRequest#startAsync()}</li>
101+
* <li>{@link HttpServletRequest#startAsync(ServletRequest, ServletResponse)}</li>
102+
* </ul>
103+
*
104+
* @return
105+
*/
106+
@Override
107+
public ElementMatcher<? super MethodDescription> getMethodMatcher() {
108+
return isPublic()
109+
.and(named("startAsync"))
110+
.and(returns(named("javax.servlet.AsyncContext")))
111+
.and(takesArguments(0)
112+
.or(
113+
takesArgument(0, named("javax.servlet.ServletRequest"))
114+
.and(takesArgument(1, named("javax.servlet.ServletResponse")))
115+
)
116+
);
117+
}
116118

117-
public interface StartAsyncAdviceHelper<T> {
118-
void onExitStartAsync(T asyncContext);
119-
}
119+
@Override
120+
public Class<?> getAdviceClass() {
121+
return StartAsyncAdvice.class;
122+
}
120123

121-
@VisibleForAdvice
122-
public static class StartAsyncAdvice {
124+
@VisibleForAdvice
125+
public static class StartAsyncAdvice {
123126

124-
@Advice.OnMethodExit(suppress = Throwable.class)
125-
private static void onExitStartAsync(@Advice.This HttpServletRequest request, @Advice.Return AsyncContext asyncContext) {
126-
if (tracer != null) {
127-
if (asyncHelper != null) {
128-
asyncHelper.getForClassLoaderOfClass(AsyncContext.class).onExitStartAsync(asyncContext);
127+
@Advice.OnMethodExit(suppress = Throwable.class)
128+
private static void onExitStartAsync(@Advice.This HttpServletRequest request, @Advice.Return AsyncContext asyncContext) {
129+
if (tracer != null) {
130+
if (asyncHelper != null) {
131+
asyncHelper.getForClassLoaderOfClass(AsyncContext.class).onExitStartAsync(asyncContext);
132+
}
129133
}
130134
}
131135
}
132136
}
133137

138+
public static class AsyncContextInstrumentation extends AsyncInstrumentation {
139+
@Override
140+
public ElementMatcher<? super NamedElement> getTypeMatcherPreFilter() {
141+
return nameContains("AsyncContext");
142+
}
143+
144+
@Override
145+
public ElementMatcher<? super TypeDescription> getTypeMatcher() {
146+
return not(isInterface())
147+
.and(hasSuperType(named("javax.servlet.AsyncContext")));
148+
}
149+
150+
@Override
151+
public ElementMatcher<? super MethodDescription> getMethodMatcher() {
152+
return isPublic()
153+
.and(named("start"))
154+
.and(takesArguments(Runnable.class));
155+
}
156+
157+
@Override
158+
public Class<?> getAdviceClass() {
159+
return AsyncContextStartAdvice.class;
160+
}
161+
162+
@VisibleForAdvice
163+
public static class AsyncContextStartAdvice {
164+
165+
@Advice.OnMethodEnter(suppress = Throwable.class)
166+
private static void onEnterAsyncContextStart(@Advice.Argument(value = 0, readOnly = false) @Nullable Runnable runnable) {
167+
if (tracer != null) {
168+
final Transaction transaction = tracer.currentTransaction();
169+
if (transaction != null && runnable != null && asyncHelper != null) {
170+
runnable = tracer.wrapRunnable(runnable, transaction);
171+
}
172+
}
173+
}
174+
}
175+
}
134176
}

‎apm-agent-plugins/apm-servlet-plugin/src/main/java/co/elastic/apm/agent/servlet/helper/StartAsyncAdviceHelperImpl.java ‎apm-agent-plugins/apm-servlet-plugin/src/main/java/co/elastic/apm/agent/servlet/helper/AsyncContextAdviceHelperImpl.java

+2-2
Original file line numberDiff line numberDiff line change
@@ -31,14 +31,14 @@
3131
import static co.elastic.apm.agent.servlet.ServletTransactionHelper.ASYNC_ATTRIBUTE;
3232
import static co.elastic.apm.agent.servlet.ServletTransactionHelper.TRANSACTION_ATTRIBUTE;
3333

34-
public class StartAsyncAdviceHelperImpl implements AsyncInstrumentation.StartAsyncAdviceHelper<AsyncContext> {
34+
public class AsyncContextAdviceHelperImpl implements AsyncInstrumentation.AsyncContextAdviceHelper<AsyncContext> {
3535

3636
private static final String ASYNC_LISTENER_ADDED = ServletApiAdvice.class.getName() + ".asyncListenerAdded";
3737

3838
private final ServletTransactionHelper servletTransactionHelper;
3939
private final ElasticApmTracer tracer;
4040

41-
public StartAsyncAdviceHelperImpl(ElasticApmTracer tracer) {
41+
public AsyncContextAdviceHelperImpl(ElasticApmTracer tracer) {
4242
this.tracer = tracer;
4343
servletTransactionHelper = new ServletTransactionHelper(tracer);
4444
}

‎apm-agent-plugins/apm-servlet-plugin/src/main/java/co/elastic/apm/agent/servlet/helper/package-info.java

+3-3
Original file line numberDiff line numberDiff line change
@@ -48,13 +48,13 @@
4848
* </p>
4949
*
5050
* <pre>
51-
* System/Bootstrap CL (Agent) provides interface StartAsyncAdviceHelper
51+
* System/Bootstrap CL (Agent) provides interface AsyncContextAdviceHelper
5252
* |
5353
* v
54-
* App CL (Servlet API) uses StartAsyncAdviceHelperImpl from Helper CL
54+
* App CL (Servlet API) uses AsyncContextAdviceHelperImpl from Helper CL
5555
* / \
5656
* v v
57-
* WebApp CL (user code/libs) Helper CL implements StartAsyncAdviceHelperImpl
57+
* WebApp CL (user code/libs) Helper CL implements AsyncContextAdviceHelperImpl
5858
* </pre>
5959
*/
6060
@NonnullApi

‎apm-agent-plugins/apm-servlet-plugin/src/main/resources/META-INF/services/co.elastic.apm.agent.bci.ElasticApmInstrumentation

+2-1
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
11
co.elastic.apm.agent.servlet.ServletInstrumentation
22
co.elastic.apm.agent.servlet.FilterChainInstrumentation
3-
co.elastic.apm.agent.servlet.AsyncInstrumentation
3+
co.elastic.apm.agent.servlet.AsyncInstrumentation$StartAsyncInstrumentation
4+
co.elastic.apm.agent.servlet.AsyncInstrumentation$AsyncContextInstrumentation

0 commit comments

Comments
 (0)