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:
Marcin Grzejszczak
2017-04-11 14:16:25 +02:00
parent 21d73f6ed6
commit cce5b18dec
2 changed files with 276 additions and 5 deletions

View File

@@ -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();
}

View File

@@ -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;
}
}