Java 类io.grpc.Context 实例源码

@Test
public void interceptorShouldFreezeContext() {
    TestService svc = new TestService();

    // Plumbing
    serverRule.getServiceRegistry().addService(ServerInterceptors.interceptForward(svc,
            new AmbientContextServerInterceptor("ctx-"),
            new AmbientContextFreezeServerInterceptor()));
    GreeterGrpc.GreeterBlockingStub stub = GreeterGrpc
            .newBlockingStub(serverRule.getChannel())
            .withInterceptors(new AmbientContextClientInterceptor("ctx-"));

    // Test
    Metadata.Key<String> key = Metadata.Key.of("ctx-k", Metadata.ASCII_STRING_MARSHALLER);
    AmbientContext.initialize(Context.current()).run(() -> {
        AmbientContext.current().put(key, "value");
        stub.sayHello(HelloRequest.newBuilder().setName("World").build());
    });

    assertThat(svc.frozen).isTrue();
}

@Override
public Statement apply(final Statement base, final Description description) {
    return new Statement() {
        @Override
        public void evaluate() throws Throwable {
            // Reset the gRPC context between test executions
            Context prev = Context.ROOT.attach();
            try {
                base.evaluate();
                if (Context.current() != Context.ROOT) {
                    Assert.fail("Test is leaking context state between tests! Ensure proper " +
                            "attach()/detach() pairing.");
                }
            } finally {
                Context.ROOT.detach(prev);
            }
        }
    };
}

@Test
public void ruleSetsContextToRoot() {
    Context.current().withValue(Context.key("foo"), "bar").run(() -> {
        assertThat(Context.current()).isNotEqualTo(Context.ROOT);

        try {
            GrpcContextRule rule = new GrpcContextRule();
            rule.apply(new Statement() {
                @Override
                public void evaluate() {
                    assertThat(Context.current()).isEqualTo(Context.ROOT);
                }
            }, Description.createTestDescription(GrpcContextRuleTest.class, "ruleSetsContextToRoot"))
            .evaluate();
        } catch (Throwable throwable) {
            fail(throwable.getMessage());
        }
    });
}

@Test
public void ruleFailsIfContextLeaks() {
    Context.current().withValue(Context.key("foo"), "bar").run(() -> {
        assertThat(Context.current()).isNotEqualTo(Context.ROOT);

        assertThatThrownBy(() -> {
            GrpcContextRule rule = new GrpcContextRule();
            rule.apply(new Statement() {
                @Override
                public void evaluate() {
                    // Leak context
                    Context.current().withValue(Context.key("cheese"), "baz").attach();
                }
            }, Description.createTestDescription(GrpcContextRuleTest.class, "ruleSetsContextToRoot"))
            .evaluate();
        }).isInstanceOf(AssertionError.class).hasMessageContaining("Test is leaking context");
    });
}

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

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

@Test
public void record_CurrentContextSet() {
  View view =
      View.create(
          VIEW_NAME,
          "description",
          MEASURE_DOUBLE,
          Sum.create(),
          Arrays.asList(KEY),
          Cumulative.create());
  viewManager.registerView(view);
  Context orig =
      Context.current()
          .withValue(ContextUtils.TAG_CONTEXT_KEY, new SimpleTagContext(Tag.create(KEY, VALUE)))
          .attach();
  try {
    statsRecorder.newMeasureMap().put(MEASURE_DOUBLE, 1.0).record();
  } finally {
    Context.current().detach(orig);
  }
  ViewData viewData = viewManager.getView(VIEW_NAME);

  // record() should have used the given TagContext.
  assertThat(viewData.getAggregationMap().keySet()).containsExactly(Arrays.asList(VALUE));
}

