io.grpc.stub.ServerCallStreamObserver Java Examples

/**
 * Implements a bidirectional stream → stream call as {@link Flowable} → {@link Flowable}, where both the client
 * and the server independently stream to each other.
 */
public static <TRequest, TResponse> StreamObserver<TRequest> manyToMany(
        final StreamObserver<TResponse> responseObserver,
        final Function<Flowable<TRequest>, Flowable<TResponse>> delegate,
        final CallOptions options) {

    final int prefetch = RxCallOptions.getPrefetch(options);
    final int lowTide = RxCallOptions.getLowTide(options);

    final RxServerStreamObserverAndPublisher<TRequest> streamObserverPublisher =
            new RxServerStreamObserverAndPublisher<TRequest>((ServerCallStreamObserver<TResponse>) responseObserver, null, prefetch, lowTide);

    try {
        final Flowable<TResponse> rxResponse = Preconditions.checkNotNull(delegate.apply(Flowable.fromPublisher(streamObserverPublisher)));
        final RxSubscriberAndServerProducer<TResponse> subscriber = new RxSubscriberAndServerProducer<TResponse>();
        subscriber.subscribe((ServerCallStreamObserver<TResponse>) responseObserver);
        // Don't try to respond if the server has already canceled the request
        rxResponse.subscribe(subscriber);
    } catch (Throwable throwable) {
        responseObserver.onError(prepareError(throwable));
    }

    return streamObserverPublisher;
}
 

@Override
public void observeEvents(Empty request, StreamObserver<SupervisorEvent> responseObserver) {
    Subscription subscription = supervisorOperations.events()
            .map(SupervisorGrpcModelConverters::toGrpcEvent)
            .subscribe(
                    responseObserver::onNext,
                    e -> responseObserver.onError(
                            new StatusRuntimeException(Status.INTERNAL
                                    .withDescription("Supervisor event stream terminated with an error")
                                    .withCause(e))
                    ),
                    responseObserver::onCompleted
            );
    ServerCallStreamObserver<SupervisorEvent> serverObserver = (ServerCallStreamObserver<SupervisorEvent>) responseObserver;
    serverObserver.setOnCancelHandler(subscription::unsubscribe);
}
 

/**
 * Implements a bidirectional stream → stream call as {@link Flux} → {@link Flux}, where both the client
 * and the server independently stream to each other.
 */
public static <TRequest, TResponse> StreamObserver<TRequest> manyToMany(
        StreamObserver<TResponse> responseObserver,
        Function<Flux<TRequest>, Flux<TResponse>> delegate,
        CallOptions options) {

    final int prefetch = ReactorCallOptions.getPrefetch(options);
    final int lowTide = ReactorCallOptions.getLowTide(options);

    ReactorServerStreamObserverAndPublisher<TRequest> streamObserverPublisher =
            new ReactorServerStreamObserverAndPublisher<>((ServerCallStreamObserver<TResponse>) responseObserver, null, prefetch, lowTide);

    try {
        Flux<TResponse> rxResponse = Preconditions.checkNotNull(delegate.apply(Flux.from(streamObserverPublisher)));
        ReactorSubscriberAndServerProducer<TResponse> subscriber = new ReactorSubscriberAndServerProducer<>();
        subscriber.subscribe((ServerCallStreamObserver<TResponse>) responseObserver);
        // Don't try to respond if the server has already canceled the request
        rxResponse.subscribe(subscriber);
    } catch (Throwable throwable) {
        responseObserver.onError(prepareError(throwable));
    }

    return streamObserverPublisher;
}
 

@Override
public void waitExecution(
    WaitExecutionRequest request, StreamObserver<Operation> responseObserver) {
  String operationName = request.getName();
  Instance instance;
  try {
    instance = instances.getFromOperationName(operationName);
  } catch (InstanceNotFoundException e) {
    responseObserver.onError(BuildFarmInstances.toStatusException(e));
    return;
  }

  ServerCallStreamObserver<Operation> serverCallStreamObserver =
      (ServerCallStreamObserver<Operation>) responseObserver;
  withCancellation(
      serverCallStreamObserver,
      instance.watchOperation(
          operationName,
          createWatcher(serverCallStreamObserver, TracingMetadataUtils.fromCurrentContext())));
}
 

