brave.sampler.SamplerFunction Java Examples

@Test public void customSampler() throws IOException {
  String path = "/foo";

  closeClient(client);

  SamplerFunction<HttpRequest> sampler = HttpRuleSampler.newBuilder()
    .putRule(pathStartsWith(path), Sampler.NEVER_SAMPLE)
    .build();

  httpTracing = httpTracing.toBuilder().clientSampler(sampler).build();
  client = newClient(server.getPort());

  server.enqueue(new MockResponse());
  get(client, path);

  assertThat(extract(takeRequest()).sampled()).isFalse();
}
 

private SamplerFunction<HttpRequest> combineUserProvidedSamplerWithSkipPatternSampler(
		@Nullable SamplerFunction<HttpRequest> serverSampler,
		@Nullable SkipPatternProvider provider) {
	SamplerFunction<HttpRequest> skipPatternSampler = provider != null
			? new SkipPatternHttpServerSampler(provider) : null;
	if (serverSampler == null && skipPatternSampler == null) {
		return SamplerFunctions.deferDecision();
	}
	else if (serverSampler == null) {
		return skipPatternSampler;
	}
	else if (skipPatternSampler == null) {
		return serverSampler;
	}
	return new CompositeHttpSampler(skipPatternSampler, serverSampler);
}
 

@Test public void customSampler() throws IOException {
  String path = "/foo";

  SamplerFunction<HttpRequest> sampler = HttpRuleSampler.newBuilder()
    .putRule(pathStartsWith(path), Sampler.NEVER_SAMPLE)
    .build();

  httpTracing = httpTracing.toBuilder().serverSampler(sampler).build();
  init();

  Request request = new Request.Builder().url(url(path)).build();
  try (Response response = client.newCall(request).execute()) {
    assertThat(response.isSuccessful()).isTrue();
  }

  // @After will check that nothing is reported
}
 

@Bean
@ConditionalOnMissingBean
// NOTE: stable bean name as might be used outside sleuth
RpcTracing rpcTracing(Tracing tracing,
		@Nullable @RpcClientSampler SamplerFunction<RpcRequest> clientSampler,
		@Nullable @RpcServerSampler SamplerFunction<RpcRequest> serverSampler,
		@Nullable List<RpcTracingCustomizer> rpcTracingCustomizers) {

	RpcTracing.Builder builder = RpcTracing.newBuilder(tracing);
	if (clientSampler != null) {
		builder.clientSampler(clientSampler);
	}
	if (serverSampler != null) {
		builder.serverSampler(serverSampler);
	}
	if (rpcTracingCustomizers != null) {
		for (RpcTracingCustomizer customizer : rpcTracingCustomizers) {
			customizer.customize(builder);
		}
	}
	return builder.build();
}
 

@Bean
@ConditionalOnMissingBean
// NOTE: stable bean name as might be used outside sleuth
MessagingTracing messagingTracing(Tracing tracing,
		@Nullable @ProducerSampler SamplerFunction<MessagingRequest> producerSampler,
		@Nullable @ConsumerSampler SamplerFunction<MessagingRequest> consumerSampler,
		@Nullable List<MessagingTracingCustomizer> messagingTracingCustomizers) {

	MessagingTracing.Builder builder = MessagingTracing.newBuilder(tracing);
	if (producerSampler != null) {
		builder.producerSampler(producerSampler);
	}
	if (consumerSampler != null) {
		builder.consumerSampler(consumerSampler);
	}
	if (messagingTracingCustomizers != null) {
		for (MessagingTracingCustomizer customizer : messagingTracingCustomizers) {
			customizer.customize(builder);
		}
	}
	return builder.build();
}
 

@Bean(name = ConsumerSampler.NAME)
SamplerFunction<MessagingRequest> myMessagingSampler() {
	return MessagingRuleSampler.newBuilder()
			.putRule(channelNameEquals("alerts"), Sampler.NEVER_SAMPLE)
			.putRule(Matchers.alwaysMatch(), RateLimitingSampler.create(100))
			.build();
}
 

@Test
public void defaultsToBraveConsumerSampler() {
	contextRunner().run((context) -> {
		SamplerFunction<MessagingRequest> consumerSampler = context
				.getBean(MessagingTracing.class).consumerSampler();

		then(consumerSampler).isSameAs(SamplerFunctions.deferDecision());
	});
}
 

@Test public void deprecatedNextSpan_samplerSeesHttpClientRequest() {
  SamplerFunction<HttpRequest> clientSampler = mock(SamplerFunction.class);
  init(httpTracingBuilder(tracingBuilder()).clientSampler(clientSampler));

  handler.nextSpan(request);

  verify(clientSampler).trySample(request);
}
 