@Test
public void testWithTagContextUsingWrap() {
  Runnable runnable;
  Scope scopedTags = CurrentTagContextUtils.withTagContext(tagContext);
  try {
    assertThat(CurrentTagContextUtils.getCurrentTagContext()).isSameAs(tagContext);
    runnable =
        Context.current()
            .wrap(
                new Runnable() {
                  @Override
                  public void run() {
                    assertThat(CurrentTagContextUtils.getCurrentTagContext())
                        .isSameAs(tagContext);
                  }
                });
  } finally {
    scopedTags.close();
  }
  assertThat(tagContextToList(CurrentTagContextUtils.getCurrentTagContext())).isEmpty();
  // When we run the runnable we will have the TagContext in the current Context.
  runnable.run();
}

@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()).isSameAs(context);
          assertThat(KEY.get()).isEqualTo("myvalue");
          tested.set(true);
        }
      };
  Thread thread = new Thread(runnable);

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

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

@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()).isSameAs(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 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()).isNotSameAs(callerThread);
          assertThat(Context.current()).isSameAs(context);
          assertThat(KEY.get()).isEqualTo("myvalue");
          tested.release();
        }
      });

  tested.acquire();
}

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

  final AtomicBoolean tested = new AtomicBoolean(false);

  executor
      .submit(
          new Callable<Void>() {
            @Override
            public Void call() throws Exception {
              assertThat(Thread.currentThread()).isNotSameAs(callerThread);
              assertThat(Context.current()).isSameAs(context);
              assertThat(KEY.get()).isEqualTo("myvalue");
              tested.set(true);

              return null;
            }
          })
      .get();

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

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

  final AtomicBoolean tested = new AtomicBoolean(false);

  executor
      .submit(
          new Runnable() {
            @Override
            public void run() {
              assertThat(Thread.currentThread()).isNotSameAs(callerThread);
              assertThat(Context.current()).isSameAs(context);
              assertThat(KEY.get()).isEqualTo("myvalue");
              tested.set(true);
            }
          })
      .get();

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

@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()).isNotSameAs(callerThread);
              assertThat(Context.current()).isNotSameAs(Context.ROOT);
              assertThat(Context.current()).isSameAs(context);
              assertThat(KEY.get()).isEqualTo("myvalue");
              tested.set(true);
            }
          },
          result);

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

@Test
public void doesPropagateContext() throws Exception {
  final Context oldContext = Context.current();
  final Context newContext = oldContext.withValues(KEY_1, VAL_1, KEY_2, VAL_2);

  newContext.attach();

  final TestSubscriber<Object> subscriber = new TestSubscriber<Object>();
  Observable.create(subscriber1 -> {
    subscriber1.onNext(KEY_1.get());
    subscriber1.onNext(KEY_2.get());
    subscriber1.onCompleted();
  }).subscribeOn(Schedulers.computation()).subscribe(subscriber);

  newContext.detach(oldContext);

  subscriber.awaitTerminalEvent();
  subscriber.assertValues(VAL_1, VAL_2);
}

@Before
public final void setUp() throws Exception {
  Chunker.setDefaultChunkSizeForTesting(1000); // Enough for everything to be one chunk.
  fs = new InMemoryFileSystem(new JavaClock(), HashFunction.SHA256);
  execRoot = fs.getPath("/exec/root");
  FileSystemUtils.createDirectoryAndParents(execRoot);
  fakeFileCache = new FakeActionInputFileCache(execRoot);

  Path stdout = fs.getPath("/tmp/stdout");
  Path stderr = fs.getPath("/tmp/stderr");
  FileSystemUtils.createDirectoryAndParents(stdout.getParentDirectory());
  FileSystemUtils.createDirectoryAndParents(stderr.getParentDirectory());
  Context withEmptyMetadata =
      TracingMetadataUtils.contextWithMetadata(
          "none", "none", DIGEST_UTIL.asActionKey(Digest.getDefaultInstance()));
  withEmptyMetadata.attach();
}