@Override
public void streamingFromServer(Messages.SimpleRequest request, StreamObserver<Messages.SimpleResponse> responseObserver) {
    final ServerCallStreamObserver<Messages.SimpleResponse> callStreamObserver = (ServerCallStreamObserver<Messages.SimpleResponse>) responseObserver;

    for (Messages.SimpleResponse response : responses) {
        if (callStreamObserver.isCancelled()) {
            return;
        }

        callStreamObserver.onNext(response);
    }

    if (!callStreamObserver.isCancelled()) {
        callStreamObserver.onCompleted();
    }
}
 

void readLimitedBlob(
    Instance instance,
    Digest digest,
    long offset,
    long limit,
    StreamObserver<ReadResponse> responseObserver) {
  ServerCallStreamObserver<ReadResponse> target =
      onErrorLogReadObserver(
          format("%s(%s)", DigestUtil.toString(digest), instance.getName()),
          offset,
          (ServerCallStreamObserver<ReadResponse>) responseObserver);
  try {
    instance.getBlob(
        digest,
        offset,
        limit,
        newChunkObserver(target),
        TracingMetadataUtils.fromCurrentContext());
  } catch (Exception e) {
    target.onError(e);
  }
}
 

/**
 * Implements a unary → stream call as {@link Single} → {@link Flowable}, where the server responds with a
 * stream of messages.
 */
public static <TRequest, TResponse> void oneToMany(
        final TRequest request,
        final StreamObserver<TResponse> responseObserver,
        final Function<Single<TRequest>, Flowable<TResponse>> delegate) {
    try {
        final Single<TRequest> rxRequest = Single.just(request);

        final Flowable<TResponse> rxResponse = Preconditions.checkNotNull(delegate.apply(rxRequest));
        final RxSubscriberAndServerProducer<TResponse> serverProducer =
                rxResponse.subscribeWith(new RxSubscriberAndServerProducer<TResponse>());
        serverProducer.subscribe((ServerCallStreamObserver<TResponse>) responseObserver);
    } catch (Throwable throwable) {
        responseObserver.onError(prepareError(throwable));
    }
}
 

StreamingPullStreamObserver(StreamObserver<StreamingPullResponse> responseObserver) {
  // Upcast to a ServerCallStreamObserver to set manual flow control
  this.responseObserver = (ServerCallStreamObserver<StreamingPullResponse>) responseObserver;
  this.responseObserver.disableAutoInboundFlowControl();

  this.responseObserver.setOnReadyHandler(
      () -> {
        if (isNull(subscriptionManager)) {
          this.responseObserver.request(1);
        }
      });
  this.responseObserver.setOnCancelHandler(
      () -> {
        logger.atInfo().log("Client cancelled StreamingPull %s", streamId);
        stopIfRunning();
      });
}
 

@Override
public void get(
    Digest digest,
    long offset,
    long limit,
    ServerCallStreamObserver<ByteString> blobObserver,
    RequestMetadata requestMetadata) {
  Blob blob = get(digest);
  if (blob == null) {
    if (delegate != null) {
      // FIXME change this to a read-through get
      delegate.get(digest, offset, limit, blobObserver, requestMetadata);
    } else {
      blobObserver.onError(io.grpc.Status.NOT_FOUND.asException());
    }
  } else {
    blobObserver.onNext(blob.getData());
    blobObserver.onCompleted();
  }
}
 

private void provideBlob(Digest digest, ByteString content) {
  blobDigests.add(digest);
  // FIXME use better answer definitions, without indexes
  doAnswer(
          new Answer<Void>() {
            @Override
            public Void answer(InvocationOnMock invocation) {
              StreamObserver<ByteString> blobObserver =
                  (StreamObserver) invocation.getArguments()[3];
              blobObserver.onNext(content);
              blobObserver.onCompleted();
              return null;
            }
          })
      .when(mockWorkerInstance)
      .getBlob(
          eq(digest),
          eq(0l),
          eq(digest.getSizeBytes()),
          any(ServerCallStreamObserver.class),
          any(RequestMetadata.class));
}
 

@Override
public void onNext(T value) {
  synchronized (lock) {
    // in theory we could implement ServerCallStreamObserver and expose isCancelled to our client,
    // but for current purposes we only need the StreamObserver API, so treat a cancelled observer
    // as something we just want to un-block from and return, and we'll trust the rest of our session
    // to shutdown accordingly.
    if (delegate instanceof ServerCallStreamObserver && ((ServerCallStreamObserver<T>) delegate).isCancelled()) {
      return;
    }
    while (!delegate.isReady()) {
      try {
        lock.wait();
      } catch (InterruptedException e) {
        Thread.currentThread().interrupt();
        return;
      }
    }
  }
  delegate.onNext(value);
}
 

