io.grpc.stub.StreamObserver Java Examples

@Test
public void unexpectedExceptionCanNotMatch() {
    GreeterGrpc.GreeterImplBase svc = new GreeterGrpc.GreeterImplBase() {
        @Override
        public void sayHello(HelloRequest request, StreamObserver<HelloResponse> responseObserver) {
            throw new ArithmeticException("Divide by zero");
        }
    };

    ServerInterceptor interceptor = new TransmitUnexpectedExceptionInterceptor().forExactType(NullPointerException.class);

    serverRule.getServiceRegistry().addService(ServerInterceptors.intercept(svc, interceptor));
    GreeterGrpc.GreeterBlockingStub stub = GreeterGrpc.newBlockingStub(serverRule.getChannel());

    assertThatThrownBy(() -> stub.sayHello(HelloRequest.newBuilder().setName("World").build()))
            .isInstanceOf(StatusRuntimeException.class)
            .matches(sre -> ((StatusRuntimeException) sre).getStatus().getCode().equals(Status.UNKNOWN.getCode()), "is Status.UNKNOWN")
            .hasMessageContaining("UNKNOWN");
}
 

/** Returns top level channel aka {@link io.grpc.ManagedChannel}. */
@Override
public void getTopChannels(
    GetTopChannelsRequest request, StreamObserver<GetTopChannelsResponse> responseObserver) {
  InternalChannelz.RootChannelList rootChannels
      = channelz.getRootChannels(request.getStartChannelId(), maxPageSize);

  GetTopChannelsResponse resp;
  try {
    resp = ChannelzProtoUtil.toGetTopChannelResponse(rootChannels);
  } catch (StatusRuntimeException e) {
    responseObserver.onError(e);
    return;
  }

  responseObserver.onNext(resp);
  responseObserver.onCompleted();
}
 

@Override
public void findAllInputsOutputs(
    FindAllInputsOutputs request,
    StreamObserver<FindAllInputsOutputs.Response> responseObserver) {
  QPSCountResource.inc();
  try {
    if (request.getItemsCount() == 0) {
      throw new ModelDBException("Items not specified", Code.INVALID_ARGUMENT);
    }
    try (RequestLatencyResource latencyResource =
        new RequestLatencyResource(ModelDBAuthInterceptor.METHOD_NAME.get())) {
      FindAllInputsOutputs.Response response = lineageDAO.findAllInputsOutputs(request);
      responseObserver.onNext(response);
      responseObserver.onCompleted();
    }
  } catch (Exception e) {
    ModelDBUtils.observeError(
        responseObserver, e, FindAllInputsOutputs.Response.getDefaultInstance());
  }
}
 

@Override
public void getQueryProfile(
  GetQueryProfileRequest request, StreamObserver<GetQueryProfileResponse> responseObserver) {
  try {
    QueryId queryId = request.getQueryId();
    Preconditions.checkNotNull(queryId);

    QueryProfile mergedProfile = fetchOrBuildMergedProfile(queryId);
    responseObserver.onNext(
        GetQueryProfileResponse.newBuilder().setProfile(mergedProfile).build());
    responseObserver.onCompleted();
  } catch (IllegalArgumentException e) {
    responseObserver.onError(
        Status.INVALID_ARGUMENT
            .withDescription("get query profile failed " + e.getMessage())
            .asRuntimeException());
  } catch (Exception ex) {
    logger.error("get query profile failed", ex);
    responseObserver.onError(
        Status.INTERNAL.withDescription(ex.getMessage()).asRuntimeException());
  }
}
 

@Test
public void recordsMultipleCalls() throws Throwable {
  startGrpcServer(CHEAP_METRICS);

  createGrpcBlockingStub().sayHello(REQUEST);
  createGrpcBlockingStub().sayHello(REQUEST);
  createGrpcBlockingStub().sayHello(REQUEST);

  StreamRecorder<HelloResponse> streamRecorder = StreamRecorder.create();
  StreamObserver<HelloRequest> requestStream =
      createGrpcStub().sayHelloBidiStream(streamRecorder);
  requestStream.onNext(REQUEST);
  requestStream.onNext(REQUEST);
  requestStream.onCompleted();
  streamRecorder.awaitCompletion();

  assertThat(findRecordedMetricOrThrow("grpc_server_started_total").samples).hasSize(2);
  assertThat(findRecordedMetricOrThrow("grpc_server_handled_total").samples).hasSize(2);
}
 

