io.opencensus.trace.unsafe.ContextUtils Java Examples

@Test
public void traceHeadersPropagateSpanContext() throws Exception {
  CensusTracingModule.ClientCallTracer callTracer =
      censusTracing.newClientCallTracer(fakeClientParentSpan, method);
  Metadata headers = new Metadata();
  callTracer.newClientStreamTracer(CallOptions.DEFAULT, headers);

  verify(mockTracingPropagationHandler).toByteArray(same(fakeClientSpanContext));
  verifyNoMoreInteractions(mockTracingPropagationHandler);
  verify(tracer).spanBuilderWithExplicitParent(
      eq("Sent.package1.service2.method3"), same(fakeClientParentSpan));
  verify(spyClientSpanBuilder).setRecordEvents(eq(true));
  verifyNoMoreInteractions(tracer);
  assertTrue(headers.containsKey(censusTracing.tracingHeader));

  ServerStreamTracer serverTracer =
      censusTracing.getServerTracerFactory().newServerStreamTracer(
          method.getFullMethodName(), headers);
  verify(mockTracingPropagationHandler).fromByteArray(same(binarySpanContext));
  verify(tracer).spanBuilderWithRemoteParent(
      eq("Recv.package1.service2.method3"), same(spyClientSpan.getContext()));
  verify(spyServerSpanBuilder).setRecordEvents(eq(true));

  Context filteredContext = serverTracer.filterContext(Context.ROOT);
  assertSame(spyServerSpan, ContextUtils.CONTEXT_SPAN_KEY.get(filteredContext));
}
 

@Override
public void run() {
  Context origContext = ContextUtils.withValue(Context.current(), span).attach();
  try {
    runnable.run();
  } catch (Throwable t) {
    setErrorStatus(span, t);
    if (t instanceof RuntimeException) {
      throw (RuntimeException) t;
    } else if (t instanceof Error) {
      throw (Error) t;
    }
    throw new RuntimeException("unexpected", t);
  } finally {
    Context.current().detach(origContext);
    if (endSpan) {
      span.end();
    }
  }
}
 

@Override
public V call() throws Exception {
  Context origContext = ContextUtils.withValue(Context.current(), span).attach();
  try {
    return callable.call();
  } catch (Exception e) {
    setErrorStatus(span, e);
    throw e;
  } catch (Throwable t) {
    setErrorStatus(span, t);
    if (t instanceof Error) {
      throw (Error) t;
    }
    throw new RuntimeException("unexpected", t);
  } finally {
    Context.current().detach(origContext);
    if (endSpan) {
      span.end();
    }
  }
}
 

@Test
public void traceHeadersPropagateSpanContext() throws Exception {
  CensusTracingModule.ClientCallTracer callTracer =
      censusTracing.newClientCallTracer(fakeClientParentSpan, method);
  Metadata headers = new Metadata();
  callTracer.newClientStreamTracer(STREAM_INFO, headers);

  verify(mockTracingPropagationHandler).toByteArray(same(fakeClientSpanContext));
  verifyNoMoreInteractions(mockTracingPropagationHandler);
  verify(tracer).spanBuilderWithExplicitParent(
      eq("Sent.package1.service2.method3"), same(fakeClientParentSpan));
  verify(spyClientSpanBuilder).setRecordEvents(eq(true));
  verifyNoMoreInteractions(tracer);
  assertTrue(headers.containsKey(censusTracing.tracingHeader));

  ServerStreamTracer serverTracer =
      censusTracing.getServerTracerFactory().newServerStreamTracer(
          method.getFullMethodName(), headers);
  verify(mockTracingPropagationHandler).fromByteArray(same(binarySpanContext));
  verify(tracer).spanBuilderWithRemoteParent(
      eq("Recv.package1.service2.method3"), same(spyClientSpan.getContext()));
  verify(spyServerSpanBuilder).setRecordEvents(eq(true));

  Context filteredContext = serverTracer.filterContext(Context.ROOT);
  assertSame(spyServerSpan, ContextUtils.getValue(filteredContext));
}
 