@Before
public final void setUp() throws Exception {
  // Use a mutable service registry for later registering the service impl for each test case.
  fakeServer =
      InProcessServerBuilder.forName(fakeServerName)
          .fallbackHandlerRegistry(serviceRegistry)
          .directExecutor()
          .build()
          .start();
  Chunker.setDefaultChunkSizeForTesting(1000); // Enough for everything to be one chunk.
  fs = new InMemoryFileSystem(new JavaClock(), HashFunction.SHA256);
  execRoot = fs.getPath("/exec/root");
  FileSystemUtils.createDirectoryAndParents(execRoot);
  fakeFileCache = new FakeActionInputFileCache(execRoot);

  Path stdout = fs.getPath("/tmp/stdout");
  Path stderr = fs.getPath("/tmp/stderr");
  FileSystemUtils.createDirectoryAndParents(stdout.getParentDirectory());
  FileSystemUtils.createDirectoryAndParents(stderr.getParentDirectory());
  outErr = new FileOutErr(stdout, stderr);
  Context withEmptyMetadata =
      TracingMetadataUtils.contextWithMetadata(
          "none", "none", DIGEST_UTIL.asActionKey(Digest.getDefaultInstance()));
  withEmptyMetadata.attach();
}

/**
 * Default constructor.
 *
 * @param deviceId   the ONOS device id
 * @param p4DeviceId the P4 device id
 * @param channel    gRPC channel
 * @param controller runtime client controller
 */
P4RuntimeClientImpl(DeviceId deviceId, long p4DeviceId, ManagedChannel channel,
                    P4RuntimeControllerImpl controller) {
    this.deviceId = deviceId;
    this.p4DeviceId = p4DeviceId;
    this.controller = controller;
    this.cancellableContext = Context.current().withCancellation();
    this.executorService = Executors.newFixedThreadPool(15, groupedThreads(
            "onos/p4runtime-client-" + deviceId.toString(),
            deviceId.toString() + "-%d"));
    this.contextExecutor = this.cancellableContext.fixedContextExecutor(executorService);
    //TODO Investigate deadline or timeout in supplyInContext Method
    this.blockingStub = P4RuntimeGrpc.newBlockingStub(channel);
    P4RuntimeGrpc.P4RuntimeStub asyncStub = P4RuntimeGrpc.newStub(channel);
    this.streamRequestObserver = asyncStub.streamChannel(new StreamChannelResponseObserver());
}

/**
 * Construct a server.
 *
 * @param builder builder with configuration for server
 * @param transportServer transport server that will create new incoming transports
 * @param rootContext context that callbacks for new RPCs should be derived from
 */
ServerImpl(
    AbstractServerImplBuilder<?> builder,
    InternalServer transportServer,
    Context rootContext) {
  this.executorPool = Preconditions.checkNotNull(builder.executorPool, "executorPool");
  this.registry = Preconditions.checkNotNull(builder.registryBuilder.build(), "registryBuilder");
  this.fallbackRegistry =
      Preconditions.checkNotNull(builder.fallbackRegistry, "fallbackRegistry");
  this.transportServer = Preconditions.checkNotNull(transportServer, "transportServer");
  // Fork from the passed in context so that it does not propagate cancellation, it only
  // inherits values.
  this.rootContext = Preconditions.checkNotNull(rootContext, "rootContext").fork();
  this.decompressorRegistry = builder.decompressorRegistry;
  this.compressorRegistry = builder.compressorRegistry;
  this.transportFilters = Collections.unmodifiableList(
      new ArrayList<ServerTransportFilter>(builder.transportFilters));
  this.interceptors =
      builder.interceptors.toArray(new ServerInterceptor[builder.interceptors.size()]);
  this.handshakeTimeoutMillis = builder.handshakeTimeoutMillis;
}

