brave.propagation.SamplingFlags Java Examples

/**
 * This tests that the parent is determined at the time the request was made, not when the request
 * was executed.
 */
@Test public void usesParentFromInvocationTime() {
  TraceContext parent = newTraceContext(SamplingFlags.SAMPLED);
  try (Scope scope = currentTraceContext.newScope(parent)) {
    RpcContext.getContext().asyncCall(() -> client.get().sayHello("jorge"));
    RpcContext.getContext().asyncCall(() -> client.get().sayHello("romeo"));
  }

  try (Scope scope = currentTraceContext.newScope(null)) {
    // complete within a different scope
    for (int i = 0; i < 2; i++) {
      TraceContext extracted = server.takeRequest().context();
      assertChildOf(extracted, parent);
    }
  }

  // The spans may report in a different order than the requests
  for (int i = 0; i < 2; i++) {
    assertChildOf(testSpanHandler.takeRemoteSpan(CLIENT), parent);
  }
}
 

/** Returns a trace context for use in propagation tests. */
protected TraceContext newTraceContext(SamplingFlags flags) {
  long id = System.nanoTime(); // Random enough as tests are run serially anyway

  // Simulate a new local root root, but without the dependency on Tracer to create it.
  TraceContext context = InternalPropagation.instance.newTraceContext(
    InternalPropagation.instance.flags(flags) | FLAG_LOCAL_ROOT,
    0L,
    id + 1L,
    id + 3L,
    id + 2L,
    id + 3L,
    Collections.emptyList()
  );

  return propagationFactory.decorate(context);
}
 

/**
 * This tests that the parent is determined at the time the request was made, not when the request
 * was executed.
 */
@Test public void usesParentFromInvocationTime() {
  server.enqueueDelay(TimeUnit.SECONDS.toMillis(1));
  GreeterGrpc.GreeterFutureStub futureStub = GreeterGrpc.newFutureStub(client);

  TraceContext parent = newTraceContext(SamplingFlags.SAMPLED);
  try (Scope scope = currentTraceContext.newScope(parent)) {
    futureStub.sayHello(HELLO_REQUEST);
    futureStub.sayHello(HELLO_REQUEST);
  }

  try (Scope scope = currentTraceContext.newScope(null)) {
    for (int i = 0; i < 2; i++) {
      TraceContext extracted = server.takeRequest().context();
      assertChildOf(extracted, parent);
    }
  }

  // The spans may report in a different order than the requests
  for (int i = 0; i < 2; i++) {
    assertChildOf(testSpanHandler.takeRemoteSpan(CLIENT), parent);
  }
}
 

/**
 * This ensures that response callbacks run in the invocation context, not the client one. This
 * allows async chaining to appear caused by the parent, not by the most recent client. Otherwise,
 * we would see a client span child of a client span, which could be confused with duplicate
 * instrumentation and affect dependency link counts.
 */
@Test public void callbackContextIsFromInvocationTime() {
  AssertableCallback<HelloReply> callback = new AssertableCallback<>();

  // Capture the current trace context when onSuccess or onError occur
  AtomicReference<TraceContext> invocationContext = new AtomicReference<>();
  callback.setListener(() -> invocationContext.set(currentTraceContext.get()));

  TraceContext parent = newTraceContext(SamplingFlags.SAMPLED);
  try (Scope scope = currentTraceContext.newScope(parent)) {
    GreeterGrpc.newStub(client).sayHello(HELLO_REQUEST, new StreamObserverAdapter(callback));
  }

  callback.join(); // ensures listener ran
  assertThat(invocationContext.get()).isSameAs(parent);
  assertChildOf(testSpanHandler.takeRemoteSpan(CLIENT), parent);
}
 

@ParameterizedTest
@MethodSource("parametersForShouldPropagateNotSamplingHeader")
@SuppressWarnings("unchecked")
public void shouldPropagateNotSamplingHeader(ResponseEntityProvider provider) {
	Span span = this.tracer
			.nextSpan(TraceContextOrSamplingFlags.create(SamplingFlags.NOT_SAMPLED))
			.name("foo").start();

	try (Tracer.SpanInScope ws = this.tracer.withSpanInScope(span)) {
		ResponseEntity<Map<String, String>> response = provider.get(this);

		then(response.getBody().get("b3")).isNotNull().endsWith("-0"); // not sampled
	}
	finally {
		span.finish();
	}

	then(this.spans).isEmpty();
	then(this.tracer.currentSpan()).isNull();
}
 