private InterruptingPredicate<QueueEntry> createOnMatch(
    Instance instance, StreamObserver<QueueEntry> responseObserver) {
  return (queueEntry) -> {
    try {
      responseObserver.onNext(queueEntry);
      responseObserver.onCompleted();
      return true;
    } catch (StatusRuntimeException e) {
      Status status = Status.fromThrowable(e);
      if (status.getCode() != Status.Code.CANCELLED) {
        responseObserver.onError(e);
      }
    }
    instance.putOperation(instance.getOperation(queueEntry.getExecuteEntry().getOperationName()));
    return false;
  };
}
 

@Override
public void findRepositories(
    FindRepositories request, StreamObserver<FindRepositories.Response> responseObserver) {
  QPSCountResource.inc();
  try {
    try (RequestLatencyResource latencyResource =
        new RequestLatencyResource(modelDBAuthInterceptor.getMethodName())) {
      FindRepositories.Response response = repositoryDAO.findRepositories(request);
      responseObserver.onNext(response);
      responseObserver.onCompleted();
    }
  } catch (Exception e) {
    ModelDBUtils.observeError(
        responseObserver, e, FindRepositories.Response.getDefaultInstance());
  }
}
 

@Override
public StreamObserver<Messages.StreamingInputCallRequest> streamingInputCall(
        StreamObserver<Messages.StreamingInputCallResponse> responseObserver) {
    List<String> list = new CopyOnWriteArrayList<>();
    return new StreamObserver<Messages.StreamingInputCallRequest>() {
        @Override
        public void onNext(Messages.StreamingInputCallRequest streamingInputCallRequest) {
            list.add(streamingInputCallRequest.getPayload().getBody().toStringUtf8());
        }

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

        @Override
        public void onCompleted() {
            assertThat(list).containsExactly("a", "b", "c", "d");
            responseObserver.onNext(Messages.StreamingInputCallResponse.newBuilder().build());
            responseObserver.onCompleted();
        }
    };
}
 

@Override
public void onTxEvent(GrpcTxEvent message, StreamObserver<GrpcAck> responseObserver) {
  boolean ok = txConsistentService.handle(new TxEvent(
      message.getServiceName(),
      message.getInstanceId(),
      new Date(),
      message.getGlobalTxId(),
      message.getLocalTxId(),
      message.getParentTxId().isEmpty() ? null : message.getParentTxId(),
      message.getType(),
      message.getCompensationMethod(),
      message.getTimeout(),
      message.getRetryMethod(),
      message.getForwardRetries(),
      message.getPayloads().toByteArray()
  ));

  responseObserver.onNext(ok ? ALLOW : REJECT);
  responseObserver.onCompleted();
}
 

@Override
public void nodeUnpublishVolume(NodeUnpublishVolumeRequest request,
    StreamObserver<NodeUnpublishVolumeResponse> responseObserver) {
  String umountCommand =
      String.format("fusermount -u %s", request.getTargetPath());
  LOG.info("Executing {}", umountCommand);

  try {
    executeCommand(umountCommand);

    responseObserver.onNext(NodeUnpublishVolumeResponse.newBuilder()
        .build());
    responseObserver.onCompleted();

  } catch (Exception e) {
    responseObserver.onError(e);
  }

}
 

@Test
public void cacheHitWithInlineOutput() throws Exception {
  serviceRegistry.addService(
      new ActionCacheImplBase() {
        @Override
        public void getActionResult(
            GetActionResultRequest request, StreamObserver<ActionResult> responseObserver) {
          responseObserver.onNext(
              ActionResult.newBuilder()
                  .addOutputFiles(DUMMY_OUTPUT)
                  .setStdoutRaw(ByteString.copyFromUtf8("stdout"))
                  .setStderrRaw(ByteString.copyFromUtf8("stderr"))
                  .build());
          responseObserver.onCompleted();
        }
      });

  FakeSpawnExecutionContext policy =
      new FakeSpawnExecutionContext(simpleSpawn, fakeFileCache, execRoot, outErr);
  SpawnResult result = client.exec(simpleSpawn, policy);
  assertThat(result.setupSuccess()).isTrue();
  assertThat(result.exitCode()).isEqualTo(0);
  assertThat(result.isCacheHit()).isTrue();
  assertThat(outErr.outAsLatin1()).isEqualTo("stdout");
  assertThat(outErr.errAsLatin1()).isEqualTo("stderr");
}
 

