Skip to content

Commit 8f5ca57

Browse files
felixbarnyfelixbarny
authored
Support async OkHttp calls (#526)
closes #425
1 parent b7d6b96 commit 8f5ca57

File tree

8 files changed

+510
-3
lines changed

8 files changed

+510
-3
lines changed

‎CHANGELOG.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
Now it supports all UTF-8 encoded plain-text content types.
2727
The option [`capture_body_content_types`](https://www.elastic.co/guide/en/apm/agent/java/master/config-http.html#config-capture-body-content-types)
2828
controls which `Content-Type`s should be captured.
29+
* Support async calls made by OkHttp client (`Call#enqueue`)
2930

3031
## Bug Fixes
3132

‎apm-agent-plugins/apm-okhttp-plugin/src/main/java/co/elastic/apm/agent/okhttp/OkHttp3ClientAsyncInstrumentation.java‎

Lines changed: 165 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,165 @@
1+
/*-
2+
* #%L
3+
* Elastic APM Java agent
4+
* %%
5+
* Copyright (C) 2018 - 2019 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.okhttp;
21+
22+
import co.elastic.apm.agent.bci.ElasticApmInstrumentation;
23+
import co.elastic.apm.agent.bci.HelperClassManager;
24+
import co.elastic.apm.agent.bci.VisibleForAdvice;
25+
import co.elastic.apm.agent.http.client.HttpClientHelper;
26+
import co.elastic.apm.agent.impl.ElasticApmTracer;
27+
import co.elastic.apm.agent.impl.transaction.Span;
28+
import co.elastic.apm.agent.impl.transaction.TraceContext;
29+
import co.elastic.apm.agent.impl.transaction.TraceContextHolder;
30+
import net.bytebuddy.asm.Advice;
31+
import net.bytebuddy.description.method.MethodDescription;
32+
import net.bytebuddy.description.type.TypeDescription;
33+
import net.bytebuddy.implementation.bytecode.assign.Assigner;
34+
import net.bytebuddy.matcher.ElementMatcher;
35+
import okhttp3.Call;
36+
import okhttp3.Callback;
37+
import okhttp3.Response;
38+
import org.slf4j.Logger;
39+
import org.slf4j.LoggerFactory;
40+
41+
import javax.annotation.Nullable;
42+
import java.io.IOException;
43+
import java.util.Arrays;
44+
import java.util.Collection;
45+
46+
import static net.bytebuddy.matcher.ElementMatchers.named;
47+
import static net.bytebuddy.matcher.ElementMatchers.returns;
48+
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
49+
import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
50+
51+
public class OkHttp3ClientAsyncInstrumentation extends ElasticApmInstrumentation {
52+
53+
@VisibleForAdvice
54+
public static final Logger logger = LoggerFactory.getLogger(OkHttp3ClientAsyncInstrumentation.class);
55+
56+
@Override
57+
public Class<?> getAdviceClass() {
58+
return OkHttpClient3ExecuteAdvice.class;
59+
}
60+
61+
@Nullable
62+
@VisibleForAdvice
63+
public static HelperClassManager<WrapperCreator<Callback>> callbackWrapperCreator;
64+
65+
66+
@Override
67+
public void init(ElasticApmTracer tracer) {
68+
callbackWrapperCreator = HelperClassManager.ForAnyClassLoader.of(tracer,
69+
OkHttp3ClientAsyncInstrumentation.class.getName() + "$CallbackWrapperCreator",
70+
OkHttp3ClientAsyncInstrumentation.class.getName() + "$CallbackWrapperCreator$CallbackWrapper");
71+
}
72+
73+
@VisibleForAdvice
74+
public static class OkHttpClient3ExecuteAdvice {
75+
76+
@Advice.OnMethodEnter(suppress = Throwable.class)
77+
private static void onBeforeEnqueue(@Advice.Origin Class<? extends Call> clazz,
78+
@Advice.FieldValue(value = "originalRequest", typing = Assigner.Typing.DYNAMIC, readOnly = false) @Nullable okhttp3.Request originalRequest,
79+
@Advice.Argument(value = 0, readOnly = false) @Nullable Callback callback,
80+
@Advice.Local("span") Span span) {
81+
if (tracer == null || tracer.getActive() == null || callbackWrapperCreator == null) {
82+
return;
83+
}
84+
85+
final WrapperCreator<Callback> wrapperCreator = callbackWrapperCreator.getForClassLoaderOfClass(clazz);
86+
if (originalRequest == null || callback == null || wrapperCreator == null) {
87+
return;
88+
}
89+
90+
final TraceContextHolder<?> parent = tracer.getActive();
91+
92+
okhttp3.Request request = originalRequest;
93+
span = HttpClientHelper.startHttpClientSpan(parent, request.method(), request.url().toString(), request.url().host());
94+
if (span != null) {
95+
span.activate().markLifecycleManagingThreadSwitchExpected();
96+
originalRequest = originalRequest.newBuilder().addHeader(TraceContext.TRACE_PARENT_HEADER, span.getTraceContext().getOutgoingTraceParentHeader().toString()).build();
97+
callback = wrapperCreator.wrap(callback, span);
98+
}
99+
}
100+
101+
@Advice.OnMethodExit(suppress = Throwable.class)
102+
private static void onAfterEnqueue(@Advice.Local("span") @Nullable Span span) {
103+
if (span != null) {
104+
span.deactivate();
105+
}
106+
}
107+
}
108+
109+
public static class CallbackWrapperCreator implements WrapperCreator<Callback> {
110+
111+
@Override
112+
public Callback wrap(final Callback delegate, Span span) {
113+
return new CallbackWrapper(span, delegate);
114+
}
115+
116+
private static class CallbackWrapper implements Callback {
117+
private final Span span;
118+
private final Callback delegate;
119+
120+
CallbackWrapper(Span span, Callback delegate) {
121+
this.span = span;
122+
this.delegate = delegate;
123+
}
124+
125+
@Override
126+
public void onFailure(Call call, IOException e) {
127+
try {
128+
span.captureException(e).end();
129+
} catch (Throwable t) {
130+
logger.error(t.getMessage(), t);
131+
} finally {
132+
delegate.onFailure(call, e);
133+
}
134+
}
135+
136+
@Override
137+
public void onResponse(Call call, Response response) throws IOException {
138+
try {
139+
span.getContext().getHttp().withStatusCode(response.code());
140+
span.end();
141+
} catch (Throwable t) {
142+
logger.error(t.getMessage(), t);
143+
} finally {
144+
delegate.onResponse(call, response);
145+
}
146+
}
147+
}
148+
}
149+
150+
@Override
151+
public ElementMatcher<? super TypeDescription> getTypeMatcher() {
152+
return named("okhttp3.RealCall");
153+
}
154+
155+
@Override
156+
public ElementMatcher<? super MethodDescription> getMethodMatcher() {
157+
return named("enqueue").and(takesArguments(1)).and(takesArgument(0, named("okhttp3.Callback"))).and(returns(void.class));
158+
}
159+
160+
@Override
161+
public Collection<String> getInstrumentationGroupNames() {
162+
return Arrays.asList("http-client", "okhttp");
163+
}
164+
165+
}

‎apm-agent-plugins/apm-okhttp-plugin/src/main/java/co/elastic/apm/agent/okhttp/OkHttpClientAsyncInstrumentation.java‎

Lines changed: 164 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,164 @@
1+
/*-
2+
* #%L
3+
* Elastic APM Java agent
4+
* %%
5+
* Copyright (C) 2018 - 2019 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.okhttp;
21+
22+
import co.elastic.apm.agent.bci.ElasticApmInstrumentation;
23+
import co.elastic.apm.agent.bci.HelperClassManager;
24+
import co.elastic.apm.agent.bci.VisibleForAdvice;
25+
import co.elastic.apm.agent.http.client.HttpClientHelper;
26+
import co.elastic.apm.agent.impl.ElasticApmTracer;
27+
import co.elastic.apm.agent.impl.transaction.Span;
28+
import co.elastic.apm.agent.impl.transaction.TraceContext;
29+
import co.elastic.apm.agent.impl.transaction.TraceContextHolder;
30+
import com.squareup.okhttp.Call;
31+
import com.squareup.okhttp.Callback;
32+
import com.squareup.okhttp.Request;
33+
import com.squareup.okhttp.Response;
34+
import net.bytebuddy.asm.Advice;
35+
import net.bytebuddy.description.method.MethodDescription;
36+
import net.bytebuddy.description.type.TypeDescription;
37+
import net.bytebuddy.implementation.bytecode.assign.Assigner;
38+
import net.bytebuddy.matcher.ElementMatcher;
39+
import org.slf4j.Logger;
40+
import org.slf4j.LoggerFactory;
41+
42+
import javax.annotation.Nullable;
43+
import java.io.IOException;
44+
import java.util.Arrays;
45+
import java.util.Collection;
46+
47+
import static net.bytebuddy.matcher.ElementMatchers.named;
48+
import static net.bytebuddy.matcher.ElementMatchers.returns;
49+
50+
public class OkHttpClientAsyncInstrumentation extends ElasticApmInstrumentation {
51+
52+
@VisibleForAdvice
53+
public static final Logger logger = LoggerFactory.getLogger(OkHttpClientAsyncInstrumentation.class);
54+
55+
@Override
56+
public Class<?> getAdviceClass() {
57+
return OkHttpClient3ExecuteAdvice.class;
58+
}
59+
60+
@Nullable
61+
@VisibleForAdvice
62+
public static HelperClassManager<WrapperCreator<Callback>> callbackWrapperCreator;
63+
64+
65+
@Override
66+
public void init(ElasticApmTracer tracer) {
67+
callbackWrapperCreator = HelperClassManager.ForAnyClassLoader.of(tracer,
68+
OkHttpClientAsyncInstrumentation.class.getName() + "$CallbackWrapperCreator",
69+
OkHttpClientAsyncInstrumentation.class.getName() + "$CallbackWrapperCreator$CallbackWrapper");
70+
}
71+
72+
@VisibleForAdvice
73+
public static class OkHttpClient3ExecuteAdvice {
74+
75+
@Advice.OnMethodEnter(suppress = Throwable.class)
76+
private static void onBeforeEnqueue(@Advice.Origin Class<? extends Call> clazz,
77+
@Advice.FieldValue(value = "originalRequest", typing = Assigner.Typing.DYNAMIC, readOnly = false) @Nullable Request originalRequest,
78+
@Advice.Argument(value = 0, readOnly = false) @Nullable Callback callback,
79+
@Advice.Local("span") Span span) {
80+
if (tracer == null || tracer.getActive() == null || callbackWrapperCreator == null) {
81+
return;
82+
}
83+
84+
final WrapperCreator<Callback> wrapperCreator = callbackWrapperCreator.getForClassLoaderOfClass(clazz);
85+
if (originalRequest == null || callback == null || wrapperCreator == null) {
86+
return;
87+
}
88+
89+
final TraceContextHolder<?> parent = tracer.getActive();
90+
91+
Request request = originalRequest;
92+
span = HttpClientHelper.startHttpClientSpan(parent, request.method(), request.url().toString(), request.url().getHost());
93+
if (span != null) {
94+
span.activate().markLifecycleManagingThreadSwitchExpected();
95+
originalRequest = originalRequest.newBuilder().addHeader(TraceContext.TRACE_PARENT_HEADER, span.getTraceContext().getOutgoingTraceParentHeader().toString()).build();
96+
callback = wrapperCreator.wrap(callback, span);
97+
}
98+
}
99+
100+
@Advice.OnMethodExit(suppress = Throwable.class)
101+
private static void onAfterEnqueue(@Advice.Local("span") @Nullable Span span) {
102+
if (span != null) {
103+
span.deactivate();
104+
}
105+
}
106+
}
107+
108+
public static class CallbackWrapperCreator implements WrapperCreator<Callback> {
109+
110+
@Override
111+
public Callback wrap(final Callback delegate, Span span) {
112+
return new CallbackWrapper(span, delegate);
113+
}
114+
115+
private static class CallbackWrapper implements Callback {
116+
private final Span span;
117+
private final Callback delegate;
118+
119+
CallbackWrapper(Span span, Callback delegate) {
120+
this.span = span;
121+
this.delegate = delegate;
122+
}
123+
124+
@Override
125+
public void onFailure(Request req, IOException e) {
126+
try {
127+
span.captureException(e).end();
128+
} catch (Throwable t) {
129+
logger.error(t.getMessage(), t);
130+
} finally {
131+
delegate.onFailure(req, e);
132+
}
133+
}
134+
135+
@Override
136+
public void onResponse(Response response) throws IOException {
137+
try {
138+
span.getContext().getHttp().withStatusCode(response.code());
139+
span.end();
140+
} catch (Throwable t) {
141+
logger.error(t.getMessage(), t);
142+
} finally {
143+
delegate.onResponse(response);
144+
}
145+
}
146+
}
147+
}
148+
149+
@Override
150+
public ElementMatcher<? super TypeDescription> getTypeMatcher() {
151+
return named("com.squareup.okhttp.Call");
152+
}
153+
154+
@Override
155+
public ElementMatcher<? super MethodDescription> getMethodMatcher() {
156+
return named("enqueue").and(returns(void.class));
157+
}
158+
159+
@Override
160+
public Collection<String> getInstrumentationGroupNames() {
161+
return Arrays.asList("http-client", "okhttp");
162+
}
163+
164+
}

‎apm-agent-plugins/apm-okhttp-plugin/src/main/java/co/elastic/apm/agent/okhttp/WrapperCreator.java‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
/*-
2+
* #%L
3+
* Elastic APM Java agent
4+
* %%
5+
* Copyright (C) 2018 - 2019 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.okhttp;
21+
22+
import co.elastic.apm.agent.impl.transaction.Span;
23+
24+
/**
25+
* Used to create a wrapper for a callback or listener
26+
*
27+
* @param <T> the type of the wrapper to create
28+
*/
29+
public interface WrapperCreator<T> {
30+
31+
/**
32+
* Wraps a callback or listener.
33+
* <p>
34+
* The implementation is supposed to create the actual wrapper which manages the lifecycle of the provided {@link Span}.
35+
* </p>
36+
*
37+
* @param delegate the actual callback which should be wrapped
38+
* @param span the currently active span
39+
* @return the wrapped callback
40+
*/
41+
T wrap(T delegate, Span span);
42+
}

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,2 +1,4 @@
11
co.elastic.apm.agent.okhttp.OkHttpClientInstrumentation
22
co.elastic.apm.agent.okhttp.OkHttp3ClientInstrumentation
3+
co.elastic.apm.agent.okhttp.OkHttpClientAsyncInstrumentation
4+
co.elastic.apm.agent.okhttp.OkHttp3ClientAsyncInstrumentation

0 commit comments

Comments
 (0)