@Test public void should_prefer_current_to_stale_b3_header() throws JMSException {
  jms.setReadOnlyProperties(message, false);
  setStringProperty(message, "b3", writeB3SingleFormat(newTraceContext(SamplingFlags.NOT_SAMPLED)));

  TraceContext parent = newTraceContext(SamplingFlags.SAMPLED);
  try (Scope scope = currentTraceContext.newScope(parent)) {
    messageProducer.send(jms.destination, message);
  }

  Message received = messageConsumer.receive();

  MutableSpan producerSpan = testSpanHandler.takeRemoteSpan(PRODUCER);
  assertChildOf(producerSpan, parent);

  assertThat(propertiesToMap(received))
    .containsAllEntriesOf(existingProperties)
    .containsEntry("b3", producerSpan.traceId() + "-" + producerSpan.id() + "-1");
}
 

void messageListener_resumesTrace(Runnable send) {
  consumer.setMessageListener(m -> {
    // clearing headers ensures later work doesn't try to use the old parent
    String b3 = getPropertyIfString(m, "b3");
    tracing.tracer().currentSpanCustomizer().tag("b3", String.valueOf(b3 != null));
  });

  TraceContext parent = newTraceContext(SamplingFlags.SAMPLED);
  producer.setProperty("b3", parent.traceIdString() + "-" + parent.spanIdString() + "-1");
  send.run();

  MutableSpan consumerSpan = testSpanHandler.takeRemoteSpan(CONSUMER);
  MutableSpan listenerSpan = testSpanHandler.takeLocalSpan();

  assertChildOf(consumerSpan, parent);
  assertChildOf(listenerSpan, consumerSpan);

  assertThat(listenerSpan.tags())
    .hasSize(1) // no redundant copy of consumer tags
    .containsEntry("b3", "false"); // b3 header not leaked to listener
}
 

/**
 * This ensures that response callbacks run in the invocation context, not the client one. This
 * allows async chaining to appear caused by the parent, not by the most recent client. Otherwise,
 * we would see a client span child of a client span, which could be confused with duplicate
 * instrumentation and affect dependency link counts.
 */
@Test public void callbackContextIsFromInvocationTime() {
  server.enqueue(new MockResponse());

  AssertableCallback<Integer> callback = new AssertableCallback<>();

  // Capture the current trace context when onSuccess or onError occur
  AtomicReference<TraceContext> invocationContext = new AtomicReference<>();
  callback.setListener(() -> invocationContext.set(currentTraceContext.get()));

  TraceContext parent = newTraceContext(SamplingFlags.SAMPLED);
  try (Scope scope = currentTraceContext.newScope(parent)) {
    get(client, "/foo", callback);
  }

  callback.join(); // ensures listener ran
  assertThat(invocationContext.get()).isSameAs(parent);
  assertChildOf(testSpanHandler.takeRemoteSpan(CLIENT), parent);
}
 

/**
 * Handle when the client doesn't actually make a client span
 *
 * <p>See https://github.com/openzipkin/brave/issues/864
 */
@Test public void cacheControl() throws IOException {
  server.enqueue(new MockResponse()
    .addHeader("Content-Type", "text/plain")
    .addHeader("Cache-Control", "max-age=600, stale-while-revalidate=1200")
    .setBody("Hello"));

  TraceContext parent = newTraceContext(SamplingFlags.SAMPLED);
  try (Scope scope = currentTraceContext.newScope(parent)) {
    get(client, "/cached");
    get(client, "/cached");
  }

  assertThat(server.getRequestCount()).isEqualTo(1);

  MutableSpan real = testSpanHandler.takeRemoteSpan(CLIENT);
  MutableSpan cached = testSpanHandler.takeLocalSpan();
  assertThat(cached.tags()).containsKey("http.cache_hit");

  for (MutableSpan child : Arrays.asList(real, cached)) {
    assertChildOf(child, parent);
  }

  assertSequential(real, cached);
}
 

