brave.propagation.TraceContextOrSamplingFlags Java Examples

@Override public ProducerRecord<K, V> onSend(ProducerRecord<K, V> record) {
  TraceContextOrSamplingFlags traceContextOrSamplingFlags = extractor.extract(record.headers());
  Span span = tracing.tracer().nextSpan(traceContextOrSamplingFlags);
  tracing.propagation().keys().forEach(key -> record.headers().remove(key));
  injector.inject(span.context(), record.headers());
  if (!span.isNoop()) {
    if (record.key() instanceof String && !"".equals(record.key())) {
      span.tag(KafkaInterceptorTagKey.KAFKA_KEY, record.key().toString());
    }
    span
      .tag(KafkaInterceptorTagKey.KAFKA_TOPIC, record.topic())
      .tag(KafkaInterceptorTagKey.KAFKA_CLIENT_ID,
        configuration.getString(ProducerConfig.CLIENT_ID_CONFIG))
      .name(SPAN_NAME)
      .kind(Span.Kind.PRODUCER)
      .remoteServiceName(remoteServiceName)
      .start()
      .finish();
  }
  return record;
}
 

@Test public void reconfigure() {
  tracing.close();
  zipkinSpanHandler.close();

  zipkinSpanHandler = (H) zipkinSpanHandler.toBuilder().alwaysReportSpans(true).build();
  tracing = Tracing.newBuilder()
      .localServiceName("Aa")
      .localIp("1.2.3.4")
      .localPort(80)
      .addSpanHandler(zipkinSpanHandler)
      .alwaysSampleLocal()
      .build();

  brave.Span unsampledRemote =
      tracing.tracer().nextSpan(TraceContextOrSamplingFlags.NOT_SAMPLED).name("test").start(1L);
  assertThat(unsampledRemote.isNoop()).isFalse();
  assertThat(unsampledRemote.context().sampled()).isFalse();
  assertThat(unsampledRemote.context().sampledLocal()).isTrue();
  unsampledRemote.finish(2L);

  triggerReport();

  // alwaysReportSpans applied
  assertThat(spans).isNotEmpty();
}
 

@Override public TraceContextOrSamplingFlags extract(R request) {
  TraceContextOrSamplingFlags.Builder builder = delegate.extract(request).toBuilder();
  BaggageFields extra = factory.baggageFactory.create();
  builder.addExtra(extra);

  if (factory.extra == null) return builder.build();

  for (BaggagePropagationConfig config : factory.configs) {
    if (config.baggageCodec == BaggageCodec.NOOP) continue; // local field

    List<String> keys = config.baggageCodec.injectKeyNames();
    for (int i = 0, length = keys.size(); i < length; i++) {
      String value = getter.get(request, keys.get(i));
      if (value != null && config.baggageCodec.decode(extra, request, value)) {
        break; // accept the first match
      }
    }
  }

  return builder.addExtra(factory.extra).build();
}
 

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

@Override public TraceContextOrSamplingFlags extract(R request) {
  if (!(request instanceof GrpcServerRequest)) return delegate.extract(request);

  Metadata metadata = ((GrpcClientRequest) request).headers;

  // First, check if we are propagating gRPC tags.
  TagsBin tagsBin = metadata.get(GRPC_TAGS_BIN);

  // Next, check to see if there is a gRPC formatted trace context: use it if parsable.
  byte[] bytes = metadata.get(GRPC_TRACE_BIN);
  if (bytes != null) {
    TraceContext maybeContext = TraceContextBinaryFormat.parseBytes(bytes, tagsBin);
    if (maybeContext != null) return TraceContextOrSamplingFlags.create(maybeContext);
  }

  // Finally, try to extract an incoming, non-gRPC trace context. If tags exist, propagate them.
  TraceContextOrSamplingFlags result = delegate.extract(request);
  if (tagsBin == null) return result;
  return result.toBuilder().addExtra(tagsBin).build();
}
 

