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

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

public ServerCapabilities get(String buildRequestId, String commandId)
    throws IOException, InterruptedException {
  Context withMetadata =
      TracingMetadataUtils.contextWithMetadata(buildRequestId, commandId, "capabilities");
  Context previous = withMetadata.attach();
  try {
    GetCapabilitiesRequest request =
        instanceName == null
            ? GetCapabilitiesRequest.getDefaultInstance()
            : GetCapabilitiesRequest.newBuilder().setInstanceName(instanceName).build();
    return retrier.execute(() -> capabilitiesBlockingStub().getCapabilities(request));
  } catch (StatusRuntimeException e) {
    if (e.getCause() instanceof IOException) {
      throw (IOException) e.getCause();
    }
    throw new IOException(e);
  } finally {
    withMetadata.detach(previous);
  }
}
 

@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
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<Void, Void>(
      method,
      MoreExecutors.directExecutor(),
      callOpts,
      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 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 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
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 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
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 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 testDefaultNone() {
  ActiveSpanSource ss = ActiveSpanSource.NONE;
  assertNull("active span should always be null", ss.getActiveSpan());

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

  assertNull("active span should always be null", ss.getActiveSpan());

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

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

  assertEquals(OpenTracingContextKey.activeSpan(), span);

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

  assertNull(OpenTracingContextKey.activeSpan());
}
 

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

@Test
public void getCurrent_blankContext() {
  Context blankCtx = Context.ROOT;
  Context origCtx = blankCtx.attach();
  try {
    assertThat(CallMetricRecorder.getCurrent().isDisabled()).isTrue();
  } finally {
    blankCtx.detach(origCtx);
  }
}
 

@Test
public void callerContextPropagatedToListener() throws Exception {
  // Attach the context which is recorded when the call is created
  final Context.Key<String> testKey = Context.key("testing");
  Context context = Context.current().withValue(testKey, "testValue");
  Context previous = context.attach();

  ClientCallImpl<Void, Void> call = new ClientCallImpl<>(
      method.toBuilder().setType(MethodType.UNKNOWN).build(),
      new SerializingExecutor(Executors.newSingleThreadExecutor()),
      baseCallOptions,
      provider,
      deadlineCancellationExecutor,
      channelCallTracer,
      /* retryEnabled= */ false)
          .setDecompressorRegistry(decompressorRegistry);

  context.detach(previous);

  // Override the value after creating the call, this should not be seen by callbacks
  context = Context.current().withValue(testKey, "badValue");
  previous = context.attach();

  final AtomicBoolean onHeadersCalled = new AtomicBoolean();
  final AtomicBoolean onMessageCalled = new AtomicBoolean();
  final AtomicBoolean onReadyCalled = new AtomicBoolean();
  final AtomicBoolean observedIncorrectContext = new AtomicBoolean();
  final CountDownLatch latch = new CountDownLatch(1);

  call.start(new ClientCall.Listener<Void>() {
    @Override
    public void onHeaders(Metadata headers) {
      onHeadersCalled.set(true);
      checkContext();
    }

    @Override
    public void onMessage(Void message) {
      onMessageCalled.set(true);
      checkContext();
    }

    @Override
    public void onClose(Status status, Metadata trailers) {
      checkContext();
      latch.countDown();
    }

    @Override
    public void onReady() {
      onReadyCalled.set(true);
      checkContext();
    }

    private void checkContext() {
      if (!"testValue".equals(testKey.get())) {
        observedIncorrectContext.set(true);
      }
    }
  }, new Metadata());

  context.detach(previous);

  verify(stream).start(listenerArgumentCaptor.capture());
  ClientStreamListener listener = listenerArgumentCaptor.getValue();
  listener.onReady();
  listener.headersRead(new Metadata());
  listener.messagesAvailable(new SingleMessageProducer(new ByteArrayInputStream(new byte[0])));
  listener.messagesAvailable(new SingleMessageProducer(new ByteArrayInputStream(new byte[0])));
  listener.closed(Status.OK, new Metadata());

  assertTrue(latch.await(5, TimeUnit.SECONDS));

  assertTrue(onHeadersCalled.get());
  assertTrue(onMessageCalled.get());
  assertTrue(onReadyCalled.get());
  assertFalse(observedIncorrectContext.get());
}
 

@Override
public <ReqT> ClientStream newRetriableStream(
    final MethodDescriptor<ReqT, ?> method,
    final CallOptions callOptions,
    final Metadata headers,
    final Context context) {
  checkState(retryEnabled, "retry should be enabled");
  final Throttle throttle = lastServiceConfig.getRetryThrottling();
  final class RetryStream extends RetriableStream<ReqT> {
    RetryStream() {
      super(
          method,
          headers,
          channelBufferUsed,
          perRpcBufferLimit,
          channelBufferLimit,
          getCallExecutor(callOptions),
          transportFactory.getScheduledExecutorService(),
          callOptions.getOption(RETRY_POLICY_KEY),
          callOptions.getOption(HEDGING_POLICY_KEY),
          throttle);
    }

    @Override
    Status prestart() {
      return uncommittedRetriableStreamsRegistry.add(this);
    }

    @Override
    void postCommit() {
      uncommittedRetriableStreamsRegistry.remove(this);
    }

    @Override
    ClientStream newSubstream(ClientStreamTracer.Factory tracerFactory, Metadata newHeaders) {
      CallOptions newOptions = callOptions.withStreamTracerFactory(tracerFactory);
      ClientTransport transport =
          get(new PickSubchannelArgsImpl(method, newHeaders, newOptions));
      Context origContext = context.attach();
      try {
        return transport.newStream(method, newHeaders, newOptions);
      } finally {
        context.detach(origContext);
      }
    }
  }

  return new RetryStream();
}
 

@Test
public void callerContextPropagatedToListener() throws Exception {
  // Attach the context which is recorded when the call is created
  final Context.Key<String> testKey = Context.key("testing");
  Context context = Context.current().withValue(testKey, "testValue");
  Context previous = context.attach();

  ClientCallImpl<Void, Void> call = new ClientCallImpl<Void, Void>(
      method,
      new SerializingExecutor(Executors.newSingleThreadExecutor()),
      baseCallOptions,
      provider,
      deadlineCancellationExecutor,
      channelCallTracer,
      false /* retryEnabled */)
          .setDecompressorRegistry(decompressorRegistry);

  context.detach(previous);

  // Override the value after creating the call, this should not be seen by callbacks
  context = Context.current().withValue(testKey, "badValue");
  previous = context.attach();

  final AtomicBoolean onHeadersCalled = new AtomicBoolean();
  final AtomicBoolean onMessageCalled = new AtomicBoolean();
  final AtomicBoolean onReadyCalled = new AtomicBoolean();
  final AtomicBoolean observedIncorrectContext = new AtomicBoolean();
  final CountDownLatch latch = new CountDownLatch(1);

  call.start(new ClientCall.Listener<Void>() {
    @Override
    public void onHeaders(Metadata headers) {
      onHeadersCalled.set(true);
      checkContext();
    }

    @Override
    public void onMessage(Void message) {
      onMessageCalled.set(true);
      checkContext();
    }

    @Override
    public void onClose(Status status, Metadata trailers) {
      checkContext();
      latch.countDown();
    }

    @Override
    public void onReady() {
      onReadyCalled.set(true);
      checkContext();
    }

    private void checkContext() {
      if (!"testValue".equals(testKey.get())) {
        observedIncorrectContext.set(true);
      }
    }
  }, new Metadata());

  context.detach(previous);

  verify(stream).start(listenerArgumentCaptor.capture());
  ClientStreamListener listener = listenerArgumentCaptor.getValue();
  listener.onReady();
  listener.headersRead(new Metadata());
  listener.messagesAvailable(new SingleMessageProducer(new ByteArrayInputStream(new byte[0])));
  listener.messagesAvailable(new SingleMessageProducer(new ByteArrayInputStream(new byte[0])));
  listener.closed(Status.OK, new Metadata());

  assertTrue(latch.await(5, TimeUnit.SECONDS));

  assertTrue(onHeadersCalled.get());
  assertTrue(onMessageCalled.get());
  assertTrue(onReadyCalled.get());
  assertFalse(observedIncorrectContext.get());
}
 

@Override
public <ReqT> RetriableStream<ReqT> newRetriableStream(
    final MethodDescriptor<ReqT, ?> method,
    final CallOptions callOptions,
    final Metadata headers,
    final Context context) {
  checkState(retryEnabled, "retry should be enabled");
  final class RetryStream extends RetriableStream<ReqT> {
    RetryStream() {
      super(
          method,
          headers,
          channelBufferUsed,
          perRpcBufferLimit,
          channelBufferLimit,
          getCallExecutor(callOptions),
          transportFactory.getScheduledExecutorService(),
          callOptions.getOption(RETRY_POLICY_KEY),
          callOptions.getOption(HEDGING_POLICY_KEY),
          throttle);
    }

    @Override
    Status prestart() {
      return uncommittedRetriableStreamsRegistry.add(this);
    }

    @Override
    void postCommit() {
      uncommittedRetriableStreamsRegistry.remove(this);
    }

    @Override
    ClientStream newSubstream(ClientStreamTracer.Factory tracerFactory, Metadata newHeaders) {
      CallOptions newOptions = callOptions.withStreamTracerFactory(tracerFactory);
      ClientTransport transport =
          get(new PickSubchannelArgsImpl(method, newHeaders, newOptions));
      Context origContext = context.attach();
      try {
        return transport.newStream(method, newHeaders, newOptions);
      } finally {
        context.detach(origContext);
      }
    }
  }

  return new RetryStream();
}