@Override
public void deleteLineage(
    DeleteLineage request, StreamObserver<DeleteLineage.Response> responseObserver) {
  QPSCountResource.inc();
  try {
    if (request.getInputCount() == 0 && request.getOutputCount() == 0) {
      throw new ModelDBException("Input and output not specified", Code.INVALID_ARGUMENT);
    } else {
      if (request.getInputCount() == 0) {
        throw new ModelDBException("Input not specified", Code.INVALID_ARGUMENT);
      } else if (request.getOutputCount() == 0) {
        throw new ModelDBException("Output not specified", Code.INVALID_ARGUMENT);
      }
    }
    try (RequestLatencyResource latencyResource =
        new RequestLatencyResource(ModelDBAuthInterceptor.METHOD_NAME.get())) {
      DeleteLineage.Response response = lineageDAO.deleteLineage(request);
      responseObserver.onNext(response);
      responseObserver.onCompleted();
    }
  } catch (Exception e) {
    ModelDBUtils.observeError(responseObserver, e, DeleteLineage.Response.getDefaultInstance());
  }
}
 

/**
 * Implements a unary → stream call as {@link Single} → {@link Flowable}, where the server responds with a
 * stream of messages.
 */
public static <TRequest, TResponse> Flowable<TResponse> oneToMany(
        final Single<TRequest> rxRequest,
        final BiConsumer<TRequest, StreamObserver<TResponse>> delegate,
        final CallOptions options) {
    try {

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

        return rxRequest
                .flatMapPublisher(new io.reactivex.functions.Function<TRequest, Publisher<? extends TResponse>>() {
                    @Override
                    public Publisher<? extends TResponse> apply(TRequest request) {
                        final RxClientStreamObserverAndPublisher<TResponse> consumerStreamObserver =
                            new RxClientStreamObserverAndPublisher<TResponse>(null, null, prefetch, lowTide);

                        delegate.accept(request, consumerStreamObserver);

                        return consumerStreamObserver;
                    }
                });
    } catch (Throwable throwable) {
        return Flowable.error(throwable);
    }
}
 

private <T> void throttle(String agentId, boolean postV09, String collectionType,
        StreamObserver<T> responseObserver, Runnable runnable) {
    Semaphore semaphore = throttlePerAgentId.getUnchecked(agentId);
    boolean acquired;
    try {
        acquired = semaphore.tryAcquire(1, MINUTES);
    } catch (InterruptedException e) {
        // probably shutdown requested
        responseObserver.onError(e);
        return;
    }
    if (!acquired) {
        logger.warn("{} - {} collection rejected due to backlog",
                getAgentIdForLogging(agentId, postV09), collectionType);
        responseObserver.onError(Status.RESOURCE_EXHAUSTED
                .withDescription("collection rejected due to backlog")
                .asRuntimeException());
        return;
    }
    try {
        runnable.run();
    } finally {
        semaphore.release();
    }
}
 

@Override
public CompletableFuture<LeaseKeepAliveResponse> keepAliveOnce(long leaseId) {
    CompletableFuture<LeaseKeepAliveResponse> future = new CompletableFuture<>();

    StreamObserver<LeaseKeepAliveRequest> requestObserver = Observers.observe(
        this.leaseStub::leaseKeepAlive,
        response -> future.complete(new LeaseKeepAliveResponse(response)),
        throwable -> future.completeExceptionally(toEtcdException(throwable)));

    // cancel grpc stream when leaseKeepAliveResponseCompletableFuture completes.
    CompletableFuture<LeaseKeepAliveResponse> answer = future
        .whenCompleteAsync((val, throwable) -> requestObserver.onCompleted(), connectionManager.getExecutorService());

    requestObserver.onNext(LeaseKeepAliveRequest.newBuilder().setID(leaseId).build());

    return answer;
}
 

