Java Code Examples for io.grpc.Context#CancellableContext

/** 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?
  statsTraceCtx.serverCallStarted(
      new ServerCallInfoImpl<ReqT, RespT>(
          methodDef.getMethodDescriptor(), // notify with original method descriptor
          stream.getAttributes(),
          stream.getAuthority()));
  ServerCallHandler<ReqT, RespT> handler = methodDef.getServerCallHandler();
  for (ServerInterceptor interceptor : interceptors) {
    handler = InternalServerInterceptors.interceptCallHandler(interceptor, handler);
  }
  ServerMethodDefinition<ReqT, RespT> interceptedDef = methodDef.withServerCallHandler(handler);
  ServerMethodDefinition<?, ?> wMethodDef = binlog == null
      ? interceptedDef : binlog.wrapMethodDefinition(interceptedDef);
  return startWrappedCall(fullMethodName, wMethodDef, stream, headers, context);
}
 

private <WReqT, WRespT> ServerStreamListener startWrappedCall(
    String fullMethodName,
    ServerMethodDefinition<WReqT, WRespT> methodDef,
    ServerStream stream,
    Metadata headers,
    Context.CancellableContext context) {
  ServerCallImpl<WReqT, WRespT> call = new ServerCallImpl<WReqT, WRespT>(
      stream,
      methodDef.getMethodDescriptor(),
      headers,
      context,
      decompressorRegistry,
      compressorRegistry,
      serverCallTracer);

  ServerCall.Listener<WReqT> listener =
      methodDef.getServerCallHandler().startCall(call, headers);
  if (listener == null) {
    throw new NullPointerException(
        "startCall() returned a null listener for method " + fullMethodName);
  }
  return call.newServerStreamListener(listener);
}
 

public static void main(String[] args) throws Exception {
  String host = args[0];
  String instanceName = args[1];
  DigestUtil digestUtil = DigestUtil.forHash(args[2]);
  String type = args[3];

  ScheduledExecutorService service = newSingleThreadScheduledExecutor();
  Context.CancellableContext ctx =
      Context.current()
          .withDeadlineAfter(deadlineSecondsForType(type), TimeUnit.SECONDS, service);
  Context prevContext = ctx.attach();
  try {
    cancellableMain(host, instanceName, digestUtil, type, args);
  } finally {
    ctx.cancel(null);
    ctx.detach(prevContext);
    if (!shutdownAndAwaitTermination(service, 1, TimeUnit.SECONDS)) {
      throw new RuntimeException("could not shut down service");
    }
  }
}
 

private Context.CancellableContext createContext(
    Metadata headers, StatsTraceContext statsTraceCtx) {
  Long timeoutNanos = headers.get(TIMEOUT_KEY);

  Context baseContext =
      statsTraceCtx
          .serverFilterContext(rootContext)
          .withValue(io.grpc.InternalServer.SERVER_CONTEXT_KEY, ServerImpl.this);

  if (timeoutNanos == null) {
    return baseContext.withCancellation();
  }

  Context.CancellableContext context =
      baseContext.withDeadline(
          Deadline.after(timeoutNanos, NANOSECONDS, ticker),
          transport.getScheduledExecutorService());

  return context;
}
 

/** Never returns {@code null}. */
private <ReqT, RespT> ServerStreamListener startCall(ServerStream stream, String fullMethodName,
    ServerMethodDefinition<ReqT, RespT> methodDef, Metadata headers,
    Context.CancellableContext context, StatsTraceContext statsTraceCtx, Tag tag) {
  // TODO(ejona86): should we update fullMethodName to have the canonical path of the method?
  statsTraceCtx.serverCallStarted(
      new ServerCallInfoImpl<>(
          methodDef.getMethodDescriptor(), // notify with original method descriptor
          stream.getAttributes(),
          stream.getAuthority()));
  ServerCallHandler<ReqT, RespT> handler = methodDef.getServerCallHandler();
  for (ServerInterceptor interceptor : interceptors) {
    handler = InternalServerInterceptors.interceptCallHandler(interceptor, handler);
  }
  ServerMethodDefinition<ReqT, RespT> interceptedDef = methodDef.withServerCallHandler(handler);
  ServerMethodDefinition<?, ?> wMethodDef = binlog == null
      ? interceptedDef : binlog.wrapMethodDefinition(interceptedDef);
  return startWrappedCall(fullMethodName, wMethodDef, stream, headers, context, tag);
}
 