@Test public void should_prefer_current_to_stale_b3_header() {
  producer.setProperty("b3", writeB3SingleFormat(newTraceContext(SamplingFlags.NOT_SAMPLED)));

  TraceContext parent = newTraceContext(SamplingFlags.SAMPLED);
  try (Scope scope = currentTraceContext.newScope(parent)) {
    producer.send(jms.queue, "foo");
  }

  Message received = consumer.receive();

  MutableSpan producerSpan = testSpanHandler.takeRemoteSpan(PRODUCER);
  assertChildOf(producerSpan, parent);

  assertThat(propertiesToMap(received))
    .containsAllEntriesOf(existingProperties)
    .containsEntry("b3", producerSpan.traceId() + "-" + producerSpan.id() + "-1");
}
 

TraceContext contextWithPassThrough() {
  extractor = BaggagePropagation.newFactoryBuilder(B3Propagation.FACTORY)
      .add(BaggagePropagationConfig.SingleBaggageField.remote(AWSPropagation.FIELD_AMZN_TRACE_ID))
      .build()
      .get()
      .extractor(Map::get);

  TraceContextOrSamplingFlags extracted = extractor.extract(carrier);

  // sanity check
  assertThat(extracted.samplingFlags()).isEqualTo(SamplingFlags.EMPTY);
  assertThat(extracted.extra()).isNotEmpty();

  // Make a context that wasn't from AWSPropagation
  return TraceContext.newBuilder()
      .traceId(1L)
      .spanId(2L)
      .sampled(true)
      .addExtra(extracted.extra().get(0))
      .build();
}
 

@Test public void createsChildWhenJoinDisabled() throws IOException {
  tracing = tracingBuilder(NEVER_SAMPLE).supportsJoin(false).build();
  httpTracing = HttpTracing.create(tracing);
  init();

  String path = "/foo";

  TraceContext parent = newTraceContext(SamplingFlags.SAMPLED);
  get(new Request.Builder().url(url(path))
    .header("b3", B3SingleFormat.writeB3SingleFormat(parent))
    .build());

  MutableSpan span = testSpanHandler.takeRemoteSpan(SERVER);
  assertChildOf(span, parent);
  assertThat(span.id()).isNotEqualTo(parent.spanIdString());
}
 

/** This prevents confusion as a blocking client should end before, the start of the next span. */
@Test public void clientTimestampAndDurationEnclosedByParent() throws IOException {
  server.enqueue(new MockResponse());

  TraceContext parent = newTraceContext(SamplingFlags.SAMPLED);
  Clock clock = tracing.clock(parent);

  long start = clock.currentTimeMicroseconds();
  try (Scope scope = currentTraceContext.newScope(parent)) {
    get(client, "/foo");
  }
  long finish = clock.currentTimeMicroseconds();

  MutableSpan clientSpan = testSpanHandler.takeRemoteSpan(CLIENT);
  assertChildOf(clientSpan, parent);
  assertSpanInInterval(clientSpan, start, finish);
}
 

@Test public void redirect() throws IOException {
  server.enqueue(new MockResponse().setResponseCode(302)
    .addHeader("Location: " + url("/bar")));
  server.enqueue(new MockResponse().setResponseCode(404)); // hehe to a bad location!

  TraceContext parent = newTraceContext(SamplingFlags.SAMPLED);
  try (Scope scope = currentTraceContext.newScope(parent)) {
    get(client, "/foo");
  } catch (RuntimeException e) {
    // some clients raise 404 as an exception such as HttpClientError
  }

  MutableSpan initial = testSpanHandler.takeRemoteSpan(CLIENT);
  MutableSpan redirected = testSpanHandler.takeRemoteSpanWithErrorTag(CLIENT, "404");

  for (MutableSpan child : Arrays.asList(initial, redirected)) {
    assertChildOf(child, parent);
  }

  assertSequential(initial, redirected);

  assertThat(initial.tags().get("http.path")).isEqualTo("/foo");
  assertThat(redirected.tags().get("http.path")).isEqualTo("/bar");
}
 

@Test public void reusesPropagatedSpanId() {
  TraceContext parent = newTraceContext(SamplingFlags.SAMPLED);
  Channel channel = clientWithB3SingleHeader(parent);
  GreeterGrpc.newBlockingStub(channel).sayHello(HELLO_REQUEST);

  assertSameIds(testSpanHandler.takeRemoteSpan(Span.Kind.SERVER), parent);
}
 

@Override public <R> TraceContext.Extractor<R> extractor(Getter<R, String> getter) {
  return request -> {
    String result = getter.get(request, TRACE_ID);
    if (result == null) return TraceContextOrSamplingFlags.create(SamplingFlags.EMPTY);
    return TraceContextOrSamplingFlags.create(TraceIdContext.newBuilder()
      .traceId(HexCodec.lowerHexToUnsignedLong(result))
      .build());
  };
}
 

