Wrapping callbacks in trace representation
without this change when user adds a callback then it's not wrapped in a trace representation, thus when executed in a separate thread, it's then missing all the tracing information
with this change we're wrapping the future in our representation that passes the tracing information around
fixes #546
This commit is contained in:
@@ -18,6 +18,9 @@ package org.springframework.cloud.sleuth.instrument.web.client;
|
||||
|
||||
import java.lang.invoke.MethodHandles;
|
||||
import java.net.URI;
|
||||
import java.util.concurrent.ExecutionException;
|
||||
import java.util.concurrent.TimeUnit;
|
||||
import java.util.concurrent.TimeoutException;
|
||||
|
||||
import org.apache.commons.logging.Log;
|
||||
import org.apache.commons.logging.LogFactory;
|
||||
@@ -28,8 +31,10 @@ import org.springframework.core.task.AsyncListenableTaskExecutor;
|
||||
import org.springframework.http.HttpMethod;
|
||||
import org.springframework.http.client.AsyncClientHttpRequestFactory;
|
||||
import org.springframework.http.client.ClientHttpRequestFactory;
|
||||
import org.springframework.util.concurrent.FailureCallback;
|
||||
import org.springframework.util.concurrent.ListenableFuture;
|
||||
import org.springframework.util.concurrent.ListenableFutureCallback;
|
||||
import org.springframework.util.concurrent.SuccessCallback;
|
||||
import org.springframework.web.client.AsyncRequestCallback;
|
||||
import org.springframework.web.client.AsyncRestTemplate;
|
||||
import org.springframework.web.client.ResponseExtractor;
|
||||
@@ -80,14 +85,144 @@ public class TraceAsyncRestTemplate extends AsyncRestTemplate {
|
||||
protected <T> ListenableFuture<T> doExecute(URI url, HttpMethod method,
|
||||
AsyncRequestCallback requestCallback, ResponseExtractor<T> responseExtractor)
|
||||
throws RestClientException {
|
||||
ListenableFuture<T> future = super.doExecute(url, method, requestCallback, responseExtractor);
|
||||
Span span = this.tracer.getCurrentSpan();
|
||||
final ListenableFuture<T> future = super.doExecute(url, method, requestCallback, responseExtractor);
|
||||
final Span span = this.tracer.getCurrentSpan();
|
||||
future.addCallback(new TraceListenableFutureCallback<>(this.tracer, span));
|
||||
// potential race can happen here
|
||||
if (span != null && span.equals(this.tracer.getCurrentSpan())) {
|
||||
this.tracer.detach(span);
|
||||
}
|
||||
return future;
|
||||
return new ListenableFuture<T>() {
|
||||
|
||||
@Override public boolean cancel(boolean mayInterruptIfRunning) {
|
||||
return future.cancel(mayInterruptIfRunning);
|
||||
}
|
||||
|
||||
@Override public boolean isCancelled() {
|
||||
return future.isCancelled();
|
||||
}
|
||||
|
||||
@Override public boolean isDone() {
|
||||
return future.isDone();
|
||||
}
|
||||
|
||||
@Override public T get() throws InterruptedException, ExecutionException {
|
||||
return future.get();
|
||||
}
|
||||
|
||||
@Override public T get(long timeout, TimeUnit unit)
|
||||
throws InterruptedException, ExecutionException, TimeoutException {
|
||||
return future.get(timeout, unit);
|
||||
}
|
||||
|
||||
@Override
|
||||
public void addCallback(ListenableFutureCallback<? super T> callback) {
|
||||
future.addCallback(new TraceListenableFutureCallbackWrapper<>(TraceAsyncRestTemplate.this.tracer, span, callback));
|
||||
}
|
||||
|
||||
@Override public void addCallback(SuccessCallback<? super T> successCallback,
|
||||
FailureCallback failureCallback) {
|
||||
future.addCallback(
|
||||
new TraceSuccessCallback<>(TraceAsyncRestTemplate.this.tracer, span, successCallback),
|
||||
new TraceFailureCallback(TraceAsyncRestTemplate.this.tracer, span, failureCallback));
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
private static class TraceSuccessCallback<T> implements SuccessCallback<T> {
|
||||
|
||||
private static final Log log = LogFactory.getLog(MethodHandles.lookup().lookupClass());
|
||||
|
||||
private final Tracer tracer;
|
||||
private final Span parent;
|
||||
private final SuccessCallback<T> delegate;
|
||||
|
||||
private TraceSuccessCallback(Tracer tracer, Span parent,
|
||||
SuccessCallback<T> delegate) {
|
||||
this.tracer = tracer;
|
||||
this.parent = parent;
|
||||
this.delegate = delegate;
|
||||
}
|
||||
|
||||
@Override public void onSuccess(T result) {
|
||||
continueSpan();
|
||||
if (log.isDebugEnabled()) {
|
||||
log.debug("Calling on success of the delegate");
|
||||
}
|
||||
this.delegate.onSuccess(result);
|
||||
finish();
|
||||
}
|
||||
|
||||
private void continueSpan() {
|
||||
this.tracer.continueSpan(this.parent);
|
||||
}
|
||||
|
||||
private void finish() {
|
||||
this.tracer.detach(currentSpan());
|
||||
}
|
||||
|
||||
private Span currentSpan() {
|
||||
return this.tracer.getCurrentSpan();
|
||||
}
|
||||
}
|
||||
|
||||
private static class TraceFailureCallback implements FailureCallback {
|
||||
|
||||
private static final Log log = LogFactory.getLog(MethodHandles.lookup().lookupClass());
|
||||
|
||||
private final Tracer tracer;
|
||||
private final Span parent;
|
||||
private final FailureCallback delegate;
|
||||
|
||||
private TraceFailureCallback(Tracer tracer, Span parent,
|
||||
FailureCallback delegate) {
|
||||
this.tracer = tracer;
|
||||
this.parent = parent;
|
||||
this.delegate = delegate;
|
||||
}
|
||||
|
||||
@Override public void onFailure(Throwable ex) {
|
||||
continueSpan();
|
||||
if (log.isDebugEnabled()) {
|
||||
log.debug("Calling on failure of the delegate");
|
||||
}
|
||||
this.delegate.onFailure(ex);
|
||||
finish();
|
||||
}
|
||||
|
||||
private void continueSpan() {
|
||||
this.tracer.continueSpan(this.parent);
|
||||
}
|
||||
|
||||
private void finish() {
|
||||
this.tracer.detach(currentSpan());
|
||||
}
|
||||
|
||||
private Span currentSpan() {
|
||||
return this.tracer.getCurrentSpan();
|
||||
}
|
||||
}
|
||||
|
||||
private static class TraceListenableFutureCallbackWrapper<T> implements ListenableFutureCallback<T> {
|
||||
|
||||
private final Tracer tracer;
|
||||
private final Span parent;
|
||||
private final ListenableFutureCallback<T> delegate;
|
||||
|
||||
private TraceListenableFutureCallbackWrapper(Tracer tracer, Span parent,
|
||||
ListenableFutureCallback<T> delegate) {
|
||||
this.tracer = tracer;
|
||||
this.parent = parent;
|
||||
this.delegate = delegate;
|
||||
}
|
||||
|
||||
@Override public void onFailure(Throwable ex) {
|
||||
new TraceFailureCallback(this.tracer, this.parent, this.delegate).onFailure(ex);
|
||||
}
|
||||
|
||||
@Override public void onSuccess(T result) {
|
||||
new TraceSuccessCallback<>(this.tracer, this.parent, this.delegate).onSuccess(result);
|
||||
}
|
||||
}
|
||||
|
||||
private static class TraceListenableFutureCallback<T> implements ListenableFutureCallback<T> {
|
||||
@@ -104,20 +239,20 @@ public class TraceAsyncRestTemplate extends AsyncRestTemplate {
|
||||
|
||||
@Override
|
||||
public void onFailure(Throwable ex) {
|
||||
continueSpan();
|
||||
if (log.isDebugEnabled()) {
|
||||
log.debug("The callback failed - will close the span");
|
||||
}
|
||||
continueSpan();
|
||||
this.tracer.addTag(Span.SPAN_ERROR_TAG_NAME, ExceptionUtils.getExceptionMessage(ex));
|
||||
finish();
|
||||
}
|
||||
|
||||
@Override
|
||||
public void onSuccess(T result) {
|
||||
continueSpan();
|
||||
if (log.isDebugEnabled()) {
|
||||
log.debug("The callback succeeded - will close the span");
|
||||
}
|
||||
continueSpan();
|
||||
finish();
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,136 @@
|
||||
/*
|
||||
* Copyright 2013-2016 the original author or authors.
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
package org.springframework.cloud.sleuth.instrument.async.issues.issue546;
|
||||
|
||||
import java.lang.invoke.MethodHandles;
|
||||
|
||||
import org.apache.commons.logging.Log;
|
||||
import org.apache.commons.logging.LogFactory;
|
||||
import org.junit.Test;
|
||||
import org.junit.runner.RunWith;
|
||||
import org.springframework.beans.factory.annotation.Autowired;
|
||||
import org.springframework.beans.factory.annotation.Value;
|
||||
import org.springframework.boot.autoconfigure.SpringBootApplication;
|
||||
import org.springframework.boot.test.context.SpringBootTest;
|
||||
import org.springframework.cloud.sleuth.Tracer;
|
||||
import org.springframework.core.env.Environment;
|
||||
import org.springframework.http.ResponseEntity;
|
||||
import org.springframework.test.context.junit4.SpringRunner;
|
||||
import org.springframework.util.concurrent.ListenableFuture;
|
||||
import org.springframework.web.bind.annotation.RequestMapping;
|
||||
import org.springframework.web.bind.annotation.RequestParam;
|
||||
import org.springframework.web.bind.annotation.RestController;
|
||||
import org.springframework.web.client.AsyncRestTemplate;
|
||||
import org.springframework.web.client.RestTemplate;
|
||||
|
||||
import static org.springframework.cloud.sleuth.assertions.SleuthAssertions.then;
|
||||
|
||||
/**
|
||||
* @author Marcin Grzejszczak
|
||||
*/
|
||||
@RunWith(SpringRunner.class)
|
||||
@SpringBootTest(classes = Issue546TestsApp.class,
|
||||
properties = {"ribbon.eureka.enabled=false", "feign.hystrix.enabled=false", "server.port=0"},
|
||||
webEnvironment = SpringBootTest.WebEnvironment.RANDOM_PORT)
|
||||
public class Issue546Tests {
|
||||
|
||||
@Autowired Environment environment;
|
||||
|
||||
@Test
|
||||
public void should_pass_tracing_info_when_using_callbacks() {
|
||||
new RestTemplate()
|
||||
.getForObject("http://localhost:" + port() + "/trace-async-rest-template",
|
||||
String.class);
|
||||
}
|
||||
|
||||
private int port() {
|
||||
return this.environment.getProperty("local.server.port", Integer.class);
|
||||
}
|
||||
}
|
||||
|
||||
@SpringBootApplication
|
||||
class Issue546TestsApp {
|
||||
|
||||
}
|
||||
|
||||
@RestController
|
||||
class Controller {
|
||||
private static final Log log = LogFactory.getLog(MethodHandles.lookup().lookupClass());
|
||||
|
||||
private final AsyncRestTemplate traceAsyncRestTemplate;
|
||||
private final Tracer tracer;
|
||||
|
||||
public Controller(AsyncRestTemplate traceAsyncRestTemplate, Tracer tracer) {
|
||||
this.traceAsyncRestTemplate = traceAsyncRestTemplate;
|
||||
this.tracer = tracer;
|
||||
}
|
||||
|
||||
@Value("${server.port}") private String port;
|
||||
|
||||
@RequestMapping(value = "/bean") public HogeBean bean() {
|
||||
log.info("(/bean) I got a request!");
|
||||
return new HogeBean("test", 18);
|
||||
}
|
||||
|
||||
@RequestMapping(value = "/trace-async-rest-template")
|
||||
public void asyncTest(@RequestParam(required = false) boolean isSleep)
|
||||
throws InterruptedException {
|
||||
log.info("(/trace-async-rest-template) I got a request!");
|
||||
final long traceId = tracer.getCurrentSpan().getTraceId();
|
||||
ListenableFuture<ResponseEntity<HogeBean>> res = traceAsyncRestTemplate
|
||||
.getForEntity("http://localhost:" + port + "/bean", HogeBean.class);
|
||||
if (isSleep) {
|
||||
Thread.sleep(1000);
|
||||
}
|
||||
res.addCallback(success -> {
|
||||
then(Controller.this.tracer.getCurrentSpan()).hasTraceIdEqualTo(traceId);
|
||||
log.info("(/trace-async-rest-template) success");
|
||||
then(Controller.this.tracer.getCurrentSpan()).hasTraceIdEqualTo(traceId);
|
||||
}, failure -> {
|
||||
then(Controller.this.tracer.getCurrentSpan()).hasTraceIdEqualTo(traceId);
|
||||
log.error("(/trace-async-rest-template) failure", failure);
|
||||
then(Controller.this.tracer.getCurrentSpan()).hasTraceIdEqualTo(traceId);
|
||||
});
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
class HogeBean {
|
||||
private String name;
|
||||
private int age;
|
||||
|
||||
public HogeBean(String name, int age) {
|
||||
this.name = name;
|
||||
this.age = age;
|
||||
}
|
||||
|
||||
public String getName() {
|
||||
return this.name;
|
||||
}
|
||||
|
||||
public void setName(String name) {
|
||||
this.name = name;
|
||||
}
|
||||
|
||||
public int getAge() {
|
||||
return this.age;
|
||||
}
|
||||
|
||||
public void setAge(int age) {
|
||||
this.age = age;
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user