brave.propagation.TraceContext Java Examples

/**
 * Creates a potentially composite injector if the input is an instance of {@link RemoteSetter}.
 * Otherwise, a deferred injector is return that examples the request parameter to decide if it is
 * remote or not.
 */
public <R> TraceContext.Injector<R> newInjector(Setter<R, String> setter) {
  if (setter == null) throw new NullPointerException("setter == null");
  if (setter instanceof RemoteSetter) {
    RemoteSetter<?> remoteSetter = (RemoteSetter<?>) setter;
    switch (remoteSetter.spanKind()) {
      case CLIENT:
        return new RemoteInjector<>(setter, clientInjectorFunction);
      case PRODUCER:
        return new RemoteInjector<>(setter, producerInjectorFunction);
      case CONSUMER:
        return new RemoteInjector<>(setter, consumerInjectorFunction);
      default: // SERVER is nonsense as it cannot be injected
    }
  }
  return new DeferredInjector<>(setter, this);
}
 

@Test public void decorate_extractedExtra_plus_emptyParent() {
  TraceContext decorated = propagationFactory.decorate(context);
  BasicMapExtra extra1 = decorated.findExtra(BasicMapExtra.class);

  BasicMapExtra extracted = factory.create();
  extracted.put("2", "three");

  context2 = propagationFactory.decorate(
      context2.toBuilder().addExtra(extra1).addExtra(extracted).build());
  BasicMapExtra extra2 = context2.findExtra(BasicMapExtra.class);
  assertThat(context2.extra()).containsExactly(extra2); // merged

  assertThat(extra2.get("2")).isEqualTo("three");

  assertExtraClaimed(context2);
}
 

@Test
public void should_pass_tracing_info_when_using_reactor() {
	final AtomicReference<TraceContext> spanInOperation = new AtomicReference<>();
	Publisher<Integer> traced = Flux.just(1, 2, 3);

	try (Scope ws = this.currentTraceContext.newScope(context)) {
		Flux.from(traced).map(d -> d + 1).map(d -> d + 1).map((d) -> {
			spanInOperation.set(this.currentTraceContext.get());
			return d + 1;
		}).map(d -> d + 1).subscribe(d -> {
		});
	}

	then(this.currentTraceContext.get()).isNull();
	then(spanInOperation.get()).isEqualTo(context);
}
 

@Test public void customParser() {
  Tag<DubboResponse> javaValue = new Tag<DubboResponse>("dubbo.result_value") {
    @Override protected String parseValue(DubboResponse input, TraceContext context) {
      Result result = input.result();
      if (result == null) return null;
      return String.valueOf(result.getValue());
    }
  };

  RpcTracing rpcTracing = RpcTracing.newBuilder(tracing)
      .serverResponseParser((res, context, span) -> {
        RpcResponseParser.DEFAULT.parse(res, context, span);
        if (res instanceof DubboResponse) {
          javaValue.tag((DubboResponse) res, span);
        }
      }).build();
  init().setRpcTracing(rpcTracing);

  String javaResult = client.get().sayHello("jorge");

  assertThat(testSpanHandler.takeRemoteSpan(SERVER).tags())
      .containsEntry("dubbo.result_value", javaResult);
}
 

@Test public void doesntCrashOnNonFatalThrowable() {
  Throwable[] toThrow = new Throwable[1];
  SpanHandler handler =
      NoopAwareSpanHandler.create(new SpanHandler[] {new SpanHandler() {
        @Override public boolean end(TraceContext context, MutableSpan span, Cause cause) {
          doThrowUnsafely(toThrow[0]);
          return true;
        }
      }}, noop);

  toThrow[0] = new RuntimeException();
  assertThat(handler.end(context, span, Cause.FINISHED)).isTrue();

  toThrow[0] = new Exception();
  assertThat(handler.end(context, span, Cause.FINISHED)).isTrue();

  toThrow[0] = new Error();
  assertThat(handler.end(context, span, Cause.FINISHED)).isTrue();

  toThrow[0] = new StackOverflowError(); // fatal
  try { // assertThatThrownBy doesn't work with StackOverflowError
    handler.end(context, span, Cause.FINISHED);
    failBecauseExceptionWasNotThrown(StackOverflowError.class);
  } catch (StackOverflowError e) {
  }
}
 