/** Never returns {@code null}. */
private <ReqT, RespT> ServerStreamListener startCall(ServerStream stream, String fullMethodName,
    ServerMethodDefinition<ReqT, RespT> methodDef, Metadata headers,
    Context.CancellableContext context, StatsTraceContext statsTraceCtx) {
  // TODO(ejona86): should we update fullMethodName to have the canonical path of the method?
  ServerCallImpl<ReqT, RespT> call = new ServerCallImpl<ReqT, RespT>(
      stream, methodDef.getMethodDescriptor(), headers, context,
      decompressorRegistry, compressorRegistry);
  ServerCallHandler<ReqT, RespT> callHandler = methodDef.getServerCallHandler();
  statsTraceCtx.serverCallStarted(
      new ServerCallInfoImpl<ReqT, RespT>(
          methodDef.getMethodDescriptor(),
          call.getAttributes(),
          call.getAuthority()));

  for (ServerInterceptor interceptor : interceptors) {
    callHandler = InternalServerInterceptors.interceptCallHandler(interceptor, callHandler);
  }
  ServerCall.Listener<ReqT> listener = callHandler.startCall(call, headers);
  if (listener == null) {
    throw new NullPointerException(
        "startCall() returned a null listener for method " + fullMethodName);
  }
  return call.newServerStreamListener(listener);
}

ClientCallImpl(
    MethodDescriptor<ReqT, RespT> method, Executor executor, CallOptions callOptions,
    ClientTransportProvider clientTransportProvider,
    ScheduledExecutorService deadlineCancellationExecutor,
    CallTracer channelCallsTracer) {
  this.method = method;
  // If we know that the executor is a direct executor, we don't need to wrap it with a
  // SerializingExecutor. This is purely for performance reasons.
  // See https://github.com/grpc/grpc-java/issues/368
  this.callExecutor = executor == directExecutor()
      ? new SerializeReentrantCallsDirectExecutor()
      : new SerializingExecutor(executor);
  this.channelCallsTracer = channelCallsTracer;
  // Propagate the context from the thread which initiated the call to all callbacks.
  this.context = Context.current();
  this.unaryRequest = method.getType() == MethodType.UNARY
      || method.getType() == MethodType.SERVER_STREAMING;
  this.callOptions = callOptions;
  this.clientTransportProvider = clientTransportProvider;
  this.deadlineCancellationExecutor = deadlineCancellationExecutor;
}

@Test
public void contextCancellationCancelsStream() throws Exception {
  // Attach the context which is recorded when the call is created
  Context.CancellableContext cancellableContext = Context.current().withCancellation();
  Context previous = cancellableContext.attach();

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

  previous.attach();

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

  Throwable t = new Throwable();
  cancellableContext.cancel(t);

  verify(stream, times(1)).cancel(statusArgumentCaptor.capture());
  Status streamStatus = statusArgumentCaptor.getValue();
  assertEquals(Status.Code.CANCELLED, streamStatus.getCode());
}

@Test
public void contextDeadlineShouldBePropagatedInMetadata() {
  long deadlineNanos = TimeUnit.SECONDS.toNanos(1);
  Context context = Context.current().withDeadlineAfter(deadlineNanos, TimeUnit.NANOSECONDS,
      deadlineCancellationExecutor);
  context.attach();

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

  Metadata headers = new Metadata();

  call.start(callListener, headers);

  assertTrue(headers.containsKey(GrpcUtil.TIMEOUT_KEY));
  Long timeout = headers.get(GrpcUtil.TIMEOUT_KEY);
  assertNotNull(timeout);

  long deltaNanos = TimeUnit.MILLISECONDS.toNanos(400);
  assertTimeoutBetween(timeout, deadlineNanos - deltaNanos, deadlineNanos);
}

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

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

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

  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 testClientClose_cancelTriggersImmediateCancellation() throws Exception {
  AtomicBoolean contextCancelled = new AtomicBoolean(false);
  AtomicReference<Context> context = new AtomicReference<Context>();
  AtomicReference<ServerCall<String, Integer>> callReference
      = new AtomicReference<ServerCall<String, Integer>>();

  ServerStreamListener streamListener = testClientClose_setup(callReference,
      context, contextCancelled);

  // For close status being non OK:
  // isCancelled is expected to be true immediately after calling closed(), without needing
  // to wait for the main executor to run any tasks.
  assertFalse(callReference.get().isCancelled());
  assertFalse(context.get().isCancelled());
  streamListener.closed(Status.CANCELLED);
  assertEquals(1, executor.runDueTasks(CONTEXT_CLOSER_TASK_FITLER));
  assertTrue(callReference.get().isCancelled());
  assertTrue(context.get().isCancelled());

  assertEquals(1, executor.runDueTasks());
  assertTrue(contextCancelled.get());
}

