diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceAsyncRestTemplate.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceAsyncRestTemplate.java index 097b87f67..dd5d6374b 100644 --- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceAsyncRestTemplate.java +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceAsyncRestTemplate.java @@ -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 ListenableFuture doExecute(URI url, HttpMethod method, AsyncRequestCallback requestCallback, ResponseExtractor responseExtractor) throws RestClientException { - ListenableFuture future = super.doExecute(url, method, requestCallback, responseExtractor); - Span span = this.tracer.getCurrentSpan(); + final ListenableFuture 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() { + + @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 callback) { + future.addCallback(new TraceListenableFutureCallbackWrapper<>(TraceAsyncRestTemplate.this.tracer, span, callback)); + } + + @Override public void addCallback(SuccessCallback successCallback, + FailureCallback failureCallback) { + future.addCallback( + new TraceSuccessCallback<>(TraceAsyncRestTemplate.this.tracer, span, successCallback), + new TraceFailureCallback(TraceAsyncRestTemplate.this.tracer, span, failureCallback)); + } + }; + } + + private static class TraceSuccessCallback implements SuccessCallback { + + private static final Log log = LogFactory.getLog(MethodHandles.lookup().lookupClass()); + + private final Tracer tracer; + private final Span parent; + private final SuccessCallback delegate; + + private TraceSuccessCallback(Tracer tracer, Span parent, + SuccessCallback 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 implements ListenableFutureCallback { + + private final Tracer tracer; + private final Span parent; + private final ListenableFutureCallback delegate; + + private TraceListenableFutureCallbackWrapper(Tracer tracer, Span parent, + ListenableFutureCallback 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 implements ListenableFutureCallback { @@ -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(); } diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/async/issues/issue546/Issue546Tests.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/async/issues/issue546/Issue546Tests.java new file mode 100644 index 000000000..60d29a2c7 --- /dev/null +++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/async/issues/issue546/Issue546Tests.java @@ -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> 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; + } +} \ No newline at end of file