@Override
@Nullable
public TraceContext get() {
    final RequestContext ctx = getRequestContextOrWarnOnce();
    if (ctx == null) {
        return THREAD_LOCAL_CONTEXT.get();
    }

    if (ctx.eventLoop().inEventLoop()) {
        return traceContext(ctx);
    } else {
        final TraceContext threadLocalContext = THREAD_LOCAL_CONTEXT.get();
        if (threadLocalContext != null) {
            return threadLocalContext;
        }
        // First span on a non-request thread will use the request's TraceContext as a parent.
        return traceContext(ctx);
    }
}
 

/** Ensure we use the same clock for traces that started in-process */
@Test
public void getOrCreate_reusesClockFromParent() {
  TraceContext trace = context;
  TraceContext traceJoin = trace.toBuilder().shared(true).build();
  TraceContext trace2 = context.toBuilder().traceId(2L).build();
  TraceContext traceChild =
    TraceContext.newBuilder().traceId(1L).parentId(trace.spanId()).spanId(3L).build();

  PendingSpan traceSpan = pendingSpans.getOrCreate(null, trace, false);
  PendingSpan traceJoinSpan = pendingSpans.getOrCreate(trace, traceJoin, false);
  PendingSpan trace2Span = pendingSpans.getOrCreate(null, trace2, false);
  PendingSpan traceChildSpan = pendingSpans.getOrCreate(trace, traceChild, false);

  assertThat(traceSpan.clock).isSameAs(traceChildSpan.clock);
  assertThat(traceSpan.clock).isSameAs(traceJoinSpan.clock);
  assertThat(traceSpan.clock).isNotSameAs(trace2Span.clock);
}
 

/**
 * This version of propagation contains at least 74 characters corresponding to identifiers and
 * the sampling bit. It will also include extra fields where present.
 *
 * <p>Ex 74 characters: {@code
 * Root=1-67891233-abcdef012345678912345678;Parent=463ac35c9f6413ad;Sampled=1}
 *
 * <p>{@inheritDoc}
 */
@Override
public void inject(TraceContext traceContext, R request) {
  AmznTraceId amznTraceId = traceContext.findExtra(AmznTraceId.class);
  CharSequence customFields = amznTraceId != null ? amznTraceId.customFields : null;
  int customFieldsLength = customFields == null ? 0 : customFields.length();
  // Root=1-67891233-abcdef012345678912345678;Parent=463ac35c9f6413ad;Sampled=1
  char[] result = new char[74 + customFieldsLength];
  System.arraycopy(ROOT, 0, result, 0, 5);
  writeRoot(traceContext, result, 5);
  System.arraycopy(PARENT, 0, result, 40, 8);
  writeHexLong(result, 48, traceContext.spanId());
  System.arraycopy(SAMPLED, 0, result, 64, 9);
  Boolean sampled = traceContext.sampled();
  // Sampled status is same as B3, but ? means downstream decides (like omitting X-B3-Sampled)
  // https://github.com/aws/aws-xray-sdk-go/blob/391885218b556c43ed05a1e736a766d70fc416f1/header/header.go#L50
  result[73] = sampled == null ? '?' : sampled ? '1' : '0';
  for (int i = 0; i < customFieldsLength; i++) {
    result[i + 74] = customFields.charAt(i);
  }
  setter.put(request, AMZN_TRACE_ID_NAME, new String(result));
}
 

@Override
public Scope newScope(TraceContext current) {
  final TraceContextHolder previous = registrySupplier.get()
      .maybeGet(TraceContextHolder.class)
      .orElse(TraceContextHolder.EMPTY);

  if (current != null) {
    registrySupplier.get().add(new TraceContextHolder(current));
    MDC.put(TRACE_ID_KEY, current.traceIdString());
  } else {
    registrySupplier.get().add(TraceContextHolder.EMPTY);
    MDC.remove(TRACE_ID_KEY);
  }

  return () -> {
    registrySupplier.get().add(previous);
    if (previous.context != null) {
      MDC.put(TRACE_ID_KEY, previous.context.traceIdString());
    } else {
      MDC.remove(TRACE_ID_KEY);
    }
  };
}
 

