brave.sampler.Sampler Java Examples

@Bean
@ConditionalOnMissingBean
// NOTE: stable bean name as might be used outside sleuth
Tracing tracing(@LocalServiceName String serviceName, Propagation.Factory factory,
		CurrentTraceContext currentTraceContext, Sampler sampler,
		SleuthProperties sleuthProperties, @Nullable List<SpanHandler> spanHandlers,
		@Nullable List<TracingCustomizer> tracingCustomizers) {
	Tracing.Builder builder = Tracing.newBuilder().sampler(sampler)
			.localServiceName(StringUtils.isEmpty(serviceName) ? DEFAULT_SERVICE_NAME
					: serviceName)
			.propagationFactory(factory).currentTraceContext(currentTraceContext)
			.traceId128Bit(sleuthProperties.isTraceId128())
			.supportsJoin(sleuthProperties.isSupportsJoin());
	if (spanHandlers != null) {
		for (SpanHandler spanHandlerFactory : spanHandlers) {
			builder.addSpanHandler(spanHandlerFactory);
		}
	}
	if (tracingCustomizers != null) {
		for (TracingCustomizer customizer : tracingCustomizers) {
			customizer.customize(builder);
		}
	}

	return builder.build();
}
 

@Test public void matches() {
  Map<Sampler, Boolean> samplerToAnswer = new LinkedHashMap<>();
  samplerToAnswer.put(Sampler.ALWAYS_SAMPLE, true);
  samplerToAnswer.put(Sampler.NEVER_SAMPLE, false);

  samplerToAnswer.forEach((sampler, answer) -> {
    this.sampler = MessagingRuleSampler.newBuilder()
      .putRule(operationEquals("receive"), sampler)
      .build();

    when(request.operation()).thenReturn("receive");

    assertThat(this.sampler.trySample(request))
      .isEqualTo(answer);

    // consistent answer
    assertThat(this.sampler.trySample(request))
      .isEqualTo(answer);
  });
}
 

protected ITRemote() {
  CurrentTraceContext.Builder currentTraceContextBuilder = currentTraceContextBuilder();
  if (currentTraceContextBuilder instanceof StrictCurrentTraceContext.Builder) {
    currentTraceContext = currentTraceContextBuilder.build();
    checkForLeakedScopes = (Closeable) currentTraceContext;
  } else {
    StrictScopeDecorator strictScopeDecorator = StrictScopeDecorator.create();
    currentTraceContext = currentTraceContextBuilder
      .addScopeDecorator(strictScopeDecorator).build();
    checkForLeakedScopes = strictScopeDecorator;
  }
  propagationFactory = BaggagePropagation.newFactoryBuilder(B3Propagation.FACTORY)
    .add(SingleBaggageField.newBuilder(BAGGAGE_FIELD)
      .addKeyName(BAGGAGE_FIELD_KEY)
      .build()).build();
  tracing = tracingBuilder(Sampler.ALWAYS_SAMPLE).build();
}
 

Tracer(
  Clock clock,
  Propagation.Factory propagationFactory,
  SpanHandler spanHandler,
  PendingSpans pendingSpans,
  Sampler sampler,
  CurrentTraceContext currentTraceContext,
  boolean traceId128Bit,
  boolean supportsJoin,
  boolean alwaysSampleLocal,
  AtomicBoolean noop
) {
  this.clock = clock;
  this.propagationFactory = propagationFactory;
  this.spanHandler = spanHandler;
  this.pendingSpans = pendingSpans;
  this.sampler = sampler;
  this.currentTraceContext = currentTraceContext;
  this.traceId128Bit = traceId128Bit;
  this.supportsJoin = supportsJoin;
  this.alwaysSampleLocal = alwaysSampleLocal;
  this.noop = noop;
}
 

/**
 * @since 4.19
 * @deprecated Since 5.8, use {@link #nextSpan(SamplerFunction, Object)}  or {@link
 * #startScopedSpan(String, SamplerFunction, Object)}
 */
@Deprecated public Tracer withSampler(Sampler sampler) {
  if (sampler == null) throw new NullPointerException("sampler == null");
  return new Tracer(
    clock,
    propagationFactory,
    spanHandler,
    pendingSpans,
    sampler,
    currentTraceContext,
    traceId128Bit,
    supportsJoin,
    alwaysSampleLocal,
    noop
  );
}
 

