Java Code Examples for io.grpc.Context#Key

@Test
public void GrpcContextPropagatesAcrossSchedulers() {
    final Context.Key<String> contextKey = Context.key("key");

    final AtomicBoolean done = new AtomicBoolean();

    Context.current().withValue(contextKey, "foo").wrap(new Runnable() {
        @Override
        public void run() {
            Observable.just(1, 2, 3)
                    .observeOn(Schedulers.computation())
                    .subscribeOn(Schedulers.io())
                    .subscribe(
                            new Consumer<Integer>() {
                                @Override
                                public void accept(Integer i) throws Exception {
                                    System.out.println(i);
                                    assertThat(contextKey.get()).isEqualTo("foo");
                                }
                            },
                            new Consumer<Throwable>() {
                                @Override
                                public void accept(Throwable throwable) throws Exception {

                                }
                            },
                            new Action() {
                                @Override
                                public void run() throws Exception {
                                    done.set(true);
                                }
                            });
        }
    }).run();

    await().atMost(Duration.FIVE_HUNDRED_MILLISECONDS).untilTrue(done);
}
 

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

Context.Key<String> getKey() {
  return key;
}
 

public static Context.Key<SofaRequest> getKeySofaRequest() {
    return keySofaRequest;
}
 

@Test
public void testGetKey() {
  Context.Key<Span> key = OpenTracingContextKey.getKey();
  assertEquals("Key should have correct name", key.toString(), (OpenTracingContextKey.KEY_NAME));
}
 

private static Context copyIntoContext(Context context, Metadata headers, Metadata.Key<String> headerKey, Context.Key<String> contextKey) {
    Object value = headers.get(headerKey);
    return value == null ? context : context.withValue(contextKey, value.toString());
}
 

public HeaderPropagationInterceptor(HeaderReader<Header> headerReader, Context.Key<Header> contextKey) {
    this.headerReader = Assert.requireNonNull(headerReader, "headerReader");
    this.contextKey = Assert.requireNonNull(contextKey, "contextKey");
}
 

public static Context.Key<Header> getAgentInfoKey() {
    return AGENT_INFO_KEY;
}
 

public static Context.Key<TransportMetadata> getTransportMetadataKey() {
    return TRANSPORT_METADATA_KEY;
}
 

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

/**
 * Retrieves the span key.
 *
 * @return the OpenTracing context key
 */
public static Context.Key<Span> getKey() {
    return key;
}
 

/**
 * Retrieves the span context key.
 *
 * @return the OpenTracing context key for span context
 */
public static Context.Key<SpanContext> getSpanContextKey() {
    return keyContext;
}
 

/**
 * Retrieves the span key.
 *
 * @return the OpenTracing context key
 */
public static Context.Key<Span> getKey() {
  return key;
}
 

/**
 * Retrieves the span context key.
 *
 * @return the OpenTracing context key for span context
 */
public static Context.Key<SpanContext> getSpanContextKey() {
  return keyContext;
}