/** Creates a potentially noop remote span representing this request */
Span nextMessagingSpan(
  SamplerFunction<MessagingRequest> sampler,
  MessagingRequest request,
  TraceContextOrSamplingFlags extracted
) {
  Boolean sampled = extracted.sampled();
  // only recreate the context if the messaging sampler made a decision
  if (sampled == null && (sampled = sampler.trySample(request)) != null) {
    extracted = extracted.sampled(sampled.booleanValue());
  }
  return tracer.nextSpan(extracted);
}
 

@Test
public void configuresUserProvidedRpcServerSampler() {
	contextRunner().withUserConfiguration(RpcServerSamplerConfig.class)
			.run((context) -> {
				SamplerFunction<RpcRequest> serverSampler = context
						.getBean(RpcTracing.class).serverSampler();

				then(serverSampler).isSameAs(RpcServerSamplerConfig.INSTANCE);
			});
}
 

@Test
public void defaultsServerSamplerToDeferWhenSkipPatternCleared() {
	contextRunner().withPropertyValues("spring.sleuth.web.skip-pattern")
			.run((context) -> {
				SamplerFunction<HttpRequest> clientSampler = context
						.getBean(HttpTracing.class).serverRequestSampler();

				then(clientSampler).isSameAs(SamplerFunctions.deferDecision());
			});
}
 

@Bean
@ConditionalOnMissingBean(name = HttpClientSampler.NAME)
SamplerFunction<HttpRequest> sleuthHttpClientSampler(
		SleuthWebProperties sleuthWebProperties) {
	String skipPattern = sleuthWebProperties.getClient().getSkipPattern();
	if (skipPattern == null) {
		return SamplerFunctions.deferDecision();
	}

	return new SkipPatternHttpClientSampler(Pattern.compile(skipPattern));
}
 

@Test
public void defaultsClientSamplerToDefer() {
	contextRunner().run((context) -> {
		SamplerFunction<HttpRequest> clientSampler = context
				.getBean(HttpTracing.class).clientRequestSampler();

		then(clientSampler).isSameAs(SamplerFunctions.deferDecision());
	});
}
 

@Test
public void configuresClientSkipPattern() throws Exception {
	contextRunner()
			.withPropertyValues("spring.sleuth.web.client.skip-pattern=foo.*|bar.*")
			.run((context) -> {
				SamplerFunction<HttpRequest> clientSampler = context
						.getBean(HttpTracing.class).clientRequestSampler();

				then(clientSampler).isInstanceOf(SkipPatternHttpClientSampler.class);
			});
}
 

@Test
public void configuresUserProvidedHttpClientSampler() {
	contextRunner().withUserConfiguration(HttpClientSamplerConfig.class)
			.run((context) -> {
				SamplerFunction<HttpRequest> clientSampler = context
						.getBean(HttpTracing.class).clientRequestSampler();

				then(clientSampler).isSameAs(HttpClientSamplerConfig.INSTANCE);
			});
}
 

@Test
public void defaultsServerSamplerToSkipPattern() {
	contextRunner().run((context) -> {
		SamplerFunction<HttpRequest> serverSampler = context
				.getBean(HttpTracing.class).serverRequestSampler();

		then(serverSampler).isInstanceOf(SkipPatternHttpServerSampler.class);
	});
}
 

private ContextConsumer<AssertableApplicationContext> thenCompositeHttpServerSamplerOf(
		SamplerFunction<HttpRequest> instance) {
	return (context) -> {

		SamplerFunction<HttpRequest> serverSampler = context
				.getBean(HttpTracing.class).serverRequestSampler();

		then(serverSampler).isInstanceOf(CompositeHttpSampler.class);

		then(((CompositeHttpSampler) serverSampler).left)
				.isInstanceOf(SkipPatternHttpServerSampler.class);
		then(((CompositeHttpSampler) serverSampler).right).isSameAs(instance);
	};
}
 

/** Creates a potentially noop remote span representing this request */
Span nextMessagingSpan(
  SamplerFunction<MessagingRequest> sampler,
  MessagingRequest request,
  TraceContextOrSamplingFlags extracted
) {
  Boolean sampled = extracted.sampled();
  // only recreate the context if the messaging sampler made a decision
  if (sampled == null && (sampled = sampler.trySample(request)) != null) {
    extracted = extracted.sampled(sampled.booleanValue());
  }
  return tracer.nextSpan(extracted);
}
 

@Test
public void defaultsToBraveProducerSampler() {
	contextRunner().run((context) -> {
		SamplerFunction<MessagingRequest> producerSampler = context
				.getBean(MessagingTracing.class).producerSampler();

		then(producerSampler).isSameAs(SamplerFunctions.deferDecision());
	});
}
 