@Test public void customSampler() throws JMSException {
  MessagingRuleSampler producerSampler = MessagingRuleSampler.newBuilder()
    .putRule(channelNameEquals(jms.queue.getQueueName()), Sampler.NEVER_SAMPLE)
    .build();

  try (MessagingTracing messagingTracing = MessagingTracing.newBuilder(tracing)
    .producerSampler(producerSampler)
    .build();
       JMSContext context = JmsTracing.create(messagingTracing)
         .connectionFactory(((ArtemisJmsTestRule) jms).factory)
         .createContext(JMSContext.AUTO_ACKNOWLEDGE)
  ) {
    context.createProducer().send(jms.queue, "foo");
  }

  Message received = queueReceiver.receive();

  assertThat(propertiesToMap(received)).containsKey("b3")
    // Check that the injected context was not sampled
    .satisfies(m -> assertThat(m.get("b3")).endsWith("-0"));

  // @After will also check that the producer was not sampled
}
 

@Override
public void run(TranslationServiceConfiguration configuration,
		Environment environment) {

	/* START TRACING INSTRUMENTATION */
	final var sender = URLConnectionSender.newBuilder()
			.endpoint(configuration.getZipkinEndpoint()).build();
	final var reporter = AsyncReporter.builder(sender).build();
	final var tracing = Tracing.newBuilder().localServiceName("translation-service")
			.sampler(Sampler.ALWAYS_SAMPLE).spanReporter(reporter).build();
	final var httpTracing = HttpTracing.newBuilder(tracing).build();
	final var jerseyTracingFilter = TracingApplicationEventListener
			.create(httpTracing);
	environment.jersey().register(jerseyTracingFilter);
	/* END TRACING INSTRUMENTATION */

	final var repository = new TranslationRepository();
	final var translationResource = new TranslationResource(repository);
	environment.jersey().register(translationResource);

	final var healthCheck = new TranslationServiceHealthCheck();
	environment.healthChecks().register("translation-service", healthCheck);
}
 

@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
}
 

@Test public void spanHandler_recordsWhenUnsampledIfContextSamplesLocal() {
  AtomicBoolean sampledLocal = new AtomicBoolean();
  try (Tracing tracing = Tracing.newBuilder()
    .propagationFactory(new Propagation.Factory() {
      @Deprecated public <K> Propagation<K> create(Propagation.KeyFactory<K> keyFactory) {
        return B3SinglePropagation.FACTORY.create(keyFactory);
      }

      @Override public TraceContext decorate(TraceContext context) {
        if (sampledLocal.getAndSet(true)) return context;
        return context.toBuilder().sampledLocal(true).build();
      }
    })
    .addSpanHandler(spans)
    .sampler(Sampler.NEVER_SAMPLE)
    .build()) {
    tracing.tracer().newTrace().start().name("one").finish();
    tracing.tracer().newTrace().start().name("two").finish();
  }

  assertThat(spans).hasSize(1);
  assertThat(spans.get(0).name()).isEqualTo("one");
}
 

@Test public void receive_customSampler() throws JMSException {
  queueReceiver.close();

  MessagingRuleSampler consumerSampler = MessagingRuleSampler.newBuilder()
    .putRule(channelNameEquals(jms.queue.getQueueName()), Sampler.NEVER_SAMPLE)
    .build();

  try (MessagingTracing messagingTracing = MessagingTracing.newBuilder(tracing)
    .consumerSampler(consumerSampler)
    .build();
       JMSContext context = JmsTracing.create(messagingTracing)
         .connectionFactory(((ArtemisJmsTestRule) jms).factory)
         .createContext(JMSContext.AUTO_ACKNOWLEDGE);
       JMSConsumer consumer = context.createConsumer(jms.queue)
  ) {
    queueSender.send(message);

    // Check that the message headers are not sampled
    assertThat(consumer.receive().getStringProperty("b3"))
      .endsWith("-0");
  }

  // @After will also check that the consumer was not sampled
}
 

@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();
}
 

@Test public void spanHandler_doesntRecordWhenUnsampled() {
  try (Tracing tracing = Tracing.newBuilder()
    .addSpanHandler(spans)
    .sampler(Sampler.NEVER_SAMPLE)
    .build()) {
    tracing.tracer().newTrace().start().name("aloha").finish();
  }

  assertThat(spans).isEmpty();
}
 

/**
 * This ensures custom span handlers can see the actual exception thrown, not just the "error"
 * tag value.
 */