/** Writes 35 characters representing the input trace ID to the buffer at the given offset */
static void writeRoot(TraceContext context, char[] result, int offset) {
  result[offset] = '1'; // version
  result[offset + 1] = '-'; // delimiter
  long high = context.traceIdHigh();
  writeHexByte(result, offset + 2, (byte) ((high >>> 56L) & 0xff));
  writeHexByte(result, offset + 4, (byte) ((high >>> 48L) & 0xff));
  writeHexByte(result, offset + 6, (byte) ((high >>> 40L) & 0xff));
  writeHexByte(result, offset + 8, (byte) ((high >>> 32L) & 0xff));
  result[offset + 10] = '-';
  writeHexByte(result, offset + 11, (byte) ((high >>> 24L) & 0xff));
  writeHexByte(result, offset + 13, (byte) ((high >>> 16L) & 0xff));
  writeHexByte(result, offset + 15, (byte) ((high >>> 8L) & 0xff));
  writeHexByte(result, offset + 17, (byte) (high & 0xff));
  writeHexLong(result, offset + 19, context.traceId());
}
 

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

/**
 * Ensures the specified {@link Tracing} uses a {@link RequestContextCurrentTraceContext}.
 *
 * @throws IllegalStateException if {@code tracing} does not use {@link RequestContextCurrentTraceContext}
 */
public static void ensureScopeUsesRequestContext(Tracing tracing) {
    requireNonNull(tracing, "tracing");
    final PingPongExtra extra = new PingPongExtra();
    // trace contexts are not recorded until Tracer.toSpan, so this won't end up as junk data
    final TraceContext dummyContext = TraceContext.newBuilder().traceId(1).spanId(1)
                                                  .addExtra(extra).build();
    final boolean scopeUsesRequestContext;
    try (Scope scope = tracing.currentTraceContext().newScope(dummyContext)) {
        scopeUsesRequestContext = extra.isPong();
    }
    if (!scopeUsesRequestContext) {
        throw new IllegalStateException(
                "Tracing.currentTraceContext is not a " + RequestContextCurrentTraceContext.class
                        .getSimpleName() + " scope. " +
                "Please call Tracing.Builder.currentTraceContext(" + RequestContextCurrentTraceContext.class
                        .getSimpleName() + ".ofDefault()).");
    }
}
 

private void handleSendMessageBatchRequest(SendMessageBatchRequest request) {
  TraceContext maybeParent = currentTraceContext.get();

  Span span;
  if (maybeParent == null) {
    span = tracer.nextSpan();
  } else {
    // If we have a span in scope assume headers were cleared before
    span = tracer.newChild(maybeParent);
  }

  span.name("publish-batch").remoteServiceName("amazon-sqs").start();
  try (SpanInScope scope = tracer.withSpanInScope(span)) {
    for (SendMessageBatchRequestEntry entry : request.getEntries()) {
      injectPerMessage(request.getQueueUrl(), entry.getMessageAttributes());
    }
  } finally {
    span.finish();
  }
}
 

@Override public Scope newScope(TraceContext currentSpan) {
  Map<Buf, MarshalledContext.Cell> saved = broadcast().env();
  Map<Buf, MarshalledContext.Cell> update;
  if (currentSpan != null) { // replace the existing trace context with this one
    update = broadcast().env().updated(
      TRACE_ID_KEY.marshalId(),
      broadcast().Real().apply(TRACE_ID_KEY, new Some(toTraceId(currentSpan)))
    );
  } else { // remove the existing trace context from scope
    update = broadcast().env().filterKeys(
      new AbstractFunction1<Buf, Object>() {
        @Override public Object apply(Buf v1) {
          return !v1.equals(TRACE_ID_KEY.marshalId());
        }
      });
  }
  broadcastLocal.set(new Some(update));
  return () -> broadcastLocal.set(new Some(saved));
}
 

@Bean
MessageListener simpleMessageListener(CurrentTraceContext current) {
	return message -> {
		log.info("Got message");
		// Didn't restart the trace
		assertThat(current.get()).isNotNull()
				.extracting(TraceContext::parentIdAsLong).isNotEqualTo(0L);
	};
}
 

public static <T> Maybe<T> wrap(
  MaybeSource<T> source, CurrentTraceContext contextScoper, TraceContext assembled) {
  if (source instanceof Callable) {
    return new TraceContextCallableMaybe<>(source, contextScoper, assembled);
  }
  return new TraceContextMaybe<>(source, contextScoper, assembled);
}
 

@Bean
SpanHandler handlerTwo() {
	return new SpanHandler() {
		@Override
		public boolean end(TraceContext traceContext, MutableSpan span,
				Cause cause) {
			span.name(span.name() + " bar");
			return true; // keep this span
		}
	};
}
 

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

/**
 * Completes the span associated with this context, if it hasn't already been finished.
 *
 * @param timestamp zero means use the current time
 * @see brave.Span#finish()
 */
