Java Code Examples for io.grpc.Context#attach()

@Test
public void contextDeadlineShouldBePropagatedToStream() {
  Context context = Context.current()
      .withDeadlineAfter(1000, TimeUnit.MILLISECONDS, deadlineCancellationExecutor);
  Context origContext = context.attach();

  ClientCallImpl<Void, Void> call = new ClientCallImpl<Void, Void>(
      method,
      MoreExecutors.directExecutor(),
      baseCallOptions,
      provider,
      deadlineCancellationExecutor,
      channelCallTracer,
      false /* retryEnabled */);
  call.start(callListener, new Metadata());

  context.detach(origContext);

  ArgumentCaptor<Deadline> deadlineCaptor = ArgumentCaptor.forClass(Deadline.class);
  verify(stream).setDeadline(deadlineCaptor.capture());

  assertTimeoutBetween(deadlineCaptor.getValue().timeRemaining(TimeUnit.MILLISECONDS), 600, 1000);
}
 

@Test
public void testMultipleContextLayers() {
  Span parentSpan = tracer.buildSpan("s0").start();
  Context parentCtx = Context.current().withValue(OpenTracingContextKey.getKey(), parentSpan);
  Context previousCtx = parentCtx.attach();

  Span childSpan = tracer.buildSpan("s1").start();
  Context childCtx = Context.current().withValue(OpenTracingContextKey.getKey(), childSpan);
  parentCtx = childCtx.attach();
  assertEquals(OpenTracingContextKey.activeSpan(), childSpan);

  childCtx.detach(parentCtx);
  childSpan.finish();
  assertEquals(OpenTracingContextKey.activeSpan(), parentSpan);

  parentCtx.detach(previousCtx);
  parentSpan.finish();
  assertNull(OpenTracingContextKey.activeSpan());
}
 

@Test(timeout = 60000)
public void execute() throws Exception {
  final Thread callerThread = Thread.currentThread();
  final Context context = Context.current().withValue(KEY, "myvalue");
  previousContext = context.attach();

  final Semaphore tested = new Semaphore(0);

  executor.execute(
      new Runnable() {
        @Override
        public void run() {
          assertThat(Thread.currentThread()).isNotSameInstanceAs(callerThread);
          assertThat(Context.current()).isSameInstanceAs(context);
          assertThat(KEY.get()).isEqualTo("myvalue");
          tested.release();
        }
      });

  tested.acquire();
}
 

@Test
public void contextDeadlineShouldBePropagatedToStream() {
  Context context = Context.current()
      .withDeadlineAfter(1000, TimeUnit.MILLISECONDS, deadlineCancellationExecutor);
  Context origContext = context.attach();

  ClientCallImpl<Void, Void> call = new ClientCallImpl<>(
      method,
      MoreExecutors.directExecutor(),
      baseCallOptions,
      provider,
      deadlineCancellationExecutor,
      channelCallTracer,
      /* retryEnabled= */ false);
  call.start(callListener, new Metadata());

  context.detach(origContext);

  ArgumentCaptor<Deadline> deadlineCaptor = ArgumentCaptor.forClass(Deadline.class);
  verify(stream).setDeadline(deadlineCaptor.capture());

  assertTimeoutBetween(deadlineCaptor.getValue().timeRemaining(TimeUnit.MILLISECONDS), 600, 1000);
}
 

@Test(timeout = 60000)
public void start_Subclass() throws Exception {
  final Context context = Context.current().withValue(KEY, "myvalue");
  previousContext = context.attach();

  final AtomicBoolean tested = new AtomicBoolean(false);

  class MyThread extends Thread {

    @Override
    public void run() {
      assertThat(Context.current()).isSameInstanceAs(context);
      assertThat(KEY.get()).isEqualTo("myvalue");
      tested.set(true);
    }
  }

  Thread thread = new MyThread();

  thread.start();
  thread.join();

  assertThat(tested.get()).isTrue();
}
 

@Test(timeout = 60000)
public void start_Runnable() throws Exception {
  final Context context = Context.current().withValue(KEY, "myvalue");
  previousContext = context.attach();

  final AtomicBoolean tested = new AtomicBoolean(false);

  Runnable runnable =
      new Runnable() {
        @Override
        public void run() {
          assertThat(Context.current()).isSameInstanceAs(context);
          assertThat(KEY.get()).isEqualTo("myvalue");
          tested.set(true);
        }
      };
  Thread thread = new Thread(runnable);

  thread.start();
  thread.join();

  assertThat(tested.get()).isTrue();
}
 

@Test
public void testDefaultGrpc() {
  ActiveSpanSource ss = ActiveSpanSource.GRPC_CONTEXT;
  assertNull("active span should be null, no span in OpenTracingContextKey", ss.getActiveSpan());

  Span span = tracer.buildSpan("s0").start();
  Context ctx = Context.current().withValue(OpenTracingContextKey.getKey(), span);
  Context previousCtx = ctx.attach();

  assertEquals(
      "active span should be OpenTracingContextKey.activeSpan()", ss.getActiveSpan(), span);

  ctx.detach(previousCtx);
  span.finish();

  assertNull("active span should be null, no span in OpenTracingContextKey", ss.getActiveSpan());
}
 