@Test
public void testClientClose_OkTriggersDelayedCancellation() throws Exception {
  AtomicBoolean contextCancelled = new AtomicBoolean(false);
  AtomicReference<Context> context = new AtomicReference<Context>();
  AtomicReference<ServerCall<String, Integer>> callReference
      = new AtomicReference<ServerCall<String, Integer>>();

  ServerStreamListener streamListener = testClientClose_setup(callReference,
      context, contextCancelled);

  // For close status OK:
  // isCancelled is expected to be true after all pending work is done
  assertFalse(callReference.get().isCancelled());
  assertFalse(context.get().isCancelled());
  streamListener.closed(Status.OK);
  assertFalse(callReference.get().isCancelled());
  assertFalse(context.get().isCancelled());

  assertEquals(1, executor.runDueTasks());
  assertTrue(callReference.get().isCancelled());
  assertTrue(context.get().isCancelled());
  assertTrue(contextCancelled.get());
}

@Test
public void messageRead_errorCancelsCall() throws Exception {
  JumpToApplicationThreadServerStreamListener listener
      = new JumpToApplicationThreadServerStreamListener(
          executor.getScheduledExecutorService(),
          executor.getScheduledExecutorService(),
          stream,
          Context.ROOT.withCancellation());
  ServerStreamListener mockListener = mock(ServerStreamListener.class);
  listener.setListener(mockListener);

  TestError expectedT = new TestError();
  doThrow(expectedT).when(mockListener)
      .messagesAvailable(any(StreamListener.MessageProducer.class));
  // Closing the InputStream is done by the delegated listener (generally ServerCallImpl)
  listener.messagesAvailable(mock(StreamListener.MessageProducer.class));
  try {
    executor.runDueTasks();
    fail("Expected exception");
  } catch (TestError t) {
    assertSame(expectedT, t);
    ensureServerStateNotLeaked();
  }
}

@Test
public void messageRead_runtimeExceptionCancelsCall() throws Exception {
  JumpToApplicationThreadServerStreamListener listener
      = new JumpToApplicationThreadServerStreamListener(
          executor.getScheduledExecutorService(),
          executor.getScheduledExecutorService(),
          stream,
          Context.ROOT.withCancellation());
  ServerStreamListener mockListener = mock(ServerStreamListener.class);
  listener.setListener(mockListener);

  RuntimeException expectedT = new RuntimeException();
  doThrow(expectedT).when(mockListener)
      .messagesAvailable(any(StreamListener.MessageProducer.class));
  // Closing the InputStream is done by the delegated listener (generally ServerCallImpl)
  listener.messagesAvailable(mock(StreamListener.MessageProducer.class));
  try {
    executor.runDueTasks();
    fail("Expected exception");
  } catch (RuntimeException t) {
    assertSame(expectedT, t);
    ensureServerStateNotLeaked();
  }
}