@Test(timeout = 10000)
public void censusContextsPropagated() {
  Assume.assumeTrue("Skip the test because server is not in the same process.", server != null);
  Span clientParentSpan = Tracing.getTracer().spanBuilder("Test.interopTest").startSpan();
  // A valid ID is guaranteed to be unique, so we can verify it is actually propagated.
  assertTrue(clientParentSpan.getContext().getTraceId().isValid());
  Context ctx =
      Context.ROOT.withValues(
          TAG_CONTEXT_KEY,
          tagger.emptyBuilder().put(
              StatsTestUtils.EXTRA_TAG, TagValue.create("extra value")).build(),
          ContextUtils.CONTEXT_SPAN_KEY,
          clientParentSpan);
  Context origCtx = ctx.attach();
  try {
    blockingStub.unaryCall(SimpleRequest.getDefaultInstance());
    Context serverCtx = contextCapture.get();
    assertNotNull(serverCtx);

    FakeTagContext statsCtx = (FakeTagContext) TAG_CONTEXT_KEY.get(serverCtx);
    assertNotNull(statsCtx);
    Map<TagKey, TagValue> tags = statsCtx.getTags();
    boolean tagFound = false;
    for (Map.Entry<TagKey, TagValue> tag : tags.entrySet()) {
      if (tag.getKey().equals(StatsTestUtils.EXTRA_TAG)) {
        assertEquals(TagValue.create("extra value"), tag.getValue());
        tagFound = true;
      }
    }
    assertTrue("tag not found", tagFound);

    Span span = CONTEXT_SPAN_KEY.get(serverCtx);
    assertNotNull(span);
    SpanContext spanContext = span.getContext();
    assertEquals(clientParentSpan.getContext().getTraceId(), spanContext.getTraceId());
  } finally {
    ctx.detach(origCtx);
  }
}
 

@Test
public void getCurrentSpan() {
  assertThat(CurrentSpanUtils.getCurrentSpan()).isEqualTo(BlankSpan.INSTANCE);
  Context origContext = ContextUtils.withValue(Context.current(), span).attach();
  // Make sure context is detached even if test fails.
  try {
    assertThat(CurrentSpanUtils.getCurrentSpan()).isSameInstanceAs(span);
  } finally {
    Context.current().detach(origContext);
  }
  assertThat(CurrentSpanUtils.getCurrentSpan()).isEqualTo(BlankSpan.INSTANCE);
}
 

private SpanContext(Span span, boolean autoClose) {
  this.span = span;
  this.autoClose = autoClose;
  this.parent = CURRENT_SPAN.get();
  this.ocSpan = new OpenCensusSleuthSpan(span);
  this.ocCurrentContext = ContextUtils.withValue(Context.current(), this.ocSpan);
}
 

@Override
public Context filterContext(Context context) {
  // Access directly the unsafe trace API to create the new Context. This is a safe usage
  // because gRPC always creates a new Context for each of the server calls and does not
  // inherit from the parent Context.
  return ContextUtils.withValue(context, span);
}
 

@Override
public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(
    MethodDescriptor<ReqT, RespT> method, CallOptions callOptions, Channel next) {
  // New RPCs on client-side inherit the tracing context from the current Context.
  // Safe usage of the unsafe trace API because CONTEXT_SPAN_KEY.get() returns the same value
  // as Tracer.getCurrentSpan() except when no value available when the return value is null
  // for the direct access and BlankSpan when Tracer API is used.
  final ClientCallTracer tracerFactory =
      newClientCallTracer(ContextUtils.getValue(Context.current()), method);
  ClientCall<ReqT, RespT> call =
      next.newCall(
          method,
          callOptions.withStreamTracerFactory(tracerFactory));
  return new SimpleForwardingClientCall<ReqT, RespT>(call) {
    @Override
    public void start(Listener<RespT> responseListener, Metadata headers) {
      delegate().start(
          new SimpleForwardingClientCallListener<RespT>(responseListener) {
            @Override
            public void onClose(io.grpc.Status status, Metadata trailers) {
              tracerFactory.callEnded(status);
              super.onClose(status, trailers);
            }
          },
          headers);
    }
  };
}
 

