brave.propagation.CurrentTraceContext.Scope Java Examples

@Override public Response intercept(Chain chain) throws IOException {
  RequestWrapper request = new RequestWrapper(chain.request());

  Span span;
  TraceContext parent = chain.request().tag(TraceContext.class);
  if (parent != null) { // TracingCallFactory setup this request
    span = handler.handleSendWithParent(request, parent != NULL_SENTINEL ? parent : null);
  } else { // This is using interceptors only
    span = handler.handleSend(request);
  }

  parseRouteAddress(chain, span);

  Response response = null;
  Throwable error = null;
  try (Scope ws = currentTraceContext.newScope(span.context())) {
    return response = chain.proceed(request.build());
  } catch (Throwable t) {
    error = t;
    throw t;
  } finally {
    // Intentionally not the same instance as chain.proceed, as properties may have changed
    if (response != null) request = new RequestWrapper(response.request());
    handler.handleReceive(new ResponseWrapper(request, response, error), span);
  }
}
 

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

protected void is_inheritable(CurrentTraceContext inheritableCurrentTraceContext)
  throws Exception {
  // use a single-threaded version of newCachedThreadPool
  ExecutorService service = new ThreadPoolExecutor(0, 1,
    60L, TimeUnit.SECONDS, new LinkedBlockingQueue<>());

  // submitting a job grows the pool, attaching the context to its thread
  try (Scope scope = inheritableCurrentTraceContext.newScope(context)) {
    assertThat(service.submit(inheritableCurrentTraceContext::get).get())
      .isEqualTo(context);
  }

  // same thread executes the next job and still has the same context (leaked and not cleaned up)
  assertThat(service.submit(inheritableCurrentTraceContext::get).get())
    .isEqualTo(context);

  service.shutdownNow();
}
 