@Test
public void halfClosed_errorCancelsCall() {
  JumpToApplicationThreadServerStreamListener listener
      = new JumpToApplicationThreadServerStreamListener(
          executor.getScheduledExecutorService(),
          executor.getScheduledExecutorService(),
          stream,
          Context.ROOT.withCancellation());
  ServerStreamListener mockListener = mock(ServerStreamListener.class);
  listener.setListener(mockListener);

  TestError expectedT = new TestError();
  doThrow(expectedT).when(mockListener).halfClosed();
  listener.halfClosed();
  try {
    executor.runDueTasks();
    fail("Expected exception");
  } catch (TestError t) {
    assertSame(expectedT, t);
    ensureServerStateNotLeaked();
  }
}

@Test
public void halfClosed_runtimeExceptionCancelsCall() {
  JumpToApplicationThreadServerStreamListener listener
      = new JumpToApplicationThreadServerStreamListener(
          executor.getScheduledExecutorService(),
          executor.getScheduledExecutorService(),
          stream,
          Context.ROOT.withCancellation());
  ServerStreamListener mockListener = mock(ServerStreamListener.class);
  listener.setListener(mockListener);

  RuntimeException expectedT = new RuntimeException();
  doThrow(expectedT).when(mockListener).halfClosed();
  listener.halfClosed();
  try {
    executor.runDueTasks();
    fail("Expected exception");
  } catch (RuntimeException t) {
    assertSame(expectedT, t);
    ensureServerStateNotLeaked();
  }
}

@Test
public void onReady_errorCancelsCall() {
  JumpToApplicationThreadServerStreamListener listener
      = new JumpToApplicationThreadServerStreamListener(
          executor.getScheduledExecutorService(),
          executor.getScheduledExecutorService(),
          stream,
          Context.ROOT.withCancellation());
  ServerStreamListener mockListener = mock(ServerStreamListener.class);
  listener.setListener(mockListener);

  TestError expectedT = new TestError();
  doThrow(expectedT).when(mockListener).onReady();
  listener.onReady();
  try {
    executor.runDueTasks();
    fail("Expected exception");
  } catch (TestError t) {
    assertSame(expectedT, t);
    ensureServerStateNotLeaked();
  }
}

@Test
public void onReady_runtimeExceptionCancelsCall() {
  JumpToApplicationThreadServerStreamListener listener
      = new JumpToApplicationThreadServerStreamListener(
          executor.getScheduledExecutorService(),
          executor.getScheduledExecutorService(),
          stream,
          Context.ROOT.withCancellation());
  ServerStreamListener mockListener = mock(ServerStreamListener.class);
  listener.setListener(mockListener);

  RuntimeException expectedT = new RuntimeException();
  doThrow(expectedT).when(mockListener).onReady();
  listener.onReady();
  try {
    executor.runDueTasks();
    fail("Expected exception");
  } catch (RuntimeException t) {
    assertSame(expectedT, t);
    ensureServerStateNotLeaked();
  }
}

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

@Test
public void serverTracingSampledToLocalSpanStore() {
  ServerStreamTracer.Factory tracerFactory = censusTracing.getServerTracerFactory();
  ServerStreamTracer serverStreamTracer =
      tracerFactory.newServerStreamTracer(sampledMethod.getFullMethodName(), new Metadata());

  serverStreamTracer.filterContext(Context.ROOT);

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

  serverStreamTracer.streamClosed(Status.CANCELLED);

  verify(spyServerSpan).end(
      EndSpanOptions.builder()
          .setStatus(io.opencensus.trace.Status.CANCELLED)
          .setSampleToLocalSpanStore(true)
          .build());
}

@Override
public <ReqT, RespT> Listener<ReqT> interceptCall(ServerCall<ReqT, RespT> call, Metadata headers, ServerCallHandler<ReqT, RespT> next) {
    String clientId = headers.get(CLIENT_ID_HEADER_KEY);
    if (clientId == null || !authenticator.authenticate(clientId)) {
        call.close(Status.UNAUTHENTICATED.withDescription("Invalid or unknown client: " + clientId), headers);
        return NOOP_LISTENER;
    }
    Context context = Context.current().withValue(CLIENT_ID_CONTEXT_KEY, clientId);
    return Contexts.interceptCall(context, call, headers, next);
}