// zero here allows us to skip overhead of using the clock when the span already finished!
public void finish(TraceContext context, long timestamp) {
  PendingSpan last = remove(context);
  if (last == null) return;
  last.span.finishTimestamp(timestamp != 0L ? timestamp : last.clock.currentTimeMicroseconds());
  spanHandler.end(last.handlerContext, last.span, Cause.FINISHED);
}
 

Single<Integer> callableSingle(TraceContext expectedCallContext) {
  return RxJavaPlugins.onAssembly(new CallableSingle() {
    @Override public Integer call() {
      assertThat(currentTraceContext.get()).isEqualTo(expectedCallContext);
      return 1;
    }
  });
}
 

@Test
public void shouldSetPongIfOnlyExtra() {
    final PingPongExtra extra = new PingPongExtra();

    final TraceContext context = TraceContext.newBuilder().traceId(1).spanId(1)
                                             .addExtra(extra).build();

    TraceContextUtil.PingPongExtra.maybeSetPong(context);

    assertThat(extra.isPong()).isTrue();
}
 

@Test public void newChild_notSampledIsNoop() {
  TraceContext notSampled =
    tracer.newTrace().context().toBuilder().sampled(false).build();

  assertThat(tracer.newChild(notSampled))
    .isInstanceOf(NoopSpan.class);
}
 

@Bean
SpanHandler handlerTwo() {
	return new SpanHandler() {
		@Override
		public boolean end(TraceContext traceContext, MutableSpan span,
				Cause cause) {
			span.name(span.name() + " bar");
			return true; // keep this span
		}
	};
}
 

/**
 * In the case of {@link Cause#ORPHANED}, the calling thread will be an arbitrary invocation of
 * {@link Span} or {@link ScopedSpan} as spans orphaned from GC are expunged inline (not on the GC
 * thread). While this class is used for troubleshooting, it should do the least work possible to
 * prevent harm to arbitrary callers.
 */
@Override public boolean end(TraceContext context, MutableSpan span, Cause cause) {
  Throwable caller = spanToCaller.remove(span);
  if (cause != Cause.ORPHANED) return true;
  boolean allocatedButNotUsed = span.equals(new MutableSpan(context, defaultSpan));
  if (caller != null) log(context, allocatedButNotUsed, caller);
  if (allocatedButNotUsed) return true; // skip adding an annotation
  span.annotate(clock.currentTimeMicroseconds(), "brave.flush");
  return true;
}
 

@Test
public void errorMessageHeadersRetained() {
	this.channel.addInterceptor(this.interceptor);
	QueueChannel deadReplyChannel = new QueueChannel();
	QueueChannel errorsReplyChannel = new QueueChannel();
	Map<String, Object> errorChannelHeaders = new HashMap<>();
	errorChannelHeaders.put(MessageHeaders.REPLY_CHANNEL, errorsReplyChannel);
	errorChannelHeaders.put(MessageHeaders.ERROR_CHANNEL, errorsReplyChannel);
	this.channel.send(new ErrorMessage(
			new MessagingException(MessageBuilder.withPayload("hi")
					.setHeader("b3", "000000000000000a-000000000000000a")
					.setReplyChannel(deadReplyChannel)
					.setErrorChannel(deadReplyChannel).build()),
			errorChannelHeaders));

	this.message = this.channel.receive();

	assertThat(this.message).isNotNull();

	// Parse fails if trace or span ID are missing
	TraceContext context = parseB3SingleFormat(
			this.message.getHeaders().get("b3", String.class)).context();

	assertThat(context.traceIdString()).isEqualTo("000000000000000a");
	assertThat(context.spanIdString()).isNotEqualTo("000000000000000a");
	assertThat(this.spans).hasSize(2);
	assertThat(this.message.getHeaders().getReplyChannel())
			.isSameAs(errorsReplyChannel);
	assertThat(this.message.getHeaders().getErrorChannel())
			.isSameAs(errorsReplyChannel);
}
 

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

@Bean
SpanHandler testSpanHandler() {
	return new SpanHandler() {
		@Override
		public boolean end(TraceContext context, MutableSpan span, Cause cause) {
			return true;
		}
	};
}
 

Observable<Integer> scalarCallableObservable(TraceContext expectedCallContext,
  RuntimeException exception) {
  return RxJavaPlugins.onAssembly(new ScalarCallableObservable() {
    @Override public Integer call() {
      assertThat(currentTraceContext.get()).isEqualTo(expectedCallContext);
      throw exception;
    }
  });
}
 

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