@Override
public void get(
    Digest digest,
    long offset,
    long count,
    ServerCallStreamObserver<ByteString> blobObserver,
    RequestMetadata requestMetadata) {
  ReadRequest request =
      ReadRequest.newBuilder()
          .setResourceName(getBlobName(digest))
          .setReadOffset(offset)
          .setReadLimit(count)
          .build();
  ByteStreamGrpc.newStub(channel)
      .withInterceptors(attachMetadataInterceptor(requestMetadata))
      .read(
          request,
          new DelegateServerCallStreamObserver<ReadResponse, ByteString>(blobObserver) {
            @Override
            public void onNext(ReadResponse response) {
              blobObserver.onNext(response.getData());
            }

            @Override
            public void onError(Throwable t) {
              blobObserver.onError(t);
            }

            @Override
            public void onCompleted() {
              blobObserver.onCompleted();
            }
          });
}
 

RxServerStreamObserverAndPublisher(
        ServerCallStreamObserver<?> serverCallStreamObserver,
        Consumer<CallStreamObserver<?>> onSubscribe,
        int prefetch,
        int lowTide) {
    super(serverCallStreamObserver, new SimpleQueueAdapter<T>(new SpscArrayQueue<T>(prefetch)), onSubscribe, prefetch, lowTide);
}
 

/**
 * Immediately responds with a payload of the type and size specified in the request.
 */
@Override
public void unaryCall(SimpleRequest req, StreamObserver<SimpleResponse> responseObserver) {
  ServerCallStreamObserver<SimpleResponse> obs =
      (ServerCallStreamObserver<SimpleResponse>) responseObserver;
  SimpleResponse.Builder responseBuilder = SimpleResponse.newBuilder();
  try {
    if (req.hasResponseCompressed() && req.getResponseCompressed().getValue()) {
      obs.setCompression("gzip");
    } else {
      obs.setCompression("identity");
    }
  } catch (IllegalArgumentException e) {
    obs.onError(Status.UNIMPLEMENTED
        .withDescription("compression not supported.")
        .withCause(e)
        .asRuntimeException());
    return;
  }

  if (req.getResponseSize() != 0) {
    // For consistency with the c++ TestServiceImpl, use a random offset for unary calls.
    // TODO(wonderfly): whether or not this is a good approach needs further discussion.
    int offset = random.nextInt(compressableBuffer.size());
    ByteString payload = generatePayload(compressableBuffer, offset, req.getResponseSize());
    responseBuilder.setPayload(
        Payload.newBuilder()
            .setBody(payload));
  }

  if (req.hasResponseStatus()) {
    obs.onError(Status.fromCodeValue(req.getResponseStatus().getCode())
        .withDescription(req.getResponseStatus().getMessage())
        .asRuntimeException());
    return;
  }

  responseObserver.onNext(responseBuilder.build());
  responseObserver.onCompleted();
}
 

@Override
public void sayHello(HelloRequest req, StreamObserver<HelloReply> plainResponseObserver) {
  ServerCallStreamObserver<HelloReply> responseObserver =
      (ServerCallStreamObserver<HelloReply>) plainResponseObserver;
  /* This line by itself enables compression for the response */
  responseObserver.setCompression("gzip");
  HelloReply reply = HelloReply.newBuilder().setMessage("Hello " + req.getName()).build();
  responseObserver.onNext(reply);
  responseObserver.onCompleted();
}
 

@Override
public void getBlob(
    Digest blobDigest,
    long offset,
    long count,
    ServerCallStreamObserver<ByteString> blobObserver,
    RequestMetadata requestMetadata) {
  contentAddressableStorage.get(blobDigest, offset, count, blobObserver, requestMetadata);
}
 

/** Retrieve a value from the CAS by streaming content when ready */
@ThreadSafe
void get(
    Digest digest,
    long offset,
    long count,
    ServerCallStreamObserver<ByteString> blobObserver,
    RequestMetadata requestMetadata);
 

private void withCancellation(
    ServerCallStreamObserver<Operation> serverCallStreamObserver, ListenableFuture<Void> future) {
  addCallback(
      future,
      new FutureCallback<Void>() {
        boolean isCancelled() {
          return serverCallStreamObserver.isCancelled() || Context.current().isCancelled();
        }

        @Override
        public void onSuccess(Void result) {
          if (!isCancelled()) {
            try {
              serverCallStreamObserver.onCompleted();
            } catch (Exception e) {
              onFailure(e);
            }
          }
        }

        @Override
        public void onFailure(Throwable t) {
          if (!isCancelled() && !(t instanceof CancellationException)) {
            logger.log(Level.WARNING, "error occurred during execution", t);
            serverCallStreamObserver.onError(Status.fromThrowable(t).asException());
          }
        }
      },
      Context.current().fixedContextExecutor(directExecutor()));
  serverCallStreamObserver.setOnCancelHandler(() -> future.cancel(false));
}
 