/** Shows we skip whitespace and extra fields like self or custom ones */
// https://aws.amazon.com/blogs/aws/application-performance-percentiles-and-request-tracing-for-aws-application-load-balancer/
@Test
public void extract_skipsSelfField() {
  // TODO: check with AWS if it is valid to have arbitrary fields in front of standard ones.
  // we currently permit them
  carrier.put(
      "x-amzn-trace-id",
      "Robot=Hello;Self=1-582113d1-1e48b74b3603af8479078ed6;  "
          + "Root=1-58211399-36d228ad5d99923122bbe354;  "
          + "TotalTimeSoFar=112ms;CalledFrom=Foo");

  TraceContextOrSamplingFlags extracted = extractor.extract(carrier);
  assertThat(extracted.traceIdContext())
      .isEqualTo(
          TraceIdContext.newBuilder()
              .traceIdHigh(lowerHexToUnsignedLong("5821139936d228ad"))
              .traceId(lowerHexToUnsignedLong("5d99923122bbe354"))
              .build());

  assertThat(((AmznTraceId) extracted.extra().get(0)).customFields)
      .contains(new StringBuilder(";Robot=Hello;TotalTimeSoFar=112ms;CalledFrom=Foo"));
}
 

@Test public void handleReceive_reusesTraceId() {
  httpTracing = HttpTracing.newBuilder(
    Tracing.newBuilder().supportsJoin(false).spanReporter(spans::add).build())
    .serverSampler(sampler).serverParser(parser).build();

  Tracer tracer = httpTracing.tracing().tracer();
  handler = HttpServerHandler.create(httpTracing, adapter);

  TraceContext incomingContext = tracer.nextSpan().context();
  when(extractor.extract(request)).thenReturn(
    TraceContextOrSamplingFlags.create(incomingContext));

  assertThat(handler.handleReceive(extractor, request).context())
    .extracting(TraceContext::traceId, TraceContext::parentId, TraceContext::shared)
    .containsOnly(incomingContext.traceId(), incomingContext.spanId(), false);
}
 

void verifyRoundTrip(TraceContextOrSamplingFlags expected) {
  TraceContextOrSamplingFlags extracted = propagation.extractor(mapEntry).extract(map);

  assertThat(extracted)
    .isEqualTo(expected);

  Map<String, String> injected = new LinkedHashMap<>();
  if (expected.context() != null) {
    propagation.injector(mapEntry).inject(expected.context(), injected);
  } else {
    inject(injected, expected.samplingFlags());
  }

  assertThat(map).isEqualTo(injected);
}
 

@RequestMapping(value = "/traceid", method = RequestMethod.GET)
public String traceId(@RequestHeader("b3") String b3Single) {
	TraceContextOrSamplingFlags traceContext = B3SingleFormat
			.parseB3SingleFormat(b3Single);
	then(traceContext.context()).isNotNull();
	return b3Single;
}
 

@Test public void nextSpan_extractedNothing_makesChildOfCurrent() {
  Span parent = tracer.newTrace();

  try (SpanInScope ws = tracer.withSpanInScope(parent)) {
    Span nextSpan = tracer.nextSpan(TraceContextOrSamplingFlags.create(EMPTY));
    assertThat(nextSpan.context().parentId())
      .isEqualTo(parent.context().spanId());
  }
}
 

TraceContextOrSamplingFlags takeRequest() {
  try {
    return requestQueue.poll(3, TimeUnit.SECONDS);
  } catch (InterruptedException e) {
    Thread.currentThread().interrupt();
    throw new AssertionError(e);
  }
}
 

<R> TraceContextOrSamplingFlags extractAndClearTraceIdHeaders(
  Extractor<R> extractor, R request, Message message
) {
  TraceContextOrSamplingFlags extracted = extractor.extract(request);
  // Clear any propagation keys present in the headers
  if (extracted.samplingFlags() == null) { // then trace IDs were extracted
    MessageProperties properties = message.getMessageProperties();
    if (properties != null) clearTraceIdHeaders(properties.getHeaders());
  }
  return extracted;
}
 

TraceContextOrSamplingFlags takeRequest() {
  try {
    return requestQueue.poll(3, TimeUnit.SECONDS);
  } catch (InterruptedException e) {
    Thread.currentThread().interrupt();
    throw new AssertionError(e);
  }
}
 

