Added rate limit sampler via a property

fixes gh-1162

docs/src/main/asciidoc/spring-cloud-sleuth.adoc

@@ -331,6 +331,8 @@ include::../../../../spring-cloud-sleuth-core/src/test/java/org/springframework/
 TIP: You can set the HTTP header `X-B3-Flags` to `1`, or, when doing messaging, you can set the `spanFlags` header to `1`.
 Doing so forces the current span to be exportable regardless of the sampling decision.
 
+In order to use the rate-limited sampler set the `spring.sleuth.sampler.rate` property to choose an amount of traces to accept on a per-second interval. The minimum number is 0 and the max is 2,147,483,647 (max int).
+
 == Propagation
 
 Propagation is needed to ensure activities originating from the same root are collected together in the same trace.

spring-cloud-sleuth-core/pom.xml

@@ -121,6 +121,11 @@
 			<groupId>org.springframework</groupId>
 			<artifactId>spring-context</artifactId>
 		</dependency>
+		<dependency>
+			<groupId>org.springframework.cloud</groupId>
+			<artifactId>spring-cloud-context</artifactId>
+			<optional>true</optional>
+		</dependency>
 		<dependency>
 			<groupId>com.netflix.hystrix</groupId>
 			<artifactId>hystrix-core</artifactId>

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/sampler/RateLimitingSampler.java

@@ -0,0 +1,47 @@
+/*
+ * Copyright 2013-2018 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.sampler;
+
+import brave.sampler.Sampler;
+
+/**
+ * The rate-limited sampler allows you to choose an amount of traces to accept on a
+ * per-second interval. The minimum number is 0 and the max is 2,147,483,647 (max int).
+ *
+ * You can read more about it in {@link brave.sampler.RateLimitingSampler}
+ *
+ * @author Marcin Grzejszczak
+ * @since 2.1.0
+ */
+class RateLimitingSampler extends Sampler {
+
+	private final Sampler sampler;
+
+	public RateLimitingSampler(SamplerProperties configuration) {
+		this.sampler = brave.sampler.RateLimitingSampler.create(rateLimit(configuration));
+	}
+
+	private Integer rateLimit(SamplerProperties configuration) {
+		return configuration.getRate() != null ? configuration.getRate() : 0;
+	}
+
+	@Override
+	public boolean isSampled(long traceId) {
+		return this.sampler.isSampled(traceId);
+	}
+
+}

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/sampler/SamplerAutoConfiguration.java

@@ -0,0 +1,72 @@
+/*
+ * Copyright 2013-2018 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.sampler;
+
+import brave.sampler.Sampler;
+
+import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
+import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
+import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
+import org.springframework.boot.context.properties.EnableConfigurationProperties;
+import org.springframework.cloud.context.config.annotation.RefreshScope;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+
+/**
+ * {@link org.springframework.boot.autoconfigure.EnableAutoConfiguration
+ * Auto-configuration} to setup sampling for Spring Cloud Sleuth.
+ *
+ * @author Marcin Grzejszczak
+ * @since 2.1.0
+ */
+@Configuration
+@ConditionalOnProperty(value = "spring.sleuth.enabled", matchIfMissing = true)
+@EnableConfigurationProperties(SamplerProperties.class)
+public class SamplerAutoConfiguration {
+
+	@Configuration
+	@ConditionalOnBean(type = "org.springframework.cloud.context.scope.refresh.RefreshScope")
+	protected static class RefreshScopedSamplerConfiguration {
+
+		@Bean
+		@RefreshScope
+		@ConditionalOnMissingBean
+		public Sampler defaultTraceSampler(SamplerProperties config) {
+			return samplerFromProps(config);
+		}
+
+	}
+
+	@Configuration
+	@ConditionalOnMissingBean(type = "org.springframework.cloud.context.scope.refresh.RefreshScope")
+	protected static class NonRefreshScopeSamplerConfiguration {
+
+		@Bean
+		@ConditionalOnMissingBean
+		public Sampler defaultTraceSampler(SamplerProperties config) {
+			return samplerFromProps(config);
+		}
+
+	}
+
+	static Sampler samplerFromProps(SamplerProperties config) {
+		if (config.getRate() != null) {
+			return new RateLimitingSampler(config);
+		}
+		return new ProbabilityBasedSampler(config);
+	}
+
+}

spring-cloud-sleuth-core/src/main/java/org/springframework/cloud/sleuth/sampler/SamplerProperties.java