public AbstractServerStreamObserverAndPublisher(
    ServerCallStreamObserver<?> serverCallStreamObserver,
    Queue<T> queue,
    Consumer<CallStreamObserver<?>> onSubscribe) {
    super(queue, onSubscribe);
    super.onSubscribe(serverCallStreamObserver);
}
 

/**
 * Immediately responds with a payload of the type and size specified in the request.
 */
@Override
public void unaryCall(SimpleRequest req, StreamObserver<SimpleResponse> responseObserver) {
  ServerCallStreamObserver<SimpleResponse> obs =
      (ServerCallStreamObserver<SimpleResponse>) responseObserver;
  SimpleResponse.Builder responseBuilder = SimpleResponse.newBuilder();
  try {
    if (req.hasResponseCompressed() && req.getResponseCompressed().getValue()) {
      obs.setCompression("gzip");
    } else {
      obs.setCompression("identity");
    }
  } catch (IllegalArgumentException e) {
    obs.onError(Status.UNIMPLEMENTED
        .withDescription("compression not supported.")
        .withCause(e)
        .asRuntimeException());
    return;
  }

  if (req.getResponseSize() != 0) {
    // For consistency with the c++ TestServiceImpl, use a random offset for unary calls.
    // TODO(wonderfly): whether or not this is a good approach needs further discussion.
    int offset = random.nextInt(compressableBuffer.size());
    ByteString payload = generatePayload(compressableBuffer, offset, req.getResponseSize());
    responseBuilder.setPayload(
        Payload.newBuilder()
            .setBody(payload));
  }

  if (req.hasResponseStatus()) {
    obs.onError(Status.fromCodeValue(req.getResponseStatus().getCode())
        .withDescription(req.getResponseStatus().getMessage())
        .asRuntimeException());
    return;
  }

  responseObserver.onNext(responseBuilder.build());
  responseObserver.onCompleted();
}
 

@Override
public void run(final RunRequest request, final StreamObserver<RunResponse> observer) {
  // Switch to our own threads so that onReadyStateHandler can be called (see class-level
  // comment).
  ServerCallStreamObserver<RunResponse> serverCallStreamObserver =
      ((ServerCallStreamObserver<RunResponse>) observer);
  BlockingStreamObserver<RunResponse> blockingStreamObserver =
      new BlockingStreamObserver<>(serverCallStreamObserver);
  commandExecutorPool.execute(() -> executeCommand(request, blockingStreamObserver));
}
 

KeepaliveWatcher createWatcher(
    ServerCallStreamObserver<Operation> serverCallStreamObserver,
    RequestMetadata requestMetadata) {
  return new KeepaliveWatcher(serverCallStreamObserver) {
    @Override
    void deliver(Operation operation) {
      if (operation != null) {
        metricsPublisher.publishRequestMetadata(operation, requestMetadata);
      }
      serverCallStreamObserver.onNext(operation);
    }
  };
}
 

ServerCallStreamObserver<ReadResponse> onErrorLogReadObserver(
    String name, long offset, ServerCallStreamObserver<ReadResponse> delegate) {
  return new UniformDelegateServerCallStreamObserver<ReadResponse>(delegate) {
    long responseCount = 0;
    long responseBytes = 0;

    @Override
    public void onNext(ReadResponse response) {
      delegate.onNext(response);
      responseCount++;
      responseBytes += response.getData().size();
    }

    @Override
    public void onError(Throwable t) {
      Status status = Status.fromThrowable(t);
      if (status.getCode() != Code.NOT_FOUND) {
        java.util.logging.Level level = Level.SEVERE;
        if (responseCount > 0 && status.getCode() == Code.DEADLINE_EXCEEDED
            || status.getCode() == Code.CANCELLED) {
          level = Level.WARNING;
        }
        String message = format("error reading %s at offset %d", name, offset);
        if (responseCount > 0) {
          message +=
              format(" after %d responses and %d bytes of content", responseCount, responseBytes);
        }
        logger.log(level, message, t);
      }
      delegate.onError(t);
    }

    @Override
    public void onCompleted() {
      delegate.onCompleted();
    }
  };
}
 