@Test
public void testDefaultGrpc() {
  ActiveSpanContextSource ss = ActiveSpanContextSource.GRPC_CONTEXT;
  assertNull(
      "active span context should be null, no span context in OpenTracingContextKey",
      ss.getActiveSpanContext());

  Span span = tracer.buildSpan("s0").start();
  Context ctx =
      Context.current().withValue(OpenTracingContextKey.getSpanContextKey(), span.context());
  Context previousCtx = ctx.attach();

  assertEquals(
      "active span context should be OpenTracingContextKey.activeSpanContext()",
      ss.getActiveSpanContext(),
      span.context());

  ctx.detach(previousCtx);
  span.finish();

  assertNull(
      "active span context should be null, no span context in OpenTracingContextKey",
      ss.getActiveSpanContext());
}
 

@Test
public void contextDeadlineShouldNotOverrideSmallerCallOptionsDeadline() {
  Context context = Context.current()
      .withDeadlineAfter(2000, TimeUnit.MILLISECONDS, deadlineCancellationExecutor);
  Context origContext = context.attach();

  CallOptions callOpts = baseCallOptions.withDeadlineAfter(1000, TimeUnit.MILLISECONDS);
  ClientCallImpl<Void, Void> call = new ClientCallImpl<>(
      method,
      MoreExecutors.directExecutor(),
      callOpts,
      provider,
      deadlineCancellationExecutor,
      channelCallTracer,
      /* retryEnabled= */ false);
  call.start(callListener, new Metadata());

  context.detach(origContext);

  ArgumentCaptor<Deadline> deadlineCaptor = ArgumentCaptor.forClass(Deadline.class);
  verify(stream).setDeadline(deadlineCaptor.capture());

  assertTimeoutBetween(deadlineCaptor.getValue().timeRemaining(TimeUnit.MILLISECONDS), 600, 1000);
}
 

@Test
public void testMultipleContextLayersForSpanContext() {
  Span parentSpan = tracer.buildSpan("s0").start();
  Context parentCtx =
      Context.current()
          .withValue(OpenTracingContextKey.getSpanContextKey(), parentSpan.context());
  Context previousCtx = parentCtx.attach();

  Span childSpan = tracer.buildSpan("s1").start();
  Context childCtx =
      Context.current().withValue(OpenTracingContextKey.getSpanContextKey(), childSpan.context());
  parentCtx = childCtx.attach();
  assertEquals(OpenTracingContextKey.activeSpanContext(), childSpan.context());

  childCtx.detach(parentCtx);
  childSpan.finish();
  assertEquals(OpenTracingContextKey.activeSpanContext(), parentSpan.context());

  parentCtx.detach(previousCtx);
  parentSpan.finish();
  assertNull(OpenTracingContextKey.activeSpanContext());
}
 

@Test
public void contextDeadlineShouldOverrideLargerCallOptionsDeadline() {
  Context context = Context.current()
      .withDeadlineAfter(1000, TimeUnit.MILLISECONDS, deadlineCancellationExecutor);
  Context origContext = context.attach();

  CallOptions callOpts = baseCallOptions.withDeadlineAfter(2000, TimeUnit.MILLISECONDS);
  ClientCallImpl<Void, Void> call = new ClientCallImpl<>(
      method,
      MoreExecutors.directExecutor(),
      callOpts,
      provider,
      deadlineCancellationExecutor,
      channelCallTracer,
      /* retryEnabled= */ false);
  call.start(callListener, new Metadata());

  context.detach(origContext);

  ArgumentCaptor<Deadline> deadlineCaptor = ArgumentCaptor.forClass(Deadline.class);
  verify(stream).setDeadline(deadlineCaptor.capture());

  assertTimeoutBetween(deadlineCaptor.getValue().timeRemaining(TimeUnit.MILLISECONDS), 600, 1000);
}
 

@Test
public void expiredDeadlineCancelsStream_Context() {
  fakeClock.forwardTime(System.nanoTime(), TimeUnit.NANOSECONDS);

  Context context = Context.current()
      .withDeadlineAfter(1, TimeUnit.SECONDS, deadlineCancellationExecutor);
  Context origContext = context.attach();

  ClientCallImpl<Void, Void> call = new ClientCallImpl<Void, Void>(
      method,
      MoreExecutors.directExecutor(),
      baseCallOptions,
      provider,
      deadlineCancellationExecutor,
      channelCallTracer,
      false /* retryEnabled */);

  context.detach(origContext);

  call.start(callListener, new Metadata());

  fakeClock.forwardNanos(TimeUnit.SECONDS.toNanos(1) + 1);

  verify(stream, times(1)).cancel(statusCaptor.capture());
  assertEquals(Status.Code.DEADLINE_EXCEEDED, statusCaptor.getValue().getCode());
}
 