void spanHandlerSeesError(Callable<Void> get) throws IOException {
  ConcurrentLinkedDeque<Throwable> caughtThrowables = new ConcurrentLinkedDeque<>();
  closeClient(client);
  httpTracing = HttpTracing.create(tracingBuilder(Sampler.ALWAYS_SAMPLE)
    .clearSpanHandlers()
    .addSpanHandler(new SpanHandler() {
      @Override public boolean end(TraceContext context, MutableSpan span, Cause cause) {
        Throwable error = span.error();
        if (error != null) {
          caughtThrowables.add(error);
        } else {
          caughtThrowables.add(new RuntimeException("Unexpected additional call to end"));
        }
        return true;
      }
    })
    // The blocking span handler goes after the error catcher, so we can assert on the errors.
    .addSpanHandler(testSpanHandler)
    .build());
  client = newClient(server.getPort());

  // If this passes, a span was reported with an error
  checkReportsSpanOnTransportException(get);

  assertThat(caughtThrowables)
      .withFailMessage("Span finished with error, but caughtThrowables empty")
      .isNotEmpty();
  if (caughtThrowables.size() > 1) {
    for (Throwable throwable : caughtThrowables) {
      Logger.getAnonymousLogger().log(Level.SEVERE, "multiple calls to finish", throwable);
    }
    assertThat(caughtThrowables).hasSize(1);
  }
}
 

@Test
void shouldSubmitSpanWithCustomRemoteName() throws Exception {
    final SpanCollector collector = new SpanCollector();

    final Tracing tracing = Tracing.newBuilder()
                                   .localServiceName(TEST_SERVICE)
                                   .addSpanHandler(collector)
                                   .sampler(Sampler.create(1.0f))
                                   .build();
    testRemoteInvocation(tracing, "fooService");

    // check span name
    final MutableSpan span = collector.spans().poll(10, TimeUnit.SECONDS);

    // check tags
    assertThat(span).isNotNull();
    assertThat(span.tags()).containsEntry("http.host", "foo.com")
                           .containsEntry("http.method", "POST")
                           .containsEntry("http.path", "/hello/armeria")
                           .containsEntry("http.url", "http://foo.com/hello/armeria")
                           .containsEntry("http.protocol", "h2c");

    // check service name
    assertThat(span.localServiceName()).isEqualTo(TEST_SERVICE);

    // check remote service name
    assertThat(span.remoteServiceName()).isEqualTo("fooService");
}
 

@Bean
@RefreshScope
@ConditionalOnMissingBean
public Sampler defaultTraceSampler(SamplerProperties config) {
	// TODO: Rewrite: refresh should replace the sampler, not change its state
	// internally
	if (config.getProbability() != null) {
		return new ProbabilityBasedSampler(config);
	}
	return new RateLimitingSampler(config);
}
 

Tracing.Builder tracingBuilder(Sampler sampler) {
  return Tracing.newBuilder()
      .spanReporter(
          s -> {
            // make sure the context was cleared prior to finish.. no leaks!
            TraceContext current = httpTracing.tracing().currentTraceContext().get();
            if (current != null) {
              Assert.assertNotEquals(current.spanId(), s.id());
            }
            spans.add(s);
          })
      .currentTraceContext(currentTraceContext)
      .sampler(sampler);
}
 

protected Tracing tracing() {
    this.tracing= Tracing
            .newBuilder()
            .localServiceName(this.serviceName)
            .sampler(Sampler.ALWAYS_SAMPLE)
            .spanReporter(spanReporter())
            .build();
    return this.tracing;
}
 

@Before public void clear() {
  tracing.set(Tracing.newBuilder()
    .currentTraceContext(currentTraceContext)
    .addSpanHandler(spans)
    .sampler(new Sampler() {
      @Override public boolean isSampled(long traceId) {
        throw new AssertionError(); // in this case, we aren't expecting a fallback
      }
    }).build());
  spans.clear();
}
 

@Test public void withNoopSpanInScope() {
  Span current = tracer.withSampler(Sampler.NEVER_SAMPLE).nextSpan();

  try (SpanInScope ws = tracer.withSpanInScope(current)) {
    assertThat(tracer.currentSpan())
      .isEqualTo(current);
    assertThat(tracer.currentSpanCustomizer())
      .isNotEqualTo(current)
      .isEqualTo(NoopSpanCustomizer.INSTANCE);
  }

  // context was cleared
  assertThat(tracer.currentSpan()).isNull();
}
 

void spanHandlerSeesError(String path) throws IOException {
  ConcurrentLinkedDeque<Throwable> caughtThrowables = new ConcurrentLinkedDeque<>();
  httpTracing = HttpTracing.create(tracingBuilder(Sampler.ALWAYS_SAMPLE)
      .clearSpanHandlers()
      .addSpanHandler(new SpanHandler() {
        @Override public boolean end(TraceContext context, MutableSpan span, Cause cause) {
          Throwable error = span.error();
          if (error != null) {
            caughtThrowables.add(error);
          } else {
            caughtThrowables.add(new RuntimeException("Unexpected additional call to end"));
          }
          return true;
        }
      })
      // The blocking span handler goes after the error catcher, so we can assert on the errors.
      .addSpanHandler(testSpanHandler)
      .build());
  init();

  // If this passes, a span was reported with an error
  httpStatusCodeTagMatchesResponse_onUncaughtException(path, ".*not ready");

  assertThat(caughtThrowables)
      .withFailMessage("Span finished with error, but caughtThrowables empty")
      .isNotEmpty();
  if (caughtThrowables.size() > 1) {
    for (Throwable throwable : caughtThrowables) {
      Logger.getAnonymousLogger().log(Level.SEVERE, "multiple calls to finish", throwable);
    }
    assertThat(caughtThrowables).hasSize(1);
  }
}
 