private <WReqT, WRespT> ServerStreamListener startWrappedCall(
    String fullMethodName,
    ServerMethodDefinition<WReqT, WRespT> methodDef,
    ServerStream stream,
    Metadata headers,
    Context.CancellableContext context,
    Tag tag) {

  ServerCallImpl<WReqT, WRespT> call = new ServerCallImpl<>(
      stream,
      methodDef.getMethodDescriptor(),
      headers,
      context,
      decompressorRegistry,
      compressorRegistry,
      serverCallTracer,
      tag);

  ServerCall.Listener<WReqT> listener =
      methodDef.getServerCallHandler().startCall(call, headers);
  if (listener == null) {
    throw new NullPointerException(
        "startCall() returned a null listener for method " + fullMethodName);
  }
  return call.newServerStreamListener(listener);
}
 

ServerCallImpl(ServerStream stream, MethodDescriptor<ReqT, RespT> method,
    Metadata inboundHeaders, Context.CancellableContext context,
    DecompressorRegistry decompressorRegistry, CompressorRegistry compressorRegistry,
    CallTracer serverCallTracer) {
  this.stream = stream;
  this.method = method;
  this.context = context;
  this.messageAcceptEncoding = inboundHeaders.get(MESSAGE_ACCEPT_ENCODING_KEY);
  this.decompressorRegistry = decompressorRegistry;
  this.compressorRegistry = compressorRegistry;
  this.serverCallTracer = serverCallTracer;
  this.serverCallTracer.reportCallStarted();
}
 

public JumpToApplicationThreadServerStreamListener(Executor executor,
    Executor cancelExecutor, ServerStream stream, Context.CancellableContext context) {
  this.callExecutor = executor;
  this.cancelExecutor = cancelExecutor;
  this.stream = stream;
  this.context = context;
}
 

@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,
      false /* retryEnabled */)
          .setDecompressorRegistry(decompressorRegistry);

  cancellableContext.detach(previous);

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

@Override
protected <T> ListenableFuture<T> expect(
    Digest digest, Parser<T> parser, Executor executor, RequestMetadata requestMetadata) {
  Context.CancellableContext withDeadline =
      Context.current().withDeadlineAfter(60, SECONDS, contextDeadlineScheduler);
  Context previousContext = withDeadline.attach();
  try {
    ListenableFuture<T> future = super.expect(digest, parser, executor, requestMetadata);
    future.addListener(() -> withDeadline.cancel(null), directExecutor());
    return future;
  } finally {
    withDeadline.detach(previousContext);
  }
}
 

ServerCallImpl(ServerStream stream, MethodDescriptor<ReqT, RespT> method,
    Metadata inboundHeaders, Context.CancellableContext context,
    DecompressorRegistry decompressorRegistry, CompressorRegistry compressorRegistry,
    CallTracer serverCallTracer, Tag tag) {
  this.stream = stream;
  this.method = method;
  this.context = context;
  this.messageAcceptEncoding = inboundHeaders.get(MESSAGE_ACCEPT_ENCODING_KEY);
  this.decompressorRegistry = decompressorRegistry;
  this.compressorRegistry = compressorRegistry;
  this.serverCallTracer = serverCallTracer;
  this.serverCallTracer.reportCallStarted();
  this.tag = tag;
}
 

public JumpToApplicationThreadServerStreamListener(Executor executor,
    Executor cancelExecutor, ServerStream stream, Context.CancellableContext context, Tag tag) {
  this.callExecutor = executor;
  this.cancelExecutor = cancelExecutor;
  this.stream = stream;
  this.context = context;
  this.tag = tag;
}
 

@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<>(
      method,
      new SerializingExecutor(Executors.newSingleThreadExecutor()),
      baseCallOptions,
      provider,
      deadlineCancellationExecutor,
      channelCallTracer,
      /* retryEnabled= */ false)
          .setDecompressorRegistry(decompressorRegistry);

  cancellableContext.detach(previous);

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

ContextCloser(Context.CancellableContext context, Throwable cause) {
  this.context = context;
  this.cause = cause;
}
 

ContextCloser(Context.CancellableContext context, Throwable cause) {
  this.context = context;
  this.cause = cause;
}
 

/**
 * Returns the context associated with this client.
 *
 * @return context
 */
protected Context.CancellableContext context() {
    return cancellableContext;
}