@Test(timeout = 60000)
public void submit_RunnableWithResult() throws Exception {
  final Thread callerThread = Thread.currentThread();
  final Context context = Context.current().withValue(KEY, "myvalue");
  previousContext = context.attach();

  final AtomicBoolean tested = new AtomicBoolean(false);
  Object result = new Object();

  Future<Object> future =
      executor.submit(
          new Runnable() {
            @Override
            public void run() {
              assertThat(Thread.currentThread()).isNotSameInstanceAs(callerThread);
              assertThat(Context.current()).isNotSameInstanceAs(Context.ROOT);
              assertThat(Context.current()).isSameInstanceAs(context);
              assertThat(KEY.get()).isEqualTo("myvalue");
              tested.set(true);
            }
          },
          result);

  assertThat(future.get()).isSameInstanceAs(result);
  assertThat(tested.get()).isTrue();
}
 

@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 testGetActiveSpanContext() {
  Span span = tracer.buildSpan("s0").start();
  Context ctx =
      Context.current().withValue(OpenTracingContextKey.getSpanContextKey(), span.context());
  Context previousCtx = ctx.attach();

  assertEquals(OpenTracingContextKey.activeSpanContext(), span.context());

  ctx.detach(previousCtx);
  span.finish();

  assertNull(OpenTracingContextKey.activeSpanContext());
}
 

@Test(timeout = 60000)
public void currentContextExecutor() throws Exception {
  final Thread callerThread = Thread.currentThread();
  final Context context = Context.current().withValue(KEY, "myvalue");
  previousContext = context.attach();

  final Semaphore tested = new Semaphore(0);

  Context.currentContextExecutor(executor)
      .execute(
          new Runnable() {
            @Override
            public void run() {
              StackTraceElement[] ste = new Exception().fillInStackTrace().getStackTrace();
              assertThat(ste[0].getClassName()).doesNotContain("Context");
              assertThat(ste[1].getClassName()).startsWith("io.grpc.Context$");
              // NB: Actually, we want the Runnable to be wrapped only once, but currently it is
              // still wrapped twice. The two places where the Runnable is wrapped are: (1) the
              // executor implementation itself, e.g. ThreadPoolExecutor, to which the Agent added
              // automatic context propagation, (2) CurrentContextExecutor.
              // ExecutorInstrumentation already avoids adding the automatic context propagation
              // to CurrentContextExecutor, but does not make it a no-op yet. Also see
              // ExecutorInstrumentation#createMatcher.
              assertThat(ste[2].getClassName()).startsWith("io.grpc.Context$");
              assertThat(ste[3].getClassName()).doesNotContain("Context");

              assertThat(Thread.currentThread()).isNotSameInstanceAs(callerThread);
              assertThat(Context.current()).isSameInstanceAs(context);
              assertThat(KEY.get()).isEqualTo("myvalue");

              tested.release();
            }
          });

  tested.acquire();
}
 

@Test
public void getCurrent_sameEnabledInstance() {
  CallMetricRecorder recorder = new CallMetricRecorder();
  Context ctx = Context.ROOT.withValue(CallMetricRecorder.CONTEXT_KEY, recorder);
  Context origCtx = ctx.attach();
  try {
    assertThat(CallMetricRecorder.getCurrent()).isSameInstanceAs(recorder);
    assertThat(recorder.isDisabled()).isFalse();
  } finally {
    ctx.detach(origCtx);
  }
}
 

@Test
public void expiredDeadlineCancelsStream_Context() {
  fakeClock.forwardTime(System.nanoTime(), TimeUnit.NANOSECONDS);

  Deadline deadline = Deadline.after(1, TimeUnit.SECONDS, fakeClock.getDeadlineTicker());
  Context context = Context.current().withDeadline(deadline, deadlineCancellationExecutor);
  Context origContext = context.attach();

  ClientCallImpl<Void, Void> call = new ClientCallImpl<>(
      method,
      MoreExecutors.directExecutor(),
      baseCallOptions,
      provider,
      deadlineCancellationExecutor,
      channelCallTracer,
      /* retryEnabled= */ false);

  context.detach(origContext);

  call.start(callListener, new Metadata());

  fakeClock.forwardTime(1000, TimeUnit.MILLISECONDS);
  verify(stream, never()).cancel(statusCaptor.capture());
  // verify app is notified.
  verify(callListener).onClose(statusCaptor.capture(), metadataArgumentCaptor.capture());
  assertThat(statusCaptor.getValue().getDescription()).contains("context timed out");
  assertThat(statusCaptor.getValue().getCode()).isEqualTo(Code.DEADLINE_EXCEEDED);

  // verify cancel send to server is delayed with DEADLINE_EXPIRATION_CANCEL_DELAY
  fakeClock.forwardNanos(DEADLINE_EXPIRATION_CANCEL_DELAY_NANOS);
  verify(stream).cancel(statusCaptor.capture());
  assertEquals(Status.Code.DEADLINE_EXCEEDED, statusCaptor.getValue().getCode());
  assertThat(statusCaptor.getValue().getDescription()).isEqualTo("context timed out");
}
 

public ContextInScope(Context context) {
  this.context = context;
  this.previous = context.attach();
}
 

@Override
public SpanContextHandle attach(SpanContext context) {
  Context current = Context.current().withValue(contextKey, context);
  Context previous = current.attach();
  return new GrpcSpanContextHandle(current, previous);
}