@Test
public void serverBasicTracingNoHeaders() {
  ServerStreamTracer.Factory tracerFactory = censusTracing.getServerTracerFactory();
  ServerStreamTracer serverStreamTracer =
      tracerFactory.newServerStreamTracer(method.getFullMethodName(), new Metadata());
  verifyZeroInteractions(mockTracingPropagationHandler);
  verify(tracer).spanBuilderWithRemoteParent(
      eq("Recv.package1.service2.method3"), isNull(SpanContext.class));
  verify(spyServerSpanBuilder).setRecordEvents(eq(true));

  Context filteredContext = serverStreamTracer.filterContext(Context.ROOT);
  assertSame(spyServerSpan, ContextUtils.CONTEXT_SPAN_KEY.get(filteredContext));

  serverStreamTracer.serverCallStarted(
      new ServerCallInfoImpl<String, String>(method, Attributes.EMPTY, null));

  verify(spyServerSpan, never()).end(any(EndSpanOptions.class));

  serverStreamTracer.outboundMessage(0);
  serverStreamTracer.outboundMessageSent(0, 882, -1);
  serverStreamTracer.inboundMessage(0);
  serverStreamTracer.outboundMessage(1);
  serverStreamTracer.outboundMessageSent(1, -1, 27);
  serverStreamTracer.inboundMessageRead(0, 255, 90);

  serverStreamTracer.streamClosed(Status.CANCELLED);

  InOrder inOrder = inOrder(spyServerSpan);
  inOrder.verify(spyServerSpan, times(3)).addMessageEvent(messageEventCaptor.capture());
  List<MessageEvent> events = messageEventCaptor.getAllValues();
  assertEquals(
      MessageEvent.builder(Type.SENT, 0).setCompressedMessageSize(882).build(), events.get(0));
  assertEquals(
      MessageEvent.builder(Type.SENT, 1).setUncompressedMessageSize(27).build(), events.get(1));
  assertEquals(
      MessageEvent.builder(Type.RECEIVED, 0)
          .setCompressedMessageSize(255)
          .setUncompressedMessageSize(90)
          .build(),
      events.get(2));
  inOrder.verify(spyServerSpan).end(
      EndSpanOptions.builder()
          .setStatus(io.opencensus.trace.Status.CANCELLED)
          .setSampleToLocalSpanStore(false)
          .build());
  verifyNoMoreInteractions(spyServerSpan);
}
 

/**
 * Constructs a new {@link ScopeInSpan}.
 *
 * @param span is the {@code Span} to be added to the current {@code io.grpc.Context}.
 */
private ScopeInSpan(Span span, boolean endSpan) {
  this.span = span;
  this.endSpan = endSpan;
  origContext = ContextUtils.withValue(Context.current(), span).attach();
}
 

private static SpanContext getCurrentSpanContext() {
  Span span = ContextUtils.getValue(Context.current());
  return span == null ? SpanContext.INVALID : span.getContext();
}
 

private static SpanContext getCurrentSpanContext() {
  Span span = ContextUtils.getValue(Context.current());
  return span == null ? SpanContext.INVALID : span.getContext();
}
 

@Test
public void serverBasicTracingNoHeaders() {
  ServerStreamTracer.Factory tracerFactory = censusTracing.getServerTracerFactory();
  ServerStreamTracer serverStreamTracer =
      tracerFactory.newServerStreamTracer(method.getFullMethodName(), new Metadata());
  verifyZeroInteractions(mockTracingPropagationHandler);
  verify(tracer).spanBuilderWithRemoteParent(
      eq("Recv.package1.service2.method3"), ArgumentMatchers.<SpanContext>isNull());
  verify(spyServerSpanBuilder).setRecordEvents(eq(true));

  Context filteredContext = serverStreamTracer.filterContext(Context.ROOT);
  assertSame(spyServerSpan, ContextUtils.getValue(filteredContext));

  serverStreamTracer.serverCallStarted(
      new CallInfo<>(method, Attributes.EMPTY, null));

  verify(spyServerSpan, never()).end(any(EndSpanOptions.class));

  serverStreamTracer.outboundMessage(0);
  serverStreamTracer.outboundMessageSent(0, 882, -1);
  serverStreamTracer.inboundMessage(0);
  serverStreamTracer.outboundMessage(1);
  serverStreamTracer.outboundMessageSent(1, -1, 27);
  serverStreamTracer.inboundMessageRead(0, 255, 90);

  serverStreamTracer.streamClosed(Status.CANCELLED);

  InOrder inOrder = inOrder(spyServerSpan);
  inOrder.verify(spyServerSpan, times(3)).addMessageEvent(messageEventCaptor.capture());
  List<MessageEvent> events = messageEventCaptor.getAllValues();
  assertEquals(
      MessageEvent.builder(Type.SENT, 0).setCompressedMessageSize(882).build(), events.get(0));
  assertEquals(
      MessageEvent.builder(Type.SENT, 1).setUncompressedMessageSize(27).build(), events.get(1));
  assertEquals(
      MessageEvent.builder(Type.RECEIVED, 0)
          .setCompressedMessageSize(255)
          .setUncompressedMessageSize(90)
          .build(),
      events.get(2));
  inOrder.verify(spyServerSpan).end(
      EndSpanOptions.builder()
          .setStatus(io.opencensus.trace.Status.CANCELLED)
          .setSampleToLocalSpanStore(false)
          .build());
  verifyNoMoreInteractions(spyServerSpan);
}
 

/**
 * Returns The {@link Span} from the current context.
 *
 * @return The {@code Span} from the current context.
 */
@Nullable
static Span getCurrentSpan() {
  return ContextUtils.getValue(Context.current());
}