@Test public void should_not_serialize_parent_span_id() throws JMSException {
  TraceContext parent = newTraceContext(SamplingFlags.SAMPLED);
  try (Scope scope = currentTraceContext.newScope(parent)) {
    messageProducer.send(jms.destination, message);
  }

  Message received = messageConsumer.receive();

  MutableSpan producerSpan = testSpanHandler.takeRemoteSpan(PRODUCER);
  assertChildOf(producerSpan, parent);

  assertThat(propertiesToMap(received))
    .containsAllEntriesOf(existingProperties)
    .containsEntry("b3", producerSpan.traceId() + "-" + producerSpan.id() + "-1");
}
 

/**
 * This tests that the parent is determined at the time the request was made, not when the request
 * was executed.
 */
@Test public void usesParentFromInvocationTime() {
  server.enqueue(new MockResponse().setBodyDelay(300, TimeUnit.MILLISECONDS));
  server.enqueue(new MockResponse());

  AssertableCallback<Integer> items1 = new AssertableCallback<>();
  AssertableCallback<Integer> items2 = new AssertableCallback<>();

  TraceContext parent = newTraceContext(SamplingFlags.SAMPLED);
  try (Scope scope = currentTraceContext.newScope(parent)) {
    get(client, "/items/1", items1);
    get(client, "/items/2", items2);
  }

  try (Scope scope = currentTraceContext.newScope(null)) {
    // complete within a different scope
    items1.join();
    items2.join();

    for (int i = 0; i < 2; i++) {
      TraceContext extracted = extract(takeRequest());
      assertChildOf(extracted, parent);
    }
  }

  // The spans may report in a different order than the requests
  for (int i = 0; i < 2; i++) {
    assertChildOf(testSpanHandler.takeRemoteSpan(CLIENT), parent);
  }
}
 

/** This prevents confusion as a blocking client should end before, the start of the next span. */
@Test public void clientTimestampAndDurationEnclosedByParent() {
  TraceContext parent = newTraceContext(SamplingFlags.SAMPLED);
  Clock clock = tracing.clock(parent);

  long start = clock.currentTimeMicroseconds();
  try (Scope scope = currentTraceContext.newScope(parent)) {
    GreeterGrpc.newBlockingStub(client).sayHello(HELLO_REQUEST);
  }
  long finish = clock.currentTimeMicroseconds();

  MutableSpan clientSpan = testSpanHandler.takeRemoteSpan(CLIENT);
  assertChildOf(clientSpan, parent);
  assertSpanInInterval(clientSpan, start, finish);
}
 

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

  TraceContext parent = newTraceContext(SamplingFlags.SAMPLED);
  get(new Request.Builder().url(url(path))
    .header("b3", B3SingleFormat.writeB3SingleFormat(parent))
    .build());

  assertSameIds(testSpanHandler.takeRemoteSpan(SERVER), parent);
}
 

/**
 * Returns span from TraceContext or dubbo attachments.
 */
public Span nextSpan(TraceContextOrSamplingFlags extracted) {
  TraceContext parent = currentTraceContext.get();
  //If spanInScope is closed we can use dubbo attachments
  if (parent == null) {
    parent = extracted.context();
  }
  if (parent != null) {
    return tracer.newChild(parent);
  }
  return tracer.newTrace(SamplingFlags.NOT_SAMPLED);
}
 

@Test public void propagatesBaggage_unsampled() throws IOException {
  server.enqueue(new MockResponse());

  TraceContext parent = newTraceContext(SamplingFlags.NOT_SAMPLED);
  try (Scope scope = currentTraceContext.newScope(parent)) {
    BAGGAGE_FIELD.updateValue(parent, "joey");
    get(client, "/foo");
  }

  TraceContext extracted = extract(takeRequest());
  assertThat(BAGGAGE_FIELD.getValue(extracted)).isEqualTo("joey");
}
 

@Test public void makesChildOfCurrentSpan() {
  TraceContext parent = newTraceContext(SamplingFlags.SAMPLED);
  try (Scope scope = currentTraceContext.newScope(parent)) {
    executeFind(COLLECTION_NAME);
  }

  MutableSpan clientSpan = testSpanHandler.takeRemoteSpan(CLIENT);
  assertChildOf(clientSpan, parent);
}
 