ServerCallStreamObserver<ByteString> newChunkObserver(
    ServerCallStreamObserver<ReadResponse> responseObserver) {
  return new DelegateServerCallStreamObserver<ByteString, ReadResponse>(responseObserver) {
    @Override
    public void onNext(ByteString data) {
      while (!data.isEmpty()) {
        ByteString slice;
        if (data.size() > CHUNK_SIZE) {
          slice = data.substring(0, CHUNK_SIZE);
          data = data.substring(CHUNK_SIZE);
        } else {
          slice = data;
          data = ByteString.EMPTY;
        }
        responseObserver.onNext(ReadResponse.newBuilder().setData(slice).build());
      }
    }

    @Override
    public void onError(Throwable t) {
      responseObserver.onError(t);
    }

    @Override
    public void onCompleted() {
      responseObserver.onCompleted();
    }
  };
}
 

private ServerCallStreamObserver<WriteResponse> initializeBackPressure(
    StreamObserver<WriteResponse> responseObserver) {
  final ServerCallStreamObserver<WriteResponse> serverCallStreamObserver =
      (ServerCallStreamObserver<WriteResponse>) responseObserver;
  serverCallStreamObserver.disableAutoInboundFlowControl();
  serverCallStreamObserver.request(1);
  return serverCallStreamObserver;
}
 

@Override
public StreamObserver<WriteRequest> write(StreamObserver<WriteResponse> responseObserver) {
  ServerCallStreamObserver<WriteResponse> serverCallStreamObserver =
      initializeBackPressure(responseObserver);
  return new WriteStreamObserver(
      instances,
      deadlineAfter,
      deadlineAfterUnits,
      () -> serverCallStreamObserver.request(1),
      responseObserver);
}
 

public ActiveThreadCountStreamObserver(PinpointGrpcServer pinpointGrpcServer, StreamObserver<Empty> connectionObserver) {
    this.pinpointGrpcServer = Objects.requireNonNull(pinpointGrpcServer, "pinpointGrpcServer");
    if (connectionObserver instanceof ServerCallStreamObserver) {
        this.connectionObserver = (ServerCallStreamObserver<Empty>) connectionObserver;
    } else {
        throw new IllegalArgumentException("streamConnectionManagerObserver can not cast to ServerCallStreamObserver");
    }
}
 

BlockingStreamObserver(CallStreamObserver<T> delegate) {
  this.delegate = delegate;
  final Runnable notifyAll = () -> {
    synchronized (lock) {
      lock.notifyAll(); // wake up our thread
    }
  };
  this.delegate.setOnReadyHandler(notifyAll);
  if (delegate instanceof ServerCallStreamObserver) {
    ((ServerCallStreamObserver<T>) delegate).setOnCancelHandler(notifyAll);
  }
}
 

@Override
public void staticUnaryCallSetsMessageCompression(SimpleRequest request,
                                                  StreamObserver<SimpleResponse> responseObserver) {
    if (!request.equals(REQUEST_MESSAGE)) {
        responseObserver.onError(new IllegalArgumentException("Unexpected request: " + request));
        return;
    }
    final ServerCallStreamObserver<SimpleResponse> callObserver =
            (ServerCallStreamObserver<SimpleResponse>) responseObserver;
    callObserver.setCompression("gzip");
    callObserver.setMessageCompression(true);
    responseObserver.onNext(RESPONSE_MESSAGE);
    responseObserver.onCompleted();
}
 

@Override
public StreamObserver<SimpleRequest> serverBackPressure(
        StreamObserver<SimpleResponse> rawResponseObserver) {
    final ServerCallStreamObserver<SimpleResponse> responseObserver =
            (ServerCallStreamObserver<SimpleResponse>) rawResponseObserver;
    responseObserver.disableAutoInboundFlowControl();
    final AtomicInteger numRequests = new AtomicInteger();
    responseObserver.request(1);
    return new StreamObserver<SimpleRequest>() {
        @Override
        public void onNext(SimpleRequest value) {
            if (numRequests.incrementAndGet() < CAPPED_NUM_MESSAGES) {
                responseObserver.request(1);
            } else {
                for (int i = 0; i < TOTAL_NUM_MESSAGES; i++) {
                    responseObserver.onNext(SimpleResponse.getDefaultInstance());
                }
                responseObserver.onCompleted();
            }
        }

        @Override
        public void onError(Throwable t) {
        }

        @Override
        public void onCompleted() {
        }
    };
}