/**
 * Attaches an empty ambient context to the provided gRPC {@code Context}.
 *
 * @throws IllegalStateException  if an ambient context has already been attached to the
 * provided gRPC {@code Context}.
 */
public static Context initialize(Context context) {
    checkNotNull(context, "context");
    checkState(DATA_KEY.get(context) == null,
            "AmbientContext has already been created in the scope of the current context");
    return context.withValue(DATA_KEY, new AmbientContext());
}

@Test
public void contextTransfersOneHopSync() throws Exception {
    Metadata.Key<String> ctxKey = Metadata.Key.of("ctx-context-key", Metadata.ASCII_STRING_MARSHALLER);
    String expectedCtxValue = "context-value";
    AtomicReference<String> ctxValue = new AtomicReference<>();

    // Service
    GreeterGrpc.GreeterImplBase svc = new GreeterGrpc.GreeterImplBase() {
        @Override
        public void sayHello(HelloRequest request, StreamObserver<HelloResponse> responseObserver) {
            ctxValue.set(AmbientContext.current().get(ctxKey));
            responseObserver.onNext(HelloResponse.newBuilder().setMessage("Hello " + request.getName()).build());
            responseObserver.onCompleted();
        }
    };

    // Plumbing
    serverRule1.getServiceRegistry().addService(ServerInterceptors
            .intercept(svc, new AmbientContextServerInterceptor("ctx-")));

    GreeterGrpc.GreeterBlockingStub stub = GreeterGrpc
            .newBlockingStub(serverRule1.getChannel())
            .withInterceptors(new AmbientContextClientInterceptor("ctx-"));

    // Test
    AmbientContext.initialize(Context.current()).run(() -> {
        AmbientContext.current().put(ctxKey, expectedCtxValue);
        stub.sayHello(HelloRequest.newBuilder().setName("world").build());
    });

    assertThat(ctxValue.get()).isEqualTo(expectedCtxValue);
}

@Test
public void multiValueContextTransfers() throws Exception {
    Metadata.Key<String> ctxKey = Metadata.Key.of("ctx-context-key", Metadata.ASCII_STRING_MARSHALLER);
    String expectedCtxValue1 = "context-value1";
    String expectedCtxValue2 = "context-value2";
    String expectedCtxValue3 = "context-value3";
    AtomicReference<Iterable<String>> ctxValue = new AtomicReference<>();

    // Service
    GreeterGrpc.GreeterImplBase svc = new GreeterGrpc.GreeterImplBase() {
        @Override
        public void sayHello(HelloRequest request, StreamObserver<HelloResponse> responseObserver) {
            ctxValue.set(AmbientContext.current().getAll(ctxKey));
            responseObserver.onNext(HelloResponse.newBuilder().setMessage("Hello " + request.getName()).build());
            responseObserver.onCompleted();
        }
    };

    // Plumbing
    serverRule1.getServiceRegistry().addService(ServerInterceptors
            .intercept(svc, new AmbientContextServerInterceptor("ctx-")));

    GreeterGrpc.GreeterBlockingStub stub = GreeterGrpc
            .newBlockingStub(serverRule1.getChannel())
            .withInterceptors(new AmbientContextClientInterceptor("ctx-"));

    // Test
    AmbientContext.initialize(Context.current()).run(() -> {
        AmbientContext.current().put(ctxKey, expectedCtxValue1);
        AmbientContext.current().put(ctxKey, expectedCtxValue2);
        AmbientContext.current().put(ctxKey, expectedCtxValue3);
        stub.sayHello(HelloRequest.newBuilder().setName("world").build());
    });

    assertThat(ctxValue.get()).containsExactlyInAnyOrder(expectedCtxValue1, expectedCtxValue2, expectedCtxValue3);
}