@Test public void propagatesBaggage() throws IOException {
  server.enqueue(new MockResponse());

  TraceContext parent = newTraceContext(SamplingFlags.SAMPLED);
  try (Scope scope = currentTraceContext.newScope(parent)) {
    BAGGAGE_FIELD.updateValue(parent, "joey");
    get(client, "/foo");
  }

  TraceContext extracted = extract(takeRequest());
  assertThat(BAGGAGE_FIELD.getValue(extracted)).isEqualTo("joey");

  testSpanHandler.takeRemoteSpan(CLIENT);
}
 

/** Unlike Brave 3, Brave 4 propagates trace ids even when unsampled */
@Test public void propagatesUnsampledContext() throws IOException {
  server.enqueue(new MockResponse());

  TraceContext parent = newTraceContext(SamplingFlags.NOT_SAMPLED);
  try (Scope scope = currentTraceContext.newScope(parent)) {
    get(client, "/foo");
  }

  TraceContext extracted = extract(takeRequest());
  assertThat(extracted.sampled()).isFalse();
  assertChildOf(extracted, parent);
}
 

@Test public void propagatesChildOfCurrentSpan() throws IOException {
  server.enqueue(new MockResponse());

  TraceContext parent = newTraceContext(SamplingFlags.SAMPLED);
  try (Scope scope = currentTraceContext.newScope(parent)) {
    get(client, "/foo");
  }

  TraceContext extracted = extract(takeRequest());
  assertThat(extracted.sampled()).isTrue();
  assertChildOf(extracted, parent);
  assertSameIds(testSpanHandler.takeRemoteSpan(CLIENT), extracted);
}
 

@Test public void createsChildWhenJoinDisabled() {
  tracing = tracingBuilder(NEVER_SAMPLE).supportsJoin(false).build();
  init();

  TraceContext parent = newTraceContext(SamplingFlags.SAMPLED);

  RpcContext.getContext().getAttachments().put("b3", B3SingleFormat.writeB3SingleFormat(parent));
  client.get().sayHello("jorge");

  MutableSpan span = testSpanHandler.takeRemoteSpan(SERVER);
  assertChildOf(span, parent);
  assertThat(span.id()).isNotEqualTo(parent.spanIdString());
}
 

@Test public void reusesPropagatedSpanId() {
  TraceContext parent = newTraceContext(SamplingFlags.SAMPLED);

  RpcContext.getContext().getAttachments().put("b3", B3SingleFormat.writeB3SingleFormat(parent));
  client.get().sayHello("jorge");

  assertSameIds(testSpanHandler.takeRemoteSpan(SERVER), parent);
}
 

/**
 * This tests that the parent is determined at the time the request was made, not when the request
 * was executed.
 */
@Test public void usesParentFromInvocationTime() {
  AssertableCallback<String> items1 = new AssertableCallback<>();
  AssertableCallback<String> items2 = new AssertableCallback<>();

  TraceContext parent = newTraceContext(SamplingFlags.SAMPLED);
  try (Scope scope = currentTraceContext.newScope(parent)) {
    RpcContext.getContext().asyncCall(() -> client.get().sayHello("jorge"))
        .whenComplete(items1);
    RpcContext.getContext().asyncCall(() -> client.get().sayHello("romeo"))
        .whenComplete(items2);
  }

  try (Scope scope = currentTraceContext.newScope(null)) {
    // complete within a different scope
    items1.join();
    items2.join();

    for (int i = 0; i < 2; i++) {
      TraceContext extracted = server.takeRequest().context();
      assertChildOf(extracted, parent);
    }
  }

  // The spans may report in a different order than the requests
  for (int i = 0; i < 2; i++) {
    assertChildOf(testSpanHandler.takeRemoteSpan(CLIENT), parent);
  }
}
 

@Test public void handleReceive_defaultsToMakeNewTrace() {
  when(extractor.extract(request))
    .thenReturn(TraceContextOrSamplingFlags.create(SamplingFlags.EMPTY));

  // request sampler abstains (trace ID sampler will say true)
  when(sampler.trySample(any(FromRequestAdapter.class))).thenReturn(null);

  Span newSpan = handler.handleReceive(extractor, request);
  assertThat(newSpan.isNoop()).isFalse();
  assertThat(newSpan.context().shared()).isFalse();
}