@Test
public void configuresUserProvidedProducerSampler() {
	contextRunner().withUserConfiguration(ProducerSamplerConfig.class)
			.run((context) -> {
				SamplerFunction<MessagingRequest> producerSampler = context
						.getBean(MessagingTracing.class).producerSampler();

				then(producerSampler).isSameAs(ProducerSamplerConfig.INSTANCE);
			});
}
 

@Test
public void configuresUserProvidedRpcClientSampler() {
	contextRunner().withUserConfiguration(RpcClientSamplerConfig.class)
			.run((context) -> {
				SamplerFunction<RpcRequest> clientSampler = context
						.getBean(RpcTracing.class).clientSampler();

				then(clientSampler).isSameAs(RpcClientSamplerConfig.INSTANCE);
			});
}
 

@Test
public void configuresUserProvidedConsumerSampler() {
	contextRunner().withUserConfiguration(ConsumerSamplerConfig.class)
			.run((context) -> {
				SamplerFunction<MessagingRequest> consumerSampler = context
						.getBean(MessagingTracing.class).consumerSampler();

				then(consumerSampler).isSameAs(ConsumerSamplerConfig.INSTANCE);
			});
}
 

@Test public void handleSend_samplerSeesHttpClientRequest() {
  SamplerFunction<HttpRequest> clientSampler = mock(SamplerFunction.class);
  init(httpTracingBuilder(tracingBuilder()).clientSampler(clientSampler));

  handler.handleSend(request);

  verify(clientSampler).trySample(request);
}
 

@Test public void handleReceive_samplerSeesHttpServerRequest() {
  SamplerFunction<HttpRequest> serverSampler = mock(SamplerFunction.class);
  init(httpTracingBuilder(tracingBuilder()).serverSampler(serverSampler));

  handler.handleReceive(request);

  verify(serverSampler).trySample(request);
}
 

static HttpSampler fromHttpRequestSampler(SamplerFunction<HttpRequest> sampler) {
  if (sampler == null) throw new NullPointerException("sampler == null");
  if (sampler.equals(SamplerFunctions.deferDecision())) return HttpSampler.TRACE_ID;
  if (sampler.equals(SamplerFunctions.neverSample())) return HttpSampler.NEVER_SAMPLE;
  return sampler instanceof HttpSampler ? (HttpSampler) sampler
    : new HttpRequestSamplerAdapter(sampler);
}
 

@Bean(name = RpcServerSampler.NAME)
SamplerFunction<RpcRequest> myRpcSampler() {
	Matcher<RpcRequest> userAuth = and(serviceEquals("users.UserService"),
			methodEquals("GetUserToken"));
	return RpcRuleSampler.newBuilder()
			.putRule(serviceEquals("grpc.health.v1.Health"), Sampler.NEVER_SAMPLE)
			.putRule(userAuth, RateLimitingSampler.create(100)).build();
}
 

@Bean(name = HttpServerSampler.NAME)
SamplerFunction<HttpRequest> myHttpSampler(SkipPatternProvider provider) {
	Pattern pattern = provider.skipPattern();
	return request -> {
		String url = request.path();
		boolean shouldSkip = pattern.matcher(url).matches();
		if (shouldSkip) {
			return false;
		}
		return null;
	};
}
 

@Test public void handleReceive_samplerSeesRpcServerRequest() {
  SamplerFunction<RpcRequest> serverSampler = mock(SamplerFunction.class);
  init(httpTracingBuilder(tracingBuilder()).serverSampler(serverSampler));

  handler.handleReceive(request);

  verify(serverSampler).trySample(request);
}
 

<T> TraceContext nextContext(SamplerFunction<T> samplerFunction, T arg, TraceContext parent) {
  if (samplerFunction == null) throw new NullPointerException("samplerFunction == null");
  if (arg == null) throw new NullPointerException("arg == null");
  if (parent != null) return decorateContext(parent, parent.spanId());

  Boolean sampled = samplerFunction.trySample(arg);
  SamplingFlags flags = sampled != null ? (sampled ? SAMPLED : NOT_SAMPLED) : EMPTY;
  return newRootContext(InternalPropagation.instance.flags(flags));
}
 

@Test public void handleSend_samplerSeesRpcClientRequest() {
  SamplerFunction<RpcRequest> clientSampler = mock(SamplerFunction.class);
  init(httpTracingBuilder(tracingBuilder()).clientSampler(clientSampler));

  handler.handleSend(request);

  verify(clientSampler).trySample(request);
}