@@ -35,6 +35,18 @@ public class SamplerProperties {
 	 */
 	private float probability = 0.1f;
 
+	/**
+	 * A rate per second can be a nice choice for low-traffic endpoints as it allows you
+	 * surge protection. For example, you may never expect the endpoint to get more than
+	 * 50 requests per second. If there was a sudden surge of traffic, to 5000 requests
+	 * per second, you would still end up with 50 traces per second. Conversely, if you
+	 * had a percentage, like 10%, the same surge would end up with 500 traces per second,
+	 * possibly overloading your storage. Amazon X-Ray includes a rate-limited sampler
+	 * (named Reservoir) for this purpose. Brave has taken the same approach via the
+	 * {@link brave.sampler.RateLimitingSampler}.
+	 */
+	private Integer rate;
+
 	public float getProbability() {
 		return this.probability;
 	}
@@ -43,4 +55,12 @@ public void setProbability(float probability) {
 		this.probability = probability;
 	}
 
+	public Integer getRate() {
+		return this.rate;
+	}
+
+	public void setRate(Integer rate) {
+		this.rate = rate;
+	}
+
 }

spring-cloud-sleuth-core/src/main/resources/META-INF/spring.factories

@@ -1,6 +1,7 @@
 # Auto Configuration
 org.springframework.boot.autoconfigure.EnableAutoConfiguration=\
 org.springframework.cloud.sleuth.annotation.SleuthAnnotationAutoConfiguration,\
+org.springframework.cloud.sleuth.sampler.SamplerAutoConfiguration,\
 org.springframework.cloud.sleuth.autoconfig.TraceAutoConfiguration,\
 org.springframework.cloud.sleuth.log.SleuthLogAutoConfiguration,\
 org.springframework.cloud.sleuth.propagation.SleuthTagPropagationAutoConfiguration,\

spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/annotation/SleuthSpanCreatorAspectMonoTests.java

@@ -16,6 +16,10 @@
 
 package org.springframework.cloud.sleuth.annotation;
 
+import java.util.ArrayList;
+import java.util.List;
+import java.util.stream.Collectors;
+
 import brave.Span;
 import brave.Tracer;
 import brave.sampler.Sampler;
@@ -24,28 +28,27 @@
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import reactor.core.publisher.Mono;
+import zipkin2.Annotation;
+import zipkin2.reporter.Reporter;
+
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
 import org.springframework.boot.test.context.SpringBootTest;
 import org.springframework.cloud.sleuth.util.ArrayListSpanReporter;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.data.util.Pair;
-import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;
-import reactor.core.publisher.Mono;
-import zipkin2.Annotation;
-import zipkin2.reporter.Reporter;
-
-import java.util.ArrayList;
-import java.util.List;
-import java.util.stream.Collectors;
+import org.springframework.test.annotation.DirtiesContext;
+import org.springframework.test.context.junit4.SpringRunner;
 
 import static org.assertj.core.api.BDDAssertions.then;
 import static org.springframework.cloud.sleuth.annotation.SleuthSpanCreatorAspectMonoTests.TestBean.TEST_STRING;
 import static reactor.core.publisher.Mono.just;
 
 @SpringBootTest(classes = SleuthSpanCreatorAspectMonoTests.TestConfiguration.class)