/** 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 extract_baggage_addsAllKeyNames_evenWhenEmpty() {
  TraceContextOrSamplingFlags extracted = extractor.extract(request);
  assertThat(extracted.extra()).hasSize(2);
  assertThat(extracted.extra().get(1))
      .asInstanceOf(InstanceOfAssertFactories.type(Extra.class))
      .extracting(a -> a.extractKeyNames)
      .asInstanceOf(InstanceOfAssertFactories.list(String.class))
      .containsExactly("x-vcap-request-id", "x-amzn-trace-id");
}
 

@Test public void onCompletion_shouldKeepContext_whenNotSampled() {
  Span span = tracing.tracer().nextSpan(TraceContextOrSamplingFlags.NOT_SAMPLED);

  Callback delegate =
    (metadata, exception) -> assertThat(tracing.tracer().currentSpan()).isEqualTo(span);
  Callback tracingCallback = TracingCallback.create(delegate, span, tracing.currentTraceContext());

  tracingCallback.onCompletion(null, null);
}
 

@Autowired
Baggage(Tracing tracing) {
	// When predefined baggage fields exist, the result !=
	// TraceContextOrSamplingFlags.EMPTY
	TraceContextOrSamplingFlags emptyExtraction = tracing.propagation()
			.extractor((c, k) -> null).extract(Boolean.TRUE);
	fields = BaggageField.getAll(emptyExtraction);
}
 

/** Creates a potentially noop span representing this request */
Span nextSpan(TraceContextOrSamplingFlags extracted, HttpServerRequest request) {
  Boolean sampled = extracted.sampled();
  // only recreate the context if the http sampler made a decision
  if (sampled == null && (sampled = sampler.trySample(request)) != null) {
    extracted = extracted.sampled(sampled.booleanValue());
  }
  return extracted.context() != null
    ? tracer.joinSpan(extracted.context())
    : tracer.nextSpan(extracted);
}
 

@Test public void parentId() {
  assertThat(BaggageFields.PARENT_ID.getValue(context))
    .isEqualTo(context.parentIdString());
  assertThat(BaggageFields.PARENT_ID.getValue(extracted))
    .isEqualTo(context.parentIdString());

  assertThat(BaggageFields.PARENT_ID.getValue(onlyMandatoryFields))
    .isNull();
  assertThat(BaggageFields.PARENT_ID.getValue(TraceContextOrSamplingFlags.EMPTY))
    .isNull();
  assertThat(BaggageFields.PARENT_ID.getValue(extractedTraceId))
    .isNull();
  assertThat(BaggageFields.PARENT_ID.getValue((TraceContext) null))
    .isNull();
}
 

/** Performs case-insensitive lookup */
@Override public TraceContextOrSamplingFlags extract(TextMapExtract entries) {
  Map<String, String> cache = new LinkedHashMap<>();
  for (Iterator<Map.Entry<String, String>> it = entries.iterator(); it.hasNext(); ) {
    Map.Entry<String, String> next = it.next();
    String inputKey = next.getKey().toLowerCase(Locale.ROOT);
    if (allNames.contains(inputKey)) {
      cache.put(inputKey, next.getValue());
    }
  }
  return delegate.extract(cache);
}
 

private TraceContextOrSamplingFlags extractAndClearHeaders(
		MessageConsumerRequest request) {
	TraceContextOrSamplingFlags extracted = extractor.extract(request);

	for (String propagationKey : tracing.propagation().keys()) {
		request.removeHeader(propagationKey);
	}

	return extracted;
}
 

/** 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);
}
 

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

<R> TraceContextOrSamplingFlags extractAndClearTraceIdHeaders(
  Extractor<R> extractor, R request, Headers headers
) {
  TraceContextOrSamplingFlags extracted = extractor.extract(request);
  // Clear any propagation keys present in the headers
  if (extracted.samplingFlags() == null) { // then trace IDs were extracted
    clearTraceIdHeaders(headers);
  }
  return extracted;
}
 

/**
 * Starts and propagates {@link Span.Kind#PRODUCER} span for each message sent.
 */
@Override
public Message<?> preSend(Message<?> message, MessageChannel channel) {
	if (emptyMessage(message)) {
		return message;
	}
	Message<?> retrievedMessage = getMessage(message);
	MessageHeaderAccessor headers = mutableHeaderAccessor(retrievedMessage);
	TraceContextOrSamplingFlags extracted = this.extractor.extract(headers);
	Span span = this.threadLocalSpan.next(extracted);
	MessageHeaderPropagation.removeAnyTraceHeaders(headers,
			this.tracing.propagation().keys());
	this.injector.inject(span.context(), headers);
	if (!span.isNoop()) {
		span.kind(Span.Kind.PRODUCER).name("send").start();
		span.remoteServiceName(toRemoteServiceName(headers));
		addTags(message, span, channel);
	}
	if (log.isDebugEnabled()) {
		log.debug("Created a new span in pre send" + span);
	}
	Message<?> outputMessage = outputMessage(message, retrievedMessage, headers);
	if (isDirectChannel(channel)) {
		beforeHandle(outputMessage, channel, null);
	}
	return outputMessage;
}
 

