GrpcWindmillServer extends WindmillServerStub { @Override public CommitWorkStream commitWorkStream() { return new GrpcCommitWorkStream(); } GrpcWindmillServer(StreamingDataflowWorkerOptions options); private  GrpcWindmillServer(String name, boolean enableStreamingEngine); @Override synchronized void setWindmillServiceEndpoints(Set<HostAndPort> endpoints); @Override synchronized boolean isReady(); @Override void appendSummaryHtml(PrintWriter writer); @Override GetWorkResponse getWork(GetWorkRequest request); @Override GetDataResponse getData(GetDataRequest request); @Override CommitWorkResponse commitWork(CommitWorkRequest request); @Override GetWorkStream getWorkStream(GetWorkRequest request, WorkItemReceiver receiver); @Override GetDataStream getDataStream(); @Override CommitWorkStream commitWorkStream(); @Override GetConfigResponse getConfig(GetConfigRequest request); @Override ReportStatsResponse reportStats(ReportStatsRequest request); @Override long getAndResetThrottleTime();  }
@Test public void testStreamingCommit() throws Exception { List<WorkItemCommitRequest> commitRequestList = new ArrayList<>(); List<CountDownLatch> latches = new ArrayList<>(); Map<Long, WorkItemCommitRequest> commitRequests = new HashMap<>(); for (int i = 0; i < 500; ++i) { WorkItemCommitRequest request = makeCommitRequest(i, i * (i < 480 ? 8 : 128)); commitRequestList.add(request); commitRequests.put((long) i, request); latches.add(new CountDownLatch(1)); } Collections.shuffle(commitRequestList); serviceRegistry.addService( new CloudWindmillServiceV1Alpha1ImplBase() { @Override public StreamObserver<StreamingCommitWorkRequest> commitWorkStream( StreamObserver<StreamingCommitResponse> responseObserver) { return new StreamObserver<StreamingCommitWorkRequest>() { boolean sawHeader = false; InputStream buffer = null; long remainingBytes = 0; ResponseErrorInjector injector = new ResponseErrorInjector(responseObserver); @Override public void onNext(StreamingCommitWorkRequest request) { maybeInjectError(responseObserver); if (!sawHeader) { errorCollector.checkThat( request.getHeader(), Matchers.equalTo( JobHeader.newBuilder() .setJobId("job") .setProjectId("project") .setWorkerId("worker") .build())); sawHeader = true; LOG.info("Received header"); } else { boolean first = true; LOG.info("Received request with {} chunks", request.getCommitChunkCount()); for (StreamingCommitRequestChunk chunk : request.getCommitChunkList()) { assertTrue(chunk.getSerializedWorkItemCommit().size() <= STREAM_CHUNK_SIZE); if (first || chunk.hasComputationId()) { errorCollector.checkThat( chunk.getComputationId(), Matchers.equalTo("computation")); } if (remainingBytes != 0) { errorCollector.checkThat(buffer, Matchers.notNullValue()); errorCollector.checkThat( remainingBytes, Matchers.is( chunk.getSerializedWorkItemCommit().size() + chunk.getRemainingBytesForWorkItem())); buffer = new SequenceInputStream( buffer, chunk.getSerializedWorkItemCommit().newInput()); } else { errorCollector.checkThat(buffer, Matchers.nullValue()); buffer = chunk.getSerializedWorkItemCommit().newInput(); } remainingBytes = chunk.getRemainingBytesForWorkItem(); if (remainingBytes == 0) { try { WorkItemCommitRequest received = WorkItemCommitRequest.parseFrom(buffer); errorCollector.checkThat( received, Matchers.equalTo(commitRequests.get(received.getWorkToken()))); try { responseObserver.onNext( StreamingCommitResponse.newBuilder() .addRequestId(chunk.getRequestId()) .build()); } catch (IllegalStateException e) { } } catch (Exception e) { errorCollector.addError(e); } buffer = null; } else { errorCollector.checkThat(first, Matchers.is(true)); } first = false; } } } @Override public void onError(Throwable throwable) {} @Override public void onCompleted() { injector.cancel(); responseObserver.onCompleted(); } }; } }); CommitWorkStream stream = client.commitWorkStream(); for (int i = 0; i < commitRequestList.size(); ) { final CountDownLatch latch = latches.get(i); if (stream.commitWorkItem( "computation", commitRequestList.get(i), (CommitStatus status) -> { assertEquals(status, CommitStatus.OK); latch.countDown(); })) { i++; } else { stream.flush(); } } stream.flush(); for (CountDownLatch latch : latches) { assertTrue(latch.await(1, TimeUnit.MINUTES)); } stream.close(); assertTrue(stream.awaitTermination(30, TimeUnit.SECONDS)); }