diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/AbstractTraceHttpRequestInterceptor.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/AbstractTraceHttpRequestInterceptor.java new file mode 100644 index 000000000..9856068b0 --- /dev/null +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/AbstractTraceHttpRequestInterceptor.java @@ -0,0 +1,99 @@ +/* + * 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.web.client; + +import org.springframework.cloud.sleuth.Span; +import org.springframework.cloud.sleuth.SpanAccessor; +import org.springframework.cloud.sleuth.event.ClientReceivedEvent; +import org.springframework.context.ApplicationEvent; +import org.springframework.context.ApplicationEventPublisher; +import org.springframework.context.ApplicationEventPublisherAware; +import org.springframework.http.HttpRequest; +import org.springframework.util.StringUtils; + +/** + * Abstraction over classes that interact with Http requests. Allows you + * to enrich the request headers with trace related information. + * + * @author Marcin Grzejszczak + */ +public abstract class AbstractTraceHttpRequestInterceptor + implements ApplicationEventPublisherAware { + + private ApplicationEventPublisher publisher; + private final SpanAccessor accessor; + + public AbstractTraceHttpRequestInterceptor(SpanAccessor accessor) { + this.accessor = accessor; + } + + @Override + public void setApplicationEventPublisher(ApplicationEventPublisher publisher) { + this.publisher = publisher; + } + + /** + * Adds trace related headers from the span to the request + */ + public void enrichWithTraceHeaders(HttpRequest request, Span span) { + setHeader(request, Span.TRACE_ID_NAME, span.getTraceId()); + setHeader(request, Span.SPAN_ID_NAME, span.getSpanId()); + if (!span.isExportable()) { + setHeader(request, Span.NOT_SAMPLED_NAME, "true"); + } + setHeader(request, Span.SPAN_NAME_NAME, span.getName().toString()); + setHeader(request, Span.PARENT_ID_NAME, getParentId(span)); + setHeader(request, Span.PROCESS_ID_NAME, span.getProcessId()); + } + + private Long getParentId(Span span) { + return !span.getParents().isEmpty() ? span.getParents().get(0) : null; + } + + public void setHeader(HttpRequest request, String name, String value) { + if (StringUtils.hasText(value) && !request.getHeaders().containsKey(name) && this.accessor.isTracing()) { + request.getHeaders().add(name, value); + } + } + + public void setHeader(HttpRequest request, String name, Long value) { + if (value != null) { + setHeader(request, name, Span.toHex(value)); + } + } + + /** + * Close the current span and emit the ClientReceivedEvent + */ + public void close() { + if (getCurrentSpan() == null) { + return; + } + publish(new ClientReceivedEvent(this, getCurrentSpan())); + } + + protected void publish(ApplicationEvent event) { + if (this.publisher != null) { + this.publisher.publishEvent(event); + } + } + + protected Span getCurrentSpan() { + return this.accessor.getCurrentSpan(); + } + +} diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceAsyncClientHttpRequestFactoryWrapper.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceAsyncClientHttpRequestFactoryWrapper.java new file mode 100644 index 000000000..3d66bf3fc --- /dev/null +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceAsyncClientHttpRequestFactoryWrapper.java @@ -0,0 +1,64 @@ +/* + * 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.web.client; + +import java.io.IOException; +import java.net.URI; + +import org.springframework.cloud.sleuth.Span; +import org.springframework.cloud.sleuth.SpanAccessor; +import org.springframework.cloud.sleuth.event.ClientSentEvent; +import org.springframework.context.ApplicationEventPublisher; +import org.springframework.http.HttpMethod; +import org.springframework.http.client.AsyncClientHttpRequest; +import org.springframework.http.client.AsyncClientHttpRequestFactory; + +/** + * Wrapper that adds trace related headers to the created AsyncClientHttpRequest + * + * @see org.springframework.web.client.RestTemplate + * @see SpanAccessor + * + * @author Marcin Grzejszczak + * @author Spencer Gibb + */ +public class TraceAsyncClientHttpRequestFactoryWrapper extends AbstractTraceHttpRequestInterceptor + implements AsyncClientHttpRequestFactory { + + private ApplicationEventPublisher publisher; + private final AsyncClientHttpRequestFactory delegate; + + public TraceAsyncClientHttpRequestFactoryWrapper(SpanAccessor accessor, + AsyncClientHttpRequestFactory delegate) { + super(accessor); + this.delegate = delegate; + } + + @Override + public AsyncClientHttpRequest createAsyncRequest(URI uri, HttpMethod httpMethod) + throws IOException { + AsyncClientHttpRequest request = this.delegate.createAsyncRequest(uri, httpMethod); + Span span = getCurrentSpan(); + if (span == null) { + setHeader(request, Span.NOT_SAMPLED_NAME, "true"); + return request; + } + enrichWithTraceHeaders(request, span); + publish(new ClientSentEvent(this, span)); + return request; + } +} diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceAsyncListenableTaskExecutor.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceAsyncListenableTaskExecutor.java new file mode 100644 index 000000000..9cda09e6d --- /dev/null +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceAsyncListenableTaskExecutor.java @@ -0,0 +1,71 @@ +/* + * 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.web.client; + +import java.util.concurrent.Callable; +import java.util.concurrent.Future; + +import org.springframework.cloud.sleuth.Tracer; +import org.springframework.core.task.AsyncListenableTaskExecutor; +import org.springframework.util.concurrent.ListenableFuture; + +/** + * AsyncListenableTaskExecutor that wraps all Runnable / Callable tasks into + * their trace related representation + */ +public class TraceAsyncListenableTaskExecutor implements AsyncListenableTaskExecutor { + + private final AsyncListenableTaskExecutor delegate; + private final Tracer tracer; + + TraceAsyncListenableTaskExecutor(AsyncListenableTaskExecutor delegate, + Tracer tracer) { + this.delegate = delegate; + this.tracer = tracer; + } + + @Override + public ListenableFuture submitListenable(Runnable task) { + return this.delegate.submitListenable(this.tracer.wrap(task)); + } + + @Override + public ListenableFuture submitListenable(Callable task) { + return this.delegate.submitListenable(this.tracer.wrap(task)); + } + + @Override + public void execute(Runnable task, long startTimeout) { + this.delegate.execute(this.tracer.wrap(task), startTimeout); + } + + @Override + public Future submit(Runnable task) { + return this.delegate.submit(this.tracer.wrap(task)); + } + + @Override + public Future submit(Callable task) { + return this.delegate.submit(this.tracer.wrap(task)); + } + + @Override + public void execute(Runnable task) { + this.delegate.execute(this.tracer.wrap(task)); + } + +} \ No newline at end of file diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceRestTemplateInterceptor.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceRestTemplateInterceptor.java index 051e837ef..81730bb82 100644 --- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceRestTemplateInterceptor.java +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceRestTemplateInterceptor.java @@ -1,5 +1,5 @@ /* - * Copyright 2013-2015 the original author or authors. + * 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. @@ -19,16 +19,11 @@ import java.io.IOException; import org.springframework.cloud.sleuth.Span; import org.springframework.cloud.sleuth.SpanAccessor; -import org.springframework.cloud.sleuth.event.ClientReceivedEvent; import org.springframework.cloud.sleuth.event.ClientSentEvent; -import org.springframework.context.ApplicationEvent; -import org.springframework.context.ApplicationEventPublisher; -import org.springframework.context.ApplicationEventPublisherAware; import org.springframework.http.HttpRequest; import org.springframework.http.client.ClientHttpRequestExecution; import org.springframework.http.client.ClientHttpRequestInterceptor; import org.springframework.http.client.ClientHttpResponse; -import org.springframework.util.StringUtils; /** * Interceptor that verifies whether the trance and span id has been set on the request @@ -37,23 +32,14 @@ import org.springframework.util.StringUtils; * @see org.springframework.web.client.RestTemplate * @see SpanAccessor * - * @author Marcin Grzejszczak, 4financeIT + * @author Marcin Grzejszczak * @author Spencer Gibb */ -public class TraceRestTemplateInterceptor - implements ClientHttpRequestInterceptor, ApplicationEventPublisherAware { - - private ApplicationEventPublisher publisher; - - private SpanAccessor accessor; +public class TraceRestTemplateInterceptor extends AbstractTraceHttpRequestInterceptor + implements ClientHttpRequestInterceptor { public TraceRestTemplateInterceptor(SpanAccessor accessor) { - this.accessor = accessor; - } - - @Override - public void setApplicationEventPublisher(ApplicationEventPublisher publisher) { - this.publisher = publisher; + super(accessor); } @Override @@ -64,49 +50,9 @@ public class TraceRestTemplateInterceptor setHeader(request, Span.NOT_SAMPLED_NAME, "true"); return execution.execute(request, body); } - setHeader(request, Span.TRACE_ID_NAME, span.getTraceId()); - setHeader(request, Span.SPAN_ID_NAME, span.getSpanId()); - if (!span.isExportable()) { - setHeader(request, Span.NOT_SAMPLED_NAME, "true"); - } - setHeader(request, Span.SPAN_NAME_NAME, span.getName().toString()); - setHeader(request, Span.PARENT_ID_NAME, getParentId(span)); - setHeader(request, Span.PROCESS_ID_NAME, span.getProcessId()); + enrichWithTraceHeaders(request, span); publish(new ClientSentEvent(this, span)); return new TraceHttpResponse(this, execution.execute(request, body)); } - public void close() { - if (getCurrentSpan() == null) { - return; - } - publish(new ClientReceivedEvent(this, getCurrentSpan())); - } - - private void publish(ApplicationEvent event) { - if (this.publisher != null) { - this.publisher.publishEvent(event); - } - } - - private Long getParentId(Span span) { - return !span.getParents().isEmpty() ? span.getParents().get(0) : null; - } - - public void setHeader(HttpRequest request, String name, String value) { - if (StringUtils.hasText(value) && !request.getHeaders().containsKey(name) && this.accessor.isTracing()) { - request.getHeaders().add(name, value); - } - } - - public void setHeader(HttpRequest request, String name, Long value) { - if (value != null) { - setHeader(request, name, Span.toHex(value)); - } - } - - private Span getCurrentSpan() { - return this.accessor.getCurrentSpan(); - } - } diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceWebClientAutoConfiguration.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceWebClientAutoConfiguration.java index 8c88cf03f..2754312e5 100644 --- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceWebClientAutoConfiguration.java +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/client/TraceWebClientAutoConfiguration.java @@ -1,5 +1,5 @@ /* - * Copyright 2013-2015 the original author or authors. + * 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. @@ -29,10 +29,16 @@ import org.springframework.boot.autoconfigure.condition.ConditionalOnClass; import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean; import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; import org.springframework.cloud.sleuth.SpanAccessor; +import org.springframework.cloud.sleuth.Tracer; import org.springframework.cloud.sleuth.autoconfig.TraceAutoConfiguration; import org.springframework.context.annotation.Bean; import org.springframework.context.annotation.Configuration; +import org.springframework.core.task.AsyncListenableTaskExecutor; +import org.springframework.http.client.AsyncClientHttpRequestFactory; import org.springframework.http.client.ClientHttpRequestInterceptor; +import org.springframework.http.client.SimpleClientHttpRequestFactory; +import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler; +import org.springframework.web.client.AsyncRestTemplate; import org.springframework.web.client.RestTemplate; /** @@ -57,6 +63,29 @@ public class TraceWebClientAutoConfiguration { return new RestTemplate(); } + + @Bean + @ConditionalOnMissingBean + public AsyncClientHttpRequestFactory asyncClientHttpRequestFactory(Tracer tracer) { + SimpleClientHttpRequestFactory simpleClientHttpRequestFactory = new SimpleClientHttpRequestFactory(); + simpleClientHttpRequestFactory.setTaskExecutor(asyncListenableTaskExecutor(tracer)); + return simpleClientHttpRequestFactory; + } + + private AsyncListenableTaskExecutor asyncListenableTaskExecutor(Tracer tracer) { + ThreadPoolTaskScheduler threadPoolTaskScheduler = new ThreadPoolTaskScheduler(); + threadPoolTaskScheduler.initialize(); + return new TraceAsyncListenableTaskExecutor(threadPoolTaskScheduler, tracer); + } + + @Bean + @ConditionalOnMissingBean + public AsyncRestTemplate asyncRestTemplate(SpanAccessor spanAccessor, + AsyncClientHttpRequestFactory asyncClientHttpRequestFactory, RestTemplate restTemplate) { + return new AsyncRestTemplate(new TraceAsyncClientHttpRequestFactoryWrapper( + spanAccessor, asyncClientHttpRequestFactory), restTemplate); + } + @Configuration protected static class TraceInterceptorConfiguration { diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/RestTemplateTraceAspectIntegrationTests.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/RestTemplateTraceAspectIntegrationTests.java index 230fad6b2..f10ec7af9 100644 --- a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/RestTemplateTraceAspectIntegrationTests.java +++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/RestTemplateTraceAspectIntegrationTests.java @@ -1,7 +1,8 @@ package org.springframework.cloud.sleuth.instrument.web; -import junitparams.JUnitParamsRunner; -import junitparams.Parameters; +import java.util.concurrent.Callable; +import java.util.concurrent.ExecutionException; + import org.junit.Before; import org.junit.ClassRule; import org.junit.Rule; @@ -11,8 +12,8 @@ import org.springframework.beans.factory.annotation.Autowired; import org.springframework.boot.test.SpringApplicationConfiguration; import org.springframework.cloud.sleuth.Span; import org.springframework.cloud.sleuth.instrument.DefaultTestAutoConfiguration; -import org.springframework.cloud.sleuth.instrument.web.common.HttpMockServer; import org.springframework.cloud.sleuth.instrument.web.common.AbstractMvcWiremockIntegrationTest; +import org.springframework.cloud.sleuth.instrument.web.common.HttpMockServer; import org.springframework.context.annotation.Import; import org.springframework.http.MediaType; import org.springframework.scheduling.annotation.EnableAsync; @@ -23,12 +24,18 @@ import org.springframework.test.web.servlet.request.MockMvcRequestBuilders; import org.springframework.web.bind.annotation.RequestMapping; import org.springframework.web.bind.annotation.RequestMethod; import org.springframework.web.bind.annotation.RestController; +import org.springframework.web.client.AsyncRestTemplate; import org.springframework.web.client.RestTemplate; import org.springframework.web.context.request.async.WebAsyncTask; -import java.util.concurrent.Callable; +import junitparams.JUnitParamsRunner; +import junitparams.Parameters; -import static com.github.tomakehurst.wiremock.client.WireMock.*; +import static com.github.tomakehurst.wiremock.client.WireMock.aResponse; +import static com.github.tomakehurst.wiremock.client.WireMock.get; +import static com.github.tomakehurst.wiremock.client.WireMock.getRequestedFor; +import static com.github.tomakehurst.wiremock.client.WireMock.matching; +import static com.github.tomakehurst.wiremock.client.WireMock.urlMatching; import static java.util.concurrent.TimeUnit.SECONDS; import static junitparams.JUnitParamsRunner.$; import static org.springframework.test.web.servlet.request.MockMvcRequestBuilders.asyncDispatch; @@ -54,6 +61,13 @@ public class RestTemplateTraceAspectIntegrationTests extends AbstractMvcWiremock thenTraceIdHasBeenSetOnARequestHeader(); } + @Test + public void should_set_span_data_on_headers_when_sending_a_request_via_async_rest_template() throws Exception { + whenARequestIsSentToAAsyncRestTemplateEndpoint(); + + thenTraceIdHasBeenSetOnARequestHeader(); + } + @Test @Parameters public void should_set_span_data_on_headers_via_aspect_in_asynchronous_call(String url) throws Exception { @@ -66,6 +80,10 @@ public class RestTemplateTraceAspectIntegrationTests extends AbstractMvcWiremock return $("/callablePing", "/webAsyncTaskPing"); } + private void whenARequestIsSentToAAsyncRestTemplateEndpoint() throws Exception { + this.mockMvc.perform(MockMvcRequestBuilders.get("/asyncRestTemplate").accept(MediaType.TEXT_PLAIN)).andReturn(); + } + private void whenARequestIsSentToASyncEndpoint() throws Exception { this.mockMvc.perform(MockMvcRequestBuilders.get("/syncPing").accept(MediaType.TEXT_PLAIN)).andReturn(); } @@ -95,6 +113,13 @@ public class RestTemplateTraceAspectIntegrationTests extends AbstractMvcWiremock @Autowired HttpMockServer httpMockServer; @Autowired RestTemplate restTemplate; + @Autowired AsyncRestTemplate asyncRestTemplate; + + @RequestMapping(value = "/asyncRestTemplate", method = RequestMethod.GET, produces = MediaType.TEXT_PLAIN_VALUE) + public String asyncRestTemplate() + throws ExecutionException, InterruptedException { + return callWiremockViaAsyncRestTemplateAndReturnOk(); + } @RequestMapping(value = "/syncPing", method = RequestMethod.GET, produces = MediaType.TEXT_PLAIN_VALUE) public String syncPing() { @@ -125,5 +150,11 @@ public class RestTemplateTraceAspectIntegrationTests extends AbstractMvcWiremock this.restTemplate.getForObject("http://localhost:" + this.httpMockServer.port(), String.class); return "OK"; } + + private String callWiremockViaAsyncRestTemplateAndReturnOk() + throws ExecutionException, InterruptedException { + this.asyncRestTemplate.getForEntity("http://localhost:" + this.httpMockServer.port(), String.class).get(); + return "OK"; + } } }