-@RunWith(SpringJUnit4ClassRunner.class)
+@RunWith(SpringRunner.class)
+@DirtiesContext
 public class SleuthSpanCreatorAspectMonoTests {
 
 	@Autowired

spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/instrument/web/TraceWebFluxTests.java

@@ -17,6 +17,7 @@
 package org.springframework.cloud.sleuth.instrument.web;
 
 import java.util.List;
+import java.util.stream.Collectors;
 
 import brave.Span;
 import brave.Tracer;
@@ -30,6 +31,8 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.slf4j.MDC;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
 
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.boot.WebApplicationType;
@@ -38,7 +41,6 @@
 import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
 import org.springframework.boot.builder.SpringApplicationBuilder;
 import org.springframework.cloud.sleuth.DisableWebFluxSecurity;
-
 import org.springframework.cloud.sleuth.annotation.ContinueSpan;
 import org.springframework.cloud.sleuth.annotation.NewSpan;
 import org.springframework.cloud.sleuth.instrument.reactor.TraceReactorAutoConfigurationAccessorConfiguration;
@@ -48,7 +50,6 @@
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.core.env.Environment;
-import org.springframework.stereotype.Component;
 import org.springframework.test.web.reactive.server.WebTestClient;
 import org.springframework.web.bind.annotation.GetMapping;
 import org.springframework.web.bind.annotation.PathVariable;
@@ -60,9 +61,6 @@
 import org.springframework.web.reactive.function.server.RouterFunction;
 import org.springframework.web.reactive.function.server.RouterFunctions;
 import org.springframework.web.reactive.function.server.ServerResponse;
-import reactor.core.publisher.Flux;
-import reactor.core.publisher.Mono;
-import zipkin2.reporter.Reporter;
 
 import static org.assertj.core.api.BDDAssertions.then;
 
@@ -105,6 +103,7 @@ public void should_instrument_web_filter() throws Exception {
 		ClientResponse nonSampledResponse = whenNonSampledRequestIsSent(port);
 		// then
 		thenNoSpanWasReported(accumulator, nonSampledResponse, controller2);
+		accumulator.clear();
 
 		// when
 		ClientResponse skippedPatternResponse = whenRequestIsSentToSkippedPattern(port);
@@ -138,21 +137,24 @@ private void thenSpanWasReportedWithTags(ArrayListSpanReporter accumulator,
 			ClientResponse response) {
 		Awaitility.await().untilAsserted(() -> {
 			then(response.statusCode().value()).isEqualTo(200);
-			then(accumulator.getSpans()).hasSize(1);
 		});
-		then(accumulator.getSpans().get(0).name()).isEqualTo("get /api/c2/{id}");
-		then(accumulator.getSpans().get(0).tags())
-				.containsEntry("mvc.controller.method", "successful")
+		List<zipkin2.Span> spans = accumulator.getSpans().stream()
+				.filter(span -> span.name().equals("get /api/c2/{id}"))
+				.collect(Collectors.toList());
+		then(spans).hasSize(1);
+		then(spans.get(0).name()).isEqualTo("get /api/c2/{id}");
+		then(spans.get(0).tags()).containsEntry("mvc.controller.method", "successful")
 				.containsEntry("mvc.controller.class", "Controller2");
 	}
 
 	private void thenSpanWasReportedForFunction(ArrayListSpanReporter accumulator,
 			ClientResponse response) {
 		Awaitility.await().untilAsserted(() -> {
 			then(response.statusCode().value()).isEqualTo(200);
-			then(accumulator.getSpans()).hasSize(1);
 		});
-		then(accumulator.getSpans().get(0).name()).isEqualTo("get");
+		List<zipkin2.Span> spans = accumulator.getSpans().stream()
+				.filter(span -> span.name().equals("get")).collect(Collectors.toList());
+		then(spans).hasSize(1);
 	}
 
 	private void thenNoSpanWasReported(ArrayListSpanReporter accumulator,

spring-cloud-sleuth-core/src/test/java/org/springframework/cloud/sleuth/sampler/SamplerAutoConfigurationTests.java

@@ -0,0 +1,32 @@
+package org.springframework.cloud.sleuth.sampler;
+
+import brave.sampler.Sampler;
+import org.assertj.core.api.BDDAssertions;
+import org.junit.Test;
+
+/**
+ * @author Marcin Grzejszczak
+ * @since
+ */
+public class SamplerAutoConfigurationTests {
+
+	@Test
+	public void should_use_rate_limit_sampler_when_property_set() {
+		SamplerProperties properties = new SamplerProperties();
+		properties.setRate(10);
+
+		Sampler sampler = SamplerAutoConfiguration.samplerFromProps(properties);
+
+		BDDAssertions.then(sampler).isInstanceOf(RateLimitingSampler.class);
+	}
+
+	@Test
+	public void should_use_probability_sampler_when_rate_limiting_not_set() {
+		SamplerProperties properties = new SamplerProperties();
+
+		Sampler sampler = SamplerAutoConfiguration.samplerFromProps(properties);
+
+		BDDAssertions.then(sampler).isInstanceOf(ProbabilityBasedSampler.class);
+	}
+
+}

spring-cloud-sleuth-zipkin/pom.xml

@@ -45,11 +45,6 @@
 			<groupId>org.springframework.cloud</groupId>
 			<artifactId>spring-cloud-commons</artifactId>
 		</dependency>
-		<dependency>
-			<groupId>org.springframework.cloud</groupId>
-			<artifactId>spring-cloud-context</artifactId>
-			<optional>true</optional>
-		</dependency>
 		<dependency>
 			<groupId>org.springframework.boot</groupId>
 			<artifactId>spring-boot-actuator</artifactId>