/**
 * 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 attachesSpanInRunnable() {
  Runnable runnable;
  try (Scope scope = currentTraceContext.newScope(context)) {
    runnable = currentTraceContext.wrap(() -> {
      assertThat(currentTraceContext.get())
        .isEqualTo(context);
      verifyImplicitContext(context);
    });

    runnable.run(); // runs assertion in the same scope
  }

  // Set another span between the time the task was made and executed.
  try (Scope scope2 = currentTraceContext.newScope(unsampledContext)) {
    runnable.run(); // runs assertion
    verifyImplicitContext(unsampledContext);
  }
}
 

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

@Override public void channelRead(ChannelHandlerContext ctx, Object msg) {
  if (!(msg instanceof HttpRequest)) {
    ctx.fireChannelRead(msg); // superclass does not throw
    return;
  }

  HttpRequestWrapper request =
    new HttpRequestWrapper((HttpRequest) msg, (InetSocketAddress) ctx.channel().remoteAddress());

  ctx.channel().attr(NettyHttpTracing.REQUEST_ATTRIBUTE).set(request);
  Span span = handler.handleReceive(request);
  ctx.channel().attr(NettyHttpTracing.SPAN_ATTRIBUTE).set(span);
  Scope scope = currentTraceContext.newScope(span.context());

  // Place the span in scope so that downstream code can read trace IDs
  Throwable error = null;
  try {
    ctx.fireChannelRead(msg);
  } catch (Throwable e) {
    error = e;
    throw e;
  } finally {
    if (error != null) span.error(error).finish();
    scope.close();
  }
}
 

@Test
public void should_not_scope_scalar_subscribe() {
	springContext.registerBean(CurrentTraceContext.class, () -> currentTraceContext);
	springContext.refresh();

	Function<? super Publisher<Integer>, ? extends Publisher<Integer>> transformer = scopePassingSpanOperator(
			this.springContext);

	try (Scope ws = this.currentTraceContext.newScope(context)) {

		transformer.apply(Mono.just(1))
				.subscribe(assertNotScopePassingSpanSubscriber);

		transformer.apply(Mono.error(new Exception()))
				.subscribe(assertNotScopePassingSpanSubscriber);

		transformer.apply(Mono.empty())
				.subscribe(assertNotScopePassingSpanSubscriber);

	}
}
 

@Test public void revertsChanges_withMultipleBaggageFields() {
  map.put("X-B3-TraceId", "000000000000000a");
  map.put(FIELD.name(), "bob");
  map.put(FIELD_2.name(), "BV");
  map.put(LOCAL_FIELD.name(), "ef01");
  Map<String, String> snapshot = new LinkedHashMap<>(map);

  FIELD.baggageField().updateValue(contextWithBaggage, "romeo");
  FIELD_2.baggageField().updateValue(contextWithBaggage, "FO");
  LOCAL_FIELD.baggageField().updateValue(contextWithBaggage, "abcd");

  try (Scope s = withBaggageFieldsDecorator.decorateScope(contextWithBaggage, mock(Scope.class))) {
    assertThat(map).containsOnly(
      entry("X-B3-TraceId", "0000000000000001"),
      entry(FIELD.name(), "romeo"),
      entry(FIELD_2.name(), "FO"),
      entry(LOCAL_FIELD.name(), "abcd")
    );
  }
  assertThat(map).isEqualTo(snapshot);
  map.clear();
}
 

void assertNestedUpdatesCoherent(ScopeDecorator decorator) {
  try (Scope s = decorator.decorateScope(contextWithBaggage, mock(Scope.class))) {
    FLUSH_FIELD.baggageField().updateValue(contextWithBaggage, "word");
    try (Scope s1 = decorator.decorateScope(contextWithBaggage, mock(Scope.class))) {
      assertThat(map).containsEntry("flushed", "word");
      FLUSH_FIELD.baggageField().updateValue(contextWithBaggage, "outlook");
      try (Scope s2 = decorator.decorateScope(contextWithBaggage, mock(Scope.class))) {
        assertThat(map).containsEntry("flushed", "outlook");
        FLUSH_FIELD.baggageField().updateValue(contextWithBaggage, "powerpoint");
        try (Scope s3 = decorator.decorateScope(contextWithBaggage, mock(Scope.class))) {
          assertThat(map).containsEntry("flushed", "powerpoint");
          FLUSH_FIELD.baggageField().updateValue(contextWithBaggage, "sharepoint");
          assertThat(map).containsEntry("flushed", "sharepoint");
        }
        assertThat(map).containsEntry("flushed", "powerpoint");
      }
      assertThat(map).containsEntry("flushed", "outlook");
    }
    assertThat(map).containsEntry("flushed", "word");
  }
}
 

@Test
public void should_support_reactor_fusion_optimization() {
	final AtomicReference<TraceContext> spanInOperation = new AtomicReference<>();

	try (Scope ws = this.currentTraceContext.newScope(context)) {
		Mono.just(1).flatMap(d -> Flux.just(d + 1).collectList().map(p -> p.get(0)))
				.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);
}
 

public ExceptionHandler exception(ExceptionHandler delegate) {
  return (error, req, res) -> {
    try {
      delegate.handle(error, req, res);
    } finally {
      Span span = req.attribute(Span.class.getName());
      if (span != null) {
        HttpServerResponse response =
          HttpServletResponseWrapper.create(req.raw(), res.raw(), error);
        handler.handleSend(response, span);
        req.raw().removeAttribute(Span.class.getName()); // prevent double-processing
        ((Scope) req.attribute(Scope.class.getName())).close();
      }
    }
  };
}
 

@Test
public void checkTraceIdDuringZipOperation() {
	final AtomicReference<TraceContext> spanInOperation = new AtomicReference<>();
	final AtomicReference<TraceContext> spanInZipOperation = new AtomicReference<>();

	try (Scope ws = this.currentTraceContext.newScope(context)) {
		Mono.fromCallable(this.currentTraceContext::get).map(span -> span)
				.doOnNext(spanInOperation::set)
				.zipWith(Mono.fromCallable(this.currentTraceContext::get)
						.map(span -> span).doOnNext(spanInZipOperation::set))
				.block();
	}

	then(spanInZipOperation).hasValue(context);
	then(spanInOperation).hasValue(context);
}
 

@Test public void equals_realSpan_sameContext() {
  Span current;
  try (Scope ws = tracing.currentTraceContext().newScope(context)) {
    current = tracing.tracer().currentSpan();
  }

  assertThat(current).isEqualTo(tracing.tracer().toSpan(context));
}
 

@Test public void equals_noopSpan_sameContext() {
  Span current;
  try (Scope ws = tracing.currentTraceContext().newScope(unsampledContext)) {
    current = tracing.tracer().currentSpan();
  }

  assertThat(current).isEqualTo(tracing.tracer().toSpan(unsampledContext));
}
 

@Test public void current_get() {
  TraceContext context = extractWithAmazonTraceId();

  try (Tracing t = Tracing.newBuilder().propagationFactory(factory).build();
       Scope scope = t.currentTraceContext().newScope(context)) {
    assertThat(ExtraFieldPropagation.get("x-amzn-trace-id"))
      .isEqualTo(awsTraceId);
  }
}
 

@Override public void onComplete() {
  Scope scope = contextScoper.maybeScope(assembled);
  try { // retrolambda can't resolve this try/finally
    downstream.onComplete();
  } finally {
    scope.close();
  }
}
 

@Override public void onError(Throwable t) {
  Scope scope = contextScoper.maybeScope(assembled);
  try { // retrolambda can't resolve this try/finally
    downstream.onError(t);
  } finally {
    scope.close();
  }
}
 

@Override public void process(HttpRequest request, HttpContext context) {
  TraceContext parent = (TraceContext) context.getAttribute(TraceContext.class.getName());

  HttpRequestWrapper wrapper = new HttpRequestWrapper(request, context);
  Span span = handler.handleSendWithParent(wrapper, parent);
  parseTargetAddress(wrapper.target, span);

  context.setAttribute(Span.class.getName(), span);
  context.setAttribute(Scope.class.getName(), currentTraceContext.newScope(span.context()));
}
 

@Override public void onSuccess(T value) {
  Scope scope = contextScoper.maybeScope(assembled);
  try { // retrolambda can't resolve this try/finally
    downstream.onSuccess(value);
  } finally {
    scope.close();
  }
}
 

ScopedSpan newScopedSpan(@Nullable TraceContext parent, TraceContext context, String name) {
  Scope scope = currentTraceContext.newScope(context);
  if (isNoop(context)) return new NoopScopedSpan(context, scope);

  PendingSpan pendingSpan = pendingSpans.getOrCreate(parent, context, true);
  Clock clock = pendingSpan.clock();
  MutableSpan state = pendingSpan.state();
  state.name(name);
  return new RealScopedSpan(context, scope, state, clock, pendingSpans);
}
 

@Test public void parallelFlowable_assembleInScope_subscribeInScope() {
  ParallelFlowable<Integer> source, errorSource;
  try (Scope scope = currentTraceContext.newScope(assemblyContext)) {
    source = Flowable.range(1, 3).parallel()
      .doOnNext(e -> assertInAssemblyContext())
      .doOnComplete(this::assertInAssemblyContext);
    errorSource = Flowable.<Integer>concat(
      Flowable.error(new IllegalStateException()), Flowable.error(new IllegalStateException()))
      .parallel()
      .doOnError(t -> assertInAssemblyContext())
      .doOnComplete(this::assertInAssemblyContext);
  }

  subscribeInDifferentContext(source, errorSource).assertResult(1, 2, 3);
}
 

TestObserver<Integer> subscribeInDifferentContext(ParallelFlowable<Integer> source,
  ParallelFlowable<Integer> errorSource) {
  source = source.doOnSubscribe(s -> assertInSubscribeContext());
  errorSource = errorSource.doOnSubscribe(s -> assertInSubscribeContext());

  try (Scope scope2 = currentTraceContext.newScope(subscribeContext)) {
    errorSource.sequential().test().assertFailure(IllegalStateException.class);
    return source.sequential().toObservable().test();
  }
}
 

public Filter afterAfter() {
  return (req, res) -> {
    Span span = req.attribute(Span.class.getName());
    if (span == null) return;
    HttpServerResponse response = HttpServletResponseWrapper.create(req.raw(), res.raw(), null);
    handler.handleSend(response, span);
    ((Scope) req.attribute(Scope.class.getName())).close();
  };
}
 

@Test
public void createAsync_observable_success() {
  rxjava_createAsync_success(
    (service, observer) -> {
      try (Scope scope = currentTraceContext.newScope(context1)) {
        service.observable().subscribe(observer);
      }
    });
}
 

@Test
public void createAsync_completable_success() {
  rxjava_createAsync_success(
    (service, observer) -> {
      try (Scope scope = currentTraceContext.newScope(context1)) {
        service.completable().subscribe(observer);
      }
    });
}
 

void retainsContext(Scope scope) {
  try {
    assertThat(scope).isNotEqualTo(Scope.NOOP);
    assertThat(currentTraceContext.get())
      .isEqualTo(context);
    verifyImplicitContext(context);
  } finally {
    scope.close();
  }
}
 

@Override public void sendHeaders(Metadata headers) {
  try (Scope scope = currentTraceContext.maybeScope(context)) {
    delegate().sendHeaders(headers);
  }
  // sendHeaders() JavaDoc mentions headers are not thread-safe, so we make a safe copy here.
  this.headers.merge(headers);
}
 

Single(
  Scope delegate,
  CorrelationContext context,
  SingleCorrelationField field,
  @Nullable String valueToRevert,
  boolean shouldRevert
) {
  super(context);
  this.delegate = delegate;
  this.field = field;
  this.valueToRevert = valueToRevert;
  this.shouldRevert = shouldRevert;
}
 

@Test
public void newScope_canClearScope() {
    try (SafeCloseable requestContextScope = ctx.push()) {
        try (Scope traceContextScope = currentTraceContext.newScope(traceContext)) {
            try (Scope traceContextScope2 = currentTraceContext.newScope(null)) {
                assertThat(currentTraceContext.get()).isNull();
            }
            assertThat(currentTraceContext.get()).isEqualTo(traceContext);
        }
    }
}