@Override public Set<Object> getSingletons() {
  return new LinkedHashSet<>(asList(new Resource(), TracingApplicationEventListener.create(
    HttpTracing.create(Tracing.newBuilder()
      .sampler(Sampler.NEVER_SAMPLE)
      .spanReporter(Reporter.NOOP)
      .build())
  )));
}
 

Tracing buildZipkinTracing(@NonNull String name, @NonNull String zipkinUrl, float samplingRate) {
  if (zipkinUrl.isEmpty() || !(samplingRate > 0f)) {
    return null;
  }
  return Tracing.newBuilder()
      .localServiceName(name)
      // puts trace IDs into logs
      .currentTraceContext(MDCCurrentTraceContext.create())
      .spanReporter(buildReporter(zipkinUrl))
      .sampler(Sampler.create(samplingRate))
      .build();
}
 

@Test public void matches() {
  Map<Sampler, Boolean> samplerToAnswer = new LinkedHashMap<>();
  samplerToAnswer.put(Sampler.ALWAYS_SAMPLE, true);
  samplerToAnswer.put(Sampler.NEVER_SAMPLE, false);

  samplerToAnswer.forEach((sampler, answer) -> {
    HttpRuleSampler ruleSampler = HttpRuleSampler.newBuilder()
      .putRule(pathStartsWith("/foo"), sampler)
      .build();

    when(adapter.path(request)).thenReturn("/foo");

    assertThat(ruleSampler.trySample(adapter, request))
      .isEqualTo(answer);

    when(httpClientRequest.path()).thenReturn("/foo");

    assertThat(ruleSampler.trySample(httpClientRequest))
      .isEqualTo(answer);

    when(httpServerRequest.path()).thenReturn("/foo");

    // consistent answer
    assertThat(ruleSampler.trySample(httpServerRequest))
      .isEqualTo(answer);
  });
}
 

@Test public void consumerSampler() {
  consumerSampler = MessagingRuleSampler.newBuilder()
    .putRule(channelNameEquals(TEST_QUEUE), Sampler.NEVER_SAMPLE)
    .build();

  produceUntracedMessage();
  awaitMessageConsumed();
  // reporter rules verify nothing was reported
}
 

public Unsampled() {
  super(Tracing.newBuilder()
    .sampler(Sampler.NEVER_SAMPLE)
    .addSpanHandler(new SpanHandler() {
      // intentionally not NOOP to ensure spans report
    })
    .build());
}
 

@Test public void putAllRules() {
  MessagingRuleSampler base = MessagingRuleSampler.newBuilder()
    .putRule(operationEquals("receive"), Sampler.NEVER_SAMPLE)
    .build();

  sampler = MessagingRuleSampler.newBuilder()
    .putAllRules(base)
    .build();

  when(request.operation()).thenReturn("send");

  assertThat(sampler.trySample(request))
    .isNull();
}
 

@Test public void currentSpanCustomizer_noop_when_notSampled() {
  ScopedSpan parent = tracer.withSampler(Sampler.NEVER_SAMPLE).startScopedSpan("parent");
  try {
    assertThat(tracer.currentSpanCustomizer())
      .isSameAs(NoopSpanCustomizer.INSTANCE);
  } finally {
    parent.finish();
  }
}
 

@Test
void samplerFromProps_rateLimitZero() {
	SamplerProperties properties = new SamplerProperties();
	properties.setRate(0);

	Sampler sampler = SamplerAutoConfiguration.samplerFromProps(properties);

	BDDAssertions.then(sampler).isSameAs(Sampler.NEVER_SAMPLE);
}
 

@Test
void samplerFromProps_prefersZeroProbability() {
	SamplerProperties properties = new SamplerProperties();
	properties.setProbability(0.0f);
	properties.setRate(20);

	Sampler sampler = SamplerAutoConfiguration.samplerFromProps(properties);

	BDDAssertions.then(sampler).isSameAs(Sampler.NEVER_SAMPLE);
}
 

protected Tracing.Builder tracingBuilder(Sampler sampler) {
  return Tracing.newBuilder()
    .localServiceName(getClass().getSimpleName())
    .localIp("127.0.0.1") // Prevent implicit lookups
    .addSpanHandler(testSpanHandler)
    .propagationFactory(propagationFactory)
    .currentTraceContext(currentTraceContext)
    .sampler(sampler);
}