diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/B3Utils.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/B3Utils.java index 3ba575a24..ee39594bc 100644 --- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/B3Utils.java +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/B3Utils.java @@ -136,6 +136,10 @@ public final class B3Utils { return null; } + public boolean isSampled() { + return this == SAMPLED || this == DEBUG; + } + @Override public String toString() { return String.valueOf(this.sampledChar); } diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/TraceHttpAutoConfiguration.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/TraceHttpAutoConfiguration.java index a49b0c3a9..d5ee41d08 100644 --- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/TraceHttpAutoConfiguration.java +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/TraceHttpAutoConfiguration.java @@ -24,6 +24,7 @@ import org.springframework.boot.autoconfigure.condition.ConditionalOnClass; import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean; import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingClass; import org.springframework.boot.context.properties.EnableConfigurationProperties; +import org.springframework.cloud.sleuth.Sampler; import org.springframework.cloud.sleuth.TraceKeys; import org.springframework.cloud.sleuth.Tracer; import org.springframework.cloud.sleuth.autoconfig.TraceAutoConfiguration; @@ -52,8 +53,9 @@ public class TraceHttpAutoConfiguration { @Bean @ConditionalOnMissingBean - public HttpSpanExtractor httpSpanExtractor(SkipPatternProvider skipPatternProvider) { - return new ZipkinHttpSpanExtractor(skipPatternProvider.skipPattern()); + public HttpSpanExtractor httpSpanExtractor(SkipPatternProvider skipPatternProvider, + Sampler sampler) { + return new ZipkinHttpSpanExtractor(skipPatternProvider.skipPattern(), sampler); } @Bean diff --git a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/ZipkinHttpSpanExtractor.java b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/ZipkinHttpSpanExtractor.java index 10137ed76..1f8658518 100644 --- a/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/ZipkinHttpSpanExtractor.java +++ b/spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/instrument/web/ZipkinHttpSpanExtractor.java @@ -1,14 +1,18 @@ package org.springframework.cloud.sleuth.instrument.web; -import java.lang.invoke.MethodHandles; import java.util.Map; import java.util.Random; import java.util.regex.Pattern; import org.apache.commons.logging.LogFactory; +import org.springframework.beans.BeansException; +import org.springframework.beans.factory.BeanFactory; +import org.springframework.beans.factory.BeanFactoryAware; import org.springframework.cloud.sleuth.B3Utils; +import org.springframework.cloud.sleuth.Sampler; import org.springframework.cloud.sleuth.Span; import org.springframework.cloud.sleuth.SpanTextMap; +import org.springframework.cloud.sleuth.sampler.AlwaysSampler; import org.springframework.util.StringUtils; /** @@ -17,10 +21,10 @@ import org.springframework.util.StringUtils; * @author Marcin Grzejszczak * @since 1.2.0 */ -public class ZipkinHttpSpanExtractor implements HttpSpanExtractor { +public class ZipkinHttpSpanExtractor implements HttpSpanExtractor, BeanFactoryAware { private static final org.apache.commons.logging.Log log = LogFactory.getLog( - MethodHandles.lookup().lookupClass()); + ZipkinHttpSpanExtractor.class); private static final String HTTP_COMPONENT = "http"; @@ -28,12 +32,24 @@ public class ZipkinHttpSpanExtractor implements HttpSpanExtractor { private final Pattern skipPattern; private final Random random; + private Sampler sampler; + private BeanFactory beanFactory; + /** + * @deprecated use {@link ZipkinHttpSpanExtractor#ZipkinHttpSpanExtractor(Pattern, Sampler)} + */ + @Deprecated public ZipkinHttpSpanExtractor(Pattern skipPattern) { this.skipPattern = skipPattern; this.random = new Random(); } + public ZipkinHttpSpanExtractor(Pattern skipPattern, Sampler sampler) { + this.skipPattern = skipPattern; + this.random = new Random(); + this.sampler = sampler; + } + @Override public Span joinTrace(SpanTextMap textMap) { Map carrier = SPAN_CARRIER_MAPPER.convert(textMap); @@ -132,23 +148,46 @@ public class ZipkinHttpSpanExtractor implements HttpSpanExtractor { // trace, span id were retrieved from the headers and span is sampled span.shared(!(skip || idToBeGenerated)); boolean debug = sampled == B3Utils.Sampled.DEBUG; - if (debug) { - span.exportable(true); - } else if (skip) { - span.exportable(false); - } for (Map.Entry entry : carrier.entrySet()) { if (entry.getKey().toLowerCase() .startsWith(ZipkinHttpSpanMapper.BAGGAGE_PREFIX)) { span.baggage(unprefixedKey(entry.getKey()), entry.getValue()); } } + if (debug) { + span.exportable(true); + } else if (skip) { + span.exportable(false); + } else { + span.exportable(sampled == null ? + sampler().isSampled(span.build()) : sampled.isSampled()); + } return span.build(); } + private Sampler sampler() { + // the new approach + if (this.sampler != null) { + return this.sampler; + } + // fallback not to break the API + if (this.beanFactory != null) { + this.sampler = this.beanFactory.getBean(Sampler.class); + } else { + // if somehow bean factory wasn't set it will behave as previously + // this however should happen only in tests + this.sampler = new AlwaysSampler(); + } + return this.sampler; + } + private String unprefixedKey(String key) { return key.substring(key.indexOf(ZipkinHttpSpanMapper.HEADER_DELIMITER) + 1) .toLowerCase(); } + @Override + public void setBeanFactory(BeanFactory beanFactory) throws BeansException { + this.beanFactory = beanFactory; + } } diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/B3UtilsTest.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/B3UtilsTest.java index 12cd29d94..2421456b0 100644 --- a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/B3UtilsTest.java +++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/B3UtilsTest.java @@ -165,6 +165,18 @@ public class B3UtilsTest { BDDAssertions.then(sampled).isEqualTo(B3Utils.Sampled.NOT_SAMPLED); } + @Test public void should_return_is_sampled_if_sampled() { + BDDAssertions.then(B3Utils.Sampled.SAMPLED.isSampled()).isTrue(); + } + + @Test public void should_return_is_not_sampled_if_not_sampled() { + BDDAssertions.then(B3Utils.Sampled.NOT_SAMPLED.isSampled()).isFalse(); + } + + @Test public void should_return_is_sampled_if_debug() { + BDDAssertions.then(B3Utils.Sampled.DEBUG.isSampled()).isTrue(); + } + @Test public void should_read_debug_id_from_fallback() { Map map = new HashMap<>(); map.put("fallbackFlags", "1"); diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/HttpServletRequestExtractorTests.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/HttpServletRequestExtractorTests.java index 7559c9e46..71751dbec 100644 --- a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/HttpServletRequestExtractorTests.java +++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/HttpServletRequestExtractorTests.java @@ -16,19 +16,22 @@ package org.springframework.cloud.sleuth.instrument.web; +import javax.servlet.http.HttpServletRequest; import java.util.Arrays; import java.util.Random; import java.util.Vector; import java.util.regex.Pattern; -import javax.servlet.http.HttpServletRequest; import org.junit.Before; import org.junit.Test; import org.junit.runner.RunWith; import org.mockito.BDDMockito; import org.mockito.Mock; +import org.mockito.Mockito; import org.mockito.runners.MockitoJUnitRunner; +import org.springframework.cloud.sleuth.Sampler; import org.springframework.cloud.sleuth.Span; +import org.springframework.cloud.sleuth.sampler.NeverSampler; import static org.springframework.cloud.sleuth.assertions.SleuthAssertions.then; @@ -271,4 +274,88 @@ public class HttpServletRequestExtractorTests { .willReturn(Span.SPAN_SAMPLED); return hex128Bits; } + + @Test + public void should_call_sampler_when_trace_id_and_span_id_are_set_and_sampled_flag_is_not() { + HttpServletRequest request = stubWithoutSampledFlag(Mockito.mock(HttpServletRequest.class)); + Sampler sampler = new FirstFalseThenTrueSampler(); + + Span span = extractorWithSampler(sampler) + .joinTrace(new HttpServletRequestTextMap(request)); + + then(span) + .isNotNull() + .hasTraceIdEqualTo(10L) + .hasSpanIdEqualTo(20L) + .isNotExportable(); + + request = stubWithoutSampledFlag(Mockito.mock(HttpServletRequest.class)); + span = extractorWithSampler(sampler) + .joinTrace(new HttpServletRequestTextMap(request)); + + then(span) + .isNotNull() + .hasTraceIdEqualTo(10L) + .hasSpanIdEqualTo(20L) + .isExportable(); + } + + @Test + public void should_not_call_sampler_when_trace_id_and_span_id_are_set_and_sampled_flag_is_set_too() { + HttpServletRequest request = stubWithSampledFlag(Mockito.mock(HttpServletRequest.class)); + Sampler sampler = new NeverSampler(); + + Span span = extractorWithSampler(sampler) + .joinTrace(new HttpServletRequestTextMap(request)); + + then(span) + .isNotNull() + .hasTraceIdEqualTo(10L) + .hasSpanIdEqualTo(20L) + .isExportable(); + } + + private HttpServletRequest stubWithoutSampledFlag(HttpServletRequest request) { + BDDMockito.given(request.getHeaderNames()) + .willReturn(new Vector<>(Arrays.asList(Span.TRACE_ID_NAME, + Span.SPAN_ID_NAME)).elements()); + BDDMockito.given(request.getHeader(Span.TRACE_ID_NAME)) + .willReturn(Span.idToHex(10L)); + BDDMockito.given(request.getHeader(Span.SPAN_ID_NAME)) + .willReturn(Span.idToHex(20L)); + BDDMockito.given(request.getRequestURI()).willReturn("http://foo.com"); + BDDMockito.given(request.getContextPath()).willReturn("/"); + return request; + } + + private HttpServletRequest stubWithSampledFlag(HttpServletRequest request) { + BDDMockito.given(request.getHeaderNames()) + .willReturn(new Vector<>(Arrays.asList(Span.TRACE_ID_NAME, + Span.SPAN_ID_NAME, Span.SAMPLED_NAME)).elements()); + BDDMockito.given(request.getHeader(Span.TRACE_ID_NAME)) + .willReturn(Span.idToHex(10L)); + BDDMockito.given(request.getHeader(Span.SPAN_ID_NAME)) + .willReturn(Span.idToHex(20L)); + BDDMockito.given(request.getHeader(Span.SAMPLED_NAME)) + .willReturn(Span.SPAN_SAMPLED); + BDDMockito.given(request.getRequestURI()).willReturn("http://foo.com"); + BDDMockito.given(request.getContextPath()).willReturn("/"); + return request; + } + + private ZipkinHttpSpanExtractor extractorWithSampler(Sampler sampler) { + return new ZipkinHttpSpanExtractor( + Pattern.compile(""), sampler); + } } + +class FirstFalseThenTrueSampler implements Sampler { + + int counter = 0; + + @Override public boolean isSampled(Span span) { + boolean sampled = counter > 0; + counter++; + return sampled; + } +} \ No newline at end of file diff --git a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/TraceFilterTests.java b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/TraceFilterTests.java index fb6de94f5..bb2fb2932 100644 --- a/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/TraceFilterTests.java +++ b/spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/TraceFilterTests.java @@ -16,6 +16,11 @@ package org.springframework.cloud.sleuth.instrument.web; +import javax.servlet.ServletException; +import javax.servlet.ServletOutputStream; +import javax.servlet.ServletRequest; +import javax.servlet.ServletResponse; +import javax.servlet.WriteListener; import java.io.IOException; import java.util.ArrayList; import java.util.List; @@ -59,12 +64,6 @@ import org.springframework.mock.web.MockHttpServletResponse; import org.springframework.mock.web.MockServletContext; import org.springframework.test.web.servlet.request.MockHttpServletRequestBuilder; -import javax.servlet.ServletException; -import javax.servlet.ServletOutputStream; -import javax.servlet.ServletRequest; -import javax.servlet.ServletResponse; -import javax.servlet.WriteListener; - import static org.junit.Assert.assertEquals; import static org.mockito.MockitoAnnotations.initMocks; import static org.springframework.cloud.sleuth.assertions.SleuthAssertions.assertThat; @@ -291,6 +290,7 @@ public class TraceFilterTests { .header(Span.TRACE_ID_NAME, 20L).buildRequest(new MockServletContext()); this.traceKeys.getHttp().getHeaders().add("x-foo"); BeanFactory beanFactory = beanFactory(); + BDDMockito.given(beanFactory.getBean(Sampler.class)).willReturn(new AlwaysSampler()); BDDMockito.given(beanFactory.getBean(SpanReporter.class)).willReturn(this.spanReporter); TraceFilter filter = new TraceFilter(beanFactory); diff --git a/spring-cloud-sleuth-samples/spring-cloud-sleuth-sample-messaging/src/test/java/integration/MessagingApplicationTests.java b/spring-cloud-sleuth-samples/spring-cloud-sleuth-sample-messaging/src/test/java/integration/MessagingApplicationTests.java index e3c47a2d7..00eb72fdc 100644 --- a/spring-cloud-sleuth-samples/spring-cloud-sleuth-sample-messaging/src/test/java/integration/MessagingApplicationTests.java +++ b/spring-cloud-sleuth-samples/spring-cloud-sleuth-sample-messaging/src/test/java/integration/MessagingApplicationTests.java @@ -25,6 +25,7 @@ import org.junit.Test; import org.junit.runner.RunWith; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.boot.test.context.SpringBootTest; +import org.springframework.cloud.sleuth.sampler.AlwaysSampler; import org.springframework.context.annotation.Bean; import org.springframework.context.annotation.Configuration; import org.springframework.test.annotation.DirtiesContext; @@ -186,5 +187,9 @@ public class MessagingApplicationTests extends AbstractIntegrationTest { Reporter integrationTestZipkinSpanReporter() { return new IntegrationTestZipkinSpanReporter(); } + + @Bean AlwaysSampler sampler() { + return new AlwaysSampler(); + } } }