@Override
public StreamObserver<Messages.StreamingOutputCallRequest> fullDuplexCall(
        StreamObserver<Messages.StreamingOutputCallResponse> responseObserver) {
    AtomicInteger counter = new AtomicInteger();
    return new StreamObserver<Messages.StreamingOutputCallRequest>() {
        @Override
        public void onNext(Messages.StreamingOutputCallRequest streamingOutputCallRequest) {
            Messages.Payload payload = streamingOutputCallRequest.getPayload();
            ByteString value = ByteString
                    .copyFromUtf8(payload.getBody().toStringUtf8() + counter.incrementAndGet());
            Messages.Payload resp = Messages.Payload.newBuilder().setBody(value).build();
            Messages.StreamingOutputCallResponse response = Messages.StreamingOutputCallResponse.newBuilder()
                    .setPayload(resp).build();
            responseObserver.onNext(response);
        }

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

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

@Override
public void consume(List<RemoteMessage> remoteMessages) {
    try {
        StreamObserver<RemoteMessage> streamObserver = createStreamObserver();
        for (RemoteMessage remoteMessage : remoteMessages) {
            remoteOutCounter.inc();
            streamObserver.onNext(remoteMessage);
        }
        streamObserver.onCompleted();
    } catch (Throwable t) {
        remoteOutErrorCounter.inc();
        logger.error(t.getMessage(), t);
    }
}
 

@Override
public Completable addLoadBalancer(AddLoadBalancerRequest addLoadBalancerRequest, CallMetadata callMetadata) {
    return validator.validateLoadBalancer(addLoadBalancerRequest.getLoadBalancerId().getId())
            .onErrorResumeNext(e -> Completable.error(TitusServiceException.invalidArgument(e.getMessage())))
            .andThen(createRequestCompletable(emitter -> {
                StreamObserver<Empty> streamObserver = GrpcUtil.createEmptyClientResponseObserver(emitter);
                createWrappedStub(client, callMetadata, configuration.getRequestTimeout()).addLoadBalancer(addLoadBalancerRequest, streamObserver);
            }, configuration.getRequestTimeout()));
}
 

@Override
public void getCustomerLabel(
    GetCustomerLabelRequest request, StreamObserver<CustomerLabel> responseObserver) {
  Object response = responses.remove();
  if (response instanceof CustomerLabel) {
    requests.add(request);
    responseObserver.onNext((CustomerLabel) response);
    responseObserver.onCompleted();
  } else if (response instanceof Exception) {
    responseObserver.onError((Exception) response);
  } else {
    responseObserver.onError(new IllegalArgumentException("Unrecognized response type"));
  }
}
 

@Override
public void acknowledge(AcknowledgeRequest request, StreamObserver<Empty> responseObserver) {
  logger.atFine().log("Acknowledging messages %s", request);
  SubscriptionManager subscriptionManager = subscriptions.get(request.getSubscription());
  if (subscriptionManager == null) {
    String message = request.getSubscription() + " is not a valid Subscription";
    logger.atWarning().log(message);
    responseObserver.onError(Status.NOT_FOUND.withDescription(message).asException());
  } else {
    List<String> ackIds = subscriptionManager.acknowledge(request.getAckIdsList());
    logger.atFine().log("Successfully acknowledged %d messages", ackIds.size());
    responseObserver.onNext(Empty.getDefaultInstance());
    responseObserver.onCompleted();
  }
}
 

/**
 * {@inheritDoc}
 */
@Override
public void ping(
    final PingRequest request,
    final StreamObserver<PongResponse> responseObserver
) {
    final PongResponse response = PongResponse.newBuilder()
        .setRequestId(request.getRequestId())
        .setTimestamp(Timestamps.fromMillis(System.currentTimeMillis()))
        .putServerMetadata(ServerMetadataKeys.SERVER_NAME, hostName)
        .build();
    responseObserver.onNext(response);
    responseObserver.onCompleted();

    final StringBuilder sb = new StringBuilder();

    sb.append(
        String.format(
            "Received ping with id: '%s' from client: '%s' and timestamp: '%s'. ",
            request.getRequestId(),
            request.getSourceName(),
            Timestamps.toString(request.getTimestamp())
        )
    );

    sb.append("Client metadata: [ ");
    for (final Map.Entry<String, String> clientMetadataEntry : request.getClientMetadataMap().entrySet()) {
        sb
            .append("{")
            .append(clientMetadataEntry.getKey())
            .append(" : ")
            .append(clientMetadataEntry.getValue())
            .append("}, ");
    }
    sb.append("]");

    log.info(sb.toString());
}
 

@Override
public void observeJobs(ObserveJobsQuery request, StreamObserver<JobChangeNotification> responseObserver) {
    Subscription subscription = jobServiceGateway.observeJobs(request, resolveCallMetadata()).subscribe(
            responseObserver::onNext,
            e -> safeOnError(logger, e, responseObserver),
            responseObserver::onCompleted
    );
    attachCancellingCallback(responseObserver, subscription);
}
 

/**
 * 用于性能测试的一个方法
 * <p>
 * 主要逻辑：字符串拼接
 * </p>
 */
@Override
public void performanceTest(TestRequest req, StreamObserver<TestReply> responseObserver) {
  int id = req.getId();
  long timestamp = req.getTimestamp();
  boolean vip = req.getVip();
  String name = req.getName();
  float price = req.getPrice();
  double summation = req.getSummation();
  Map<Integer, String> args = req.getArgsMap();
  List<String> messages = req.getMessagesList();

  boolean success = vip;

  // 字符串拼接，作为返回结果
  StringBuilder sb = new StringBuilder();
  sb.append(id).append(timestamp).append(name).append(price).append(summation);

  Set<Map.Entry<Integer, String>> entrySet = args.entrySet();
  for (Map.Entry<Integer, String> entry : entrySet) {
    sb.append(entry.getKey()).append(entry.getValue());
  }

  for (String msg : messages) {
    sb.append(msg);
  }

  String message = sb.toString();

  TestReply reply = TestReply.newBuilder().setSuccess(success).setMessage(message).build();
  responseObserver.onNext(reply);
  responseObserver.onCompleted();
}
 

@Test
public void emptyStream() throws Exception {
  StreamRecorder<StreamingOutputCallResponse> responseObserver = StreamRecorder.create();
  StreamObserver<StreamingOutputCallRequest> requestObserver
      = asyncStub.fullDuplexCall(responseObserver);
  requestObserver.onCompleted();
  responseObserver.awaitCompletion(operationTimeoutMillis(), TimeUnit.MILLISECONDS);
}
 

@Override
public void setType(FilterSetTypeRequest request, StreamObserver<FilterSetTypeResponse> responseObserver) {
    FilterEntity filter = getFilterEntity(request.getFilter());
    filterDao.updateSetFilterType(filter, request.getType());
    responseObserver.onNext(FilterSetTypeResponse.newBuilder().build());
    responseObserver.onCompleted();
}
 

@Override
public void noteTimestampFromWriter(Controller.TimestampFromWriter request, StreamObserver<Controller.TimestampResponse> responseObserver) {
    StreamInfo streamInfo = request.getPosition().getStreamInfo();
    log.info("noteWriterMark called for stream {}/{}, writer={} time={}", streamInfo.getScope(),
            streamInfo.getStream(), request.getWriter(), request.getTimestamp());
    authenticateExecuteAndProcessResults(() -> this.grpcAuthHelper.checkAuthorization(
            AuthResourceRepresentation.ofStreamInScope(streamInfo.getScope(), streamInfo.getStream()),
            AuthHandler.Permissions.READ_UPDATE),
            delegationToken  -> controllerService.noteTimestampFromWriter(streamInfo.getScope(),
                    streamInfo.getStream(), request.getWriter(), request.getTimestamp(), request.getPosition().getCutMap()),
            responseObserver);
}
 

@Override
public void setThreadable(LayerSetThreadableRequest request, StreamObserver<LayerSetThreadableResponse> responseObserver) {
    updateLayer(request.getLayer());
    layerDao.updateThreadable(layer, request.getThreadable());
    responseObserver.onNext(LayerSetThreadableResponse.newBuilder().build());
    responseObserver.onCompleted();
}
 

@Override
public void shareAlbum(
    ShareAlbumRequest request, StreamObserver<ShareAlbumResponse> responseObserver) {
  Object response = responses.remove();
  if (response instanceof ShareAlbumResponse) {
    requests.add(request);
    responseObserver.onNext((ShareAlbumResponse) response);
    responseObserver.onCompleted();
  } else if (response instanceof Exception) {
    responseObserver.onError((Exception) response);
  } else {
    responseObserver.onError(new IllegalArgumentException("Unrecognized response type"));
  }
}
 

@Override
public void commandActiveThreadDump(PCmdActiveThreadDumpRes activeThreadDumpRes, StreamObserver<Empty> responseObserver) {
    final Long transportId = getTransportId();
    PinpointGrpcServer pinpointGrpcServer = grpcServerRepository.get(transportId);
    if (pinpointGrpcServer != null) {
        activeThreadDumpService.handle(pinpointGrpcServer, activeThreadDumpRes, responseObserver);
        responseObserver.onNext(Empty.getDefaultInstance());
        responseObserver.onCompleted();
    } else {
        logger.info("{} => local. Can't find PinpointGrpcServer(transportId={})", getAgentInfo().getAgentKey(), transportId);
        responseObserver.onError(new StatusException(Status.NOT_FOUND));
    }
}
 

@Override
public void getMempoolTransactionList(RequestAddress req, StreamObserver<TransactionHashList> observer)
{
  AddressSpecHash spec_hash = new AddressSpecHash(req.getAddressSpecHash());

  TransactionHashList.Builder list = TransactionHashList.newBuilder();
  for(ChainHash h : node.getMemPool().getTransactionsForAddress(spec_hash))
  {
    list.addTxHashes(h.getBytes());
  }

  observer.onNext( list.build());
  observer.onCompleted();
}