@Test public void getValue_extracted_doesntExist() {
  assertThat(AMZN_TRACE_ID.getValue(requestIdExtraction))
    .isNull();
  assertThat(AMZN_TRACE_ID.getValue(emptyExtraction))
    .isNull();
  assertThat(AMZN_TRACE_ID.getValue(TraceContextOrSamplingFlags.EMPTY))
    .isNull();
}
 

@Test public void spanId() {
  assertThat(BaggageFields.SPAN_ID.getValue(context))
    .isEqualTo(context.spanIdString());
  assertThat(BaggageFields.SPAN_ID.getValue(extracted))
    .isEqualTo(context.spanIdString());

  assertThat(BaggageFields.SPAN_ID.getValue(TraceContextOrSamplingFlags.EMPTY))
    .isNull();
  assertThat(BaggageFields.SPAN_ID.getValue(extractedTraceId))
    .isNull();
  assertThat(BaggageFields.SPAN_ID.getValue((TraceContext) null))
    .isNull();
}
 

@Override public Span apply(WrappedRequestSpec req, TraceContext parent) {
    if (parent != null) return span.next(TraceContextOrSamplingFlags.create(parent));
    Boolean sampled = sampler.trySample(ADAPTER, req);
    return span.next(TraceContextOrSamplingFlags.create(
        new SamplingFlags.Builder().sampled(sampled).build())
    );
}
 

/**
 * Use this to create a span for processing the given message. Note: the result has no name and is
 * not started.
 *
 * <p>In general, prefer {@link MessageListener} for processing messages, as it is more efficient
 * and less lossy with regards to context data.
 *
 * <p>This creates a child from identifiers extracted from the message properties, or a new span
 * if one couldn't be extracted.
 */
public Span nextSpan(Message message) {
  TraceContextOrSamplingFlags extracted =
    extractAndClearTraceIdProperties(processorExtractor, message, message);
  Span result = tracer.nextSpan(extracted); // Processor spans use the normal sampler.

  // When an upstream context was not present, lookup keys are unlikely added
  if (extracted.context() == null && !result.isNoop()) {
    // simplify code by re-using an existing MessagingRequest impl
    tagQueueOrTopic(new MessageConsumerRequest(message, destination(message)), result);
  }
  return result;
}
 

/**
 * This starts a consumer span as a child of the incoming message or the current trace
 * context. It then creates a span for the handler, placing it in scope.
 */
@Override
public Message<?> beforeHandle(Message<?> message, MessageChannel channel,
		MessageHandler handler) {
	if (emptyMessage(message)) {
		return message;
	}
	MessageHeaderAccessor headers = mutableHeaderAccessor(message);
	TraceContextOrSamplingFlags extracted = this.extractor.extract(headers);
	// Start and finish a consumer span as we will immediately process it.
	Span consumerSpan = this.tracer.nextSpan(extracted);
	if (!consumerSpan.isNoop()) {
		consumerSpan.kind(Span.Kind.CONSUMER).start();
		consumerSpan.remoteServiceName(REMOTE_SERVICE_NAME);
		addTags(message, consumerSpan, channel);
		consumerSpan.finish();
	}
	// create and scope a span for the message processor
	this.threadLocalSpan
			.next(TraceContextOrSamplingFlags.create(consumerSpan.context()))
			.name("handle").start();
	// remove any trace headers, but don't re-inject as we are synchronously
	// processing the
	// message and can rely on scoping to access this span later.
	MessageHeaderPropagation.removeAnyTraceHeaders(headers,
			this.tracing.propagation().keys());
	if (log.isDebugEnabled()) {
		log.debug("Created a new span in before handle" + consumerSpan);
	}
	if (message instanceof ErrorMessage) {
		return new ErrorMessage((Throwable) message.getPayload(),
				headers.getMessageHeaders());
	}
	headers.setImmutable();
	return new GenericMessage<>(message.getPayload(), headers.getMessageHeaders());
}