@Test
public void contextTransfersOneHopAsync() throws Exception {
    Metadata.Key<String> ctxKey = Metadata.Key.of("ctx-context-key", Metadata.ASCII_STRING_MARSHALLER);
    String expectedCtxValue = "context-value";
    AtomicReference<String> ctxValue = new AtomicReference<>();

    // Service
    GreeterGrpc.GreeterImplBase svc = new GreeterGrpc.GreeterImplBase() {
        @Override
        public void sayHello(HelloRequest request, StreamObserver<HelloResponse> responseObserver) {
            ctxValue.set(AmbientContext.current().get(ctxKey));
            responseObserver.onNext(HelloResponse.newBuilder().setMessage("Hello " + request.getName()).build());
            responseObserver.onCompleted();
        }
    };

    // Plumbing
    serverRule1.getServiceRegistry().addService(ServerInterceptors
            .intercept(svc, new AmbientContextServerInterceptor("ctx-")));
    GreeterGrpc.GreeterFutureStub stub = GreeterGrpc
            .newFutureStub(serverRule1.getChannel())
            .withInterceptors(new AmbientContextClientInterceptor("ctx-"));

    // Test
    AmbientContext.initialize(Context.current()).run(() -> {
        AmbientContext.current().put(ctxKey, expectedCtxValue);
        ListenableFuture<HelloResponse> futureResponse = stub.sayHello(HelloRequest.newBuilder().setName("world").build());

        // Verify response callbacks still have context
        MoreFutures.onSuccess(
                futureResponse,
                response -> assertThat(AmbientContext.current().get(ctxKey)).isEqualTo(expectedCtxValue),
                Context.currentContextExecutor(Executors.newSingleThreadExecutor()));

        await().atMost(Duration.ONE_SECOND).until(futureResponse::isDone);
    });

    assertThat(ctxValue.get()).isEqualTo(expectedCtxValue);
}

@Test
public void multipleContextTransfersOneHopSync() throws Exception {
    Metadata.Key<String> ctxKey = Metadata.Key.of("ctx-context-key", Metadata.ASCII_STRING_MARSHALLER);
    Metadata.Key<String> l5dKey = Metadata.Key.of("l5d-context-key", Metadata.ASCII_STRING_MARSHALLER);
    String expectedCtxValue = "context-value";
    AtomicReference<String> ctxValue = new AtomicReference<>();
    AtomicReference<String> l5dValue = new AtomicReference<>();

    // Service
    GreeterGrpc.GreeterImplBase svc = new GreeterGrpc.GreeterImplBase() {
        @Override
        public void sayHello(HelloRequest request, StreamObserver<HelloResponse> responseObserver) {
            ctxValue.set(AmbientContext.current().get(ctxKey));
            l5dValue.set(AmbientContext.current().get(l5dKey));
            responseObserver.onNext(HelloResponse.newBuilder().setMessage("Hello " + request.getName()).build());
            responseObserver.onCompleted();
        }
    };

    // Plumbing
    serverRule1.getServiceRegistry().addService(ServerInterceptors.intercept(svc,
            new AmbientContextServerInterceptor("ctx-"),
            new AmbientContextServerInterceptor("l5d-")));

    GreeterGrpc.GreeterBlockingStub stub = GreeterGrpc
            .newBlockingStub(serverRule1.getChannel())
            .withInterceptors(
                    new AmbientContextClientInterceptor("ctx-"),
                    new AmbientContextClientInterceptor("l5d-"));

    // Test
    AmbientContext.initialize(Context.current()).run(() -> {
        AmbientContext.current().put(ctxKey, expectedCtxValue);
        AmbientContext.current().put(l5dKey, expectedCtxValue);
        stub.sayHello(HelloRequest.newBuilder().setName("world").build());
    });

    assertThat(ctxValue.get()).isEqualTo(expectedCtxValue);
    assertThat(l5dValue.get()).isEqualTo(expectedCtxValue);
}