Examples with CloudWindmillServiceV1Alpha1ImplBase org.apache.beam.runners.dataflow.worker.windmill.CloudWindmillServiceV1Alpha1Grpc.CloudWindmillServiceV1Alpha1ImplBase used on opensource projects

Example 1 with CloudWindmillServiceV1Alpha1ImplBase

use of org.apache.beam.runners.dataflow.worker.windmill.CloudWindmillServiceV1Alpha1Grpc.CloudWindmillServiceV1Alpha1ImplBase in project beam by apache.

the class GrpcWindmillServerTest method testStreamingGetWork.

@Test
public void testStreamingGetWork() throws Exception {
    // This fake server returns an infinite stream of identical WorkItems, obeying the request size
    // limits set by the client.
    serviceRegistry.addService(new CloudWindmillServiceV1Alpha1ImplBase() {

        @Override
        public StreamObserver<StreamingGetWorkRequest> getWorkStream(StreamObserver<StreamingGetWorkResponseChunk> responseObserver) {
            return new StreamObserver<StreamingGetWorkRequest>() {

                boolean sawHeader = false;

                ResponseErrorInjector injector = new ResponseErrorInjector(responseObserver);

                @Override
                public void onNext(StreamingGetWorkRequest request) {
                    maybeInjectError(responseObserver);
                    try {
                        long maxItems;
                        if (!sawHeader) {
                            errorCollector.checkThat(request.getRequest(), Matchers.equalTo(GetWorkRequest.newBuilder().setClientId(10).setJobId("job").setProjectId("project").setWorkerId("worker").setMaxItems(3).setMaxBytes(10000).build()));
                            sawHeader = true;
                            maxItems = request.getRequest().getMaxItems();
                        } else {
                            maxItems = request.getRequestExtension().getMaxItems();
                        }
                        for (int item = 0; item < maxItems; item++) {
                            long id = ThreadLocalRandom.current().nextLong();
                            ByteString serializedResponse = WorkItem.newBuilder().setKey(ByteString.copyFromUtf8("somewhat_long_key")).setWorkToken(id).setShardingKey(id).build().toByteString();
                            // Break the WorkItem into smaller chunks to test chunking code.
                            for (int i = 0; i < serializedResponse.size(); i += 10) {
                                int end = Math.min(serializedResponse.size(), i + 10);
                                StreamingGetWorkResponseChunk.Builder builder = StreamingGetWorkResponseChunk.newBuilder().setStreamId(id).setSerializedWorkItem(serializedResponse.substring(i, end)).setRemainingBytesForWorkItem(serializedResponse.size() - end);
                                if (i == 0) {
                                    builder.setComputationMetadata(ComputationWorkItemMetadata.newBuilder().setComputationId("comp").setDependentRealtimeInputWatermark(17000).setInputDataWatermark(18000));
                                }
                                try {
                                    responseObserver.onNext(builder.build());
                                } catch (IllegalStateException e) {
                                    // Client closed stream, we're done.
                                    return;
                                }
                            }
                        }
                    } catch (Exception e) {
                        errorCollector.addError(e);
                    }
                }

                @Override
                public void onError(Throwable throwable) {
                }

                @Override
                public void onCompleted() {
                    injector.cancel();
                    responseObserver.onCompleted();
                }
            };
        }
    });
    // Read the stream of WorkItems until 100 of them are received.
    CountDownLatch latch = new CountDownLatch(100);
    GetWorkStream stream = client.getWorkStream(GetWorkRequest.newBuilder().setClientId(10).setMaxItems(3).setMaxBytes(10000).build(), (String computation, @Nullable Instant inputDataWatermark, Instant synchronizedProcessingTime, Windmill.WorkItem workItem) -> {
        latch.countDown();
        assertEquals(inputDataWatermark, new Instant(18));
        assertEquals(synchronizedProcessingTime, new Instant(17));
        assertEquals(workItem.getKey(), ByteString.copyFromUtf8("somewhat_long_key"));
    });
    assertTrue(latch.await(30, TimeUnit.SECONDS));
    stream.close();
    assertTrue(stream.awaitTermination(30, TimeUnit.SECONDS));
}

Example 2 with CloudWindmillServiceV1Alpha1ImplBase

use of org.apache.beam.runners.dataflow.worker.windmill.CloudWindmillServiceV1Alpha1Grpc.CloudWindmillServiceV1Alpha1ImplBase in project beam by apache.

the class GrpcWindmillServerTest method testStreamingCommit.

@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) {
        // Build some requests of varying size with a few big ones.
        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);
    // This server receives WorkItemCommitRequests, and verifies they are equal to the above
    // commitRequest.
    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) {
                                    // Stream is closed.
                                    }
                                } 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();
                }
            };
        }
    });
    // Make the commit requests, waiting for each of them to be verified and acknowledged.
    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));
}

Example 3 with CloudWindmillServiceV1Alpha1ImplBase

use of org.apache.beam.runners.dataflow.worker.windmill.CloudWindmillServiceV1Alpha1Grpc.CloudWindmillServiceV1Alpha1ImplBase in project beam by apache.

the class GrpcWindmillServerTest method testStreamingGetDataHeartbeats.

@Test
public void testStreamingGetDataHeartbeats() throws Exception {
    // This server records the heartbeats observed but doesn't respond.
    final Map<String, List<KeyedGetDataRequest>> heartbeats = new HashMap<>();
    serviceRegistry.addService(new CloudWindmillServiceV1Alpha1ImplBase() {

        @Override
        public StreamObserver<StreamingGetDataRequest> getDataStream(StreamObserver<StreamingGetDataResponse> responseObserver) {
            return new StreamObserver<StreamingGetDataRequest>() {

                boolean sawHeader = false;

                @Override
                public void onNext(StreamingGetDataRequest chunk) {
                    try {
                        if (!sawHeader) {
                            LOG.info("Received header");
                            errorCollector.checkThat(chunk.getHeader(), Matchers.equalTo(JobHeader.newBuilder().setJobId("job").setProjectId("project").setWorkerId("worker").build()));
                            sawHeader = true;
                        } else {
                            LOG.info("Received {} heartbeats", chunk.getStateRequestCount());
                            errorCollector.checkThat(chunk.getSerializedSize(), Matchers.lessThanOrEqualTo(STREAM_CHUNK_SIZE));
                            errorCollector.checkThat(chunk.getRequestIdCount(), Matchers.is(0));
                            synchronized (heartbeats) {
                                for (ComputationGetDataRequest request : chunk.getStateRequestList()) {
                                    errorCollector.checkThat(request.getRequestsCount(), Matchers.is(1));
                                    heartbeats.putIfAbsent(request.getComputationId(), new ArrayList<>());
                                    heartbeats.get(request.getComputationId()).add(request.getRequestsList().get(0));
                                }
                            }
                        }
                    } catch (Exception e) {
                        errorCollector.addError(e);
                    }
                }

                @Override
                public void onError(Throwable throwable) {
                }

                @Override
                public void onCompleted() {
                    responseObserver.onCompleted();
                }
            };
        }
    });
    Map<String, List<KeyedGetDataRequest>> activeMap = new HashMap<>();
    List<String> computation1Keys = new ArrayList<>();
    List<String> computation2Keys = new ArrayList<>();
    for (int i = 0; i < 100; ++i) {
        computation1Keys.add("Computation1Key" + i);
        computation2Keys.add("Computation2Key" + largeString(i * 20));
    }
    activeMap.put("Computation1", makeHeartbeatRequest(computation1Keys));
    activeMap.put("Computation2", makeHeartbeatRequest(computation2Keys));
    GetDataStream stream = client.getDataStream();
    stream.refreshActiveWork(activeMap);
    stream.close();
    assertTrue(stream.awaitTermination(60, TimeUnit.SECONDS));
    while (true) {
        Thread.sleep(100);
        synchronized (heartbeats) {
            if (heartbeats.size() != activeMap.size()) {
                continue;
            }
            assertEquals(heartbeats, activeMap);
            break;
        }
    }
}

Example 4 with CloudWindmillServiceV1Alpha1ImplBase

use of org.apache.beam.runners.dataflow.worker.windmill.CloudWindmillServiceV1Alpha1Grpc.CloudWindmillServiceV1Alpha1ImplBase in project beam by apache.

the class GrpcWindmillServerTest method testStreamingGetData.

@Test
@SuppressWarnings("FutureReturnValueIgnored")
public void testStreamingGetData() throws Exception {
    // This server responds to GetDataRequests with responses that mirror the requests.
    serviceRegistry.addService(new CloudWindmillServiceV1Alpha1ImplBase() {

        @Override
        public StreamObserver<StreamingGetDataRequest> getDataStream(StreamObserver<StreamingGetDataResponse> responseObserver) {
            return new StreamObserver<StreamingGetDataRequest>() {

                boolean sawHeader = false;

                HashSet<Long> seenIds = new HashSet<>();

                ResponseErrorInjector injector = new ResponseErrorInjector(responseObserver);

                StreamingGetDataResponse.Builder responseBuilder = StreamingGetDataResponse.newBuilder();

                @Override
                public void onNext(StreamingGetDataRequest chunk) {
                    maybeInjectError(responseObserver);
                    try {
                        if (!sawHeader) {
                            LOG.info("Received header");
                            errorCollector.checkThat(chunk.getHeader(), Matchers.equalTo(JobHeader.newBuilder().setJobId("job").setProjectId("project").setWorkerId("worker").build()));
                            sawHeader = true;
                        } else {
                            LOG.info("Received get data of {} global data, {} data requests", chunk.getGlobalDataRequestCount(), chunk.getStateRequestCount());
                            errorCollector.checkThat(chunk.getSerializedSize(), Matchers.lessThanOrEqualTo(STREAM_CHUNK_SIZE));
                            int i = 0;
                            for (GlobalDataRequest request : chunk.getGlobalDataRequestList()) {
                                long requestId = chunk.getRequestId(i++);
                                errorCollector.checkThat(seenIds.add(requestId), Matchers.is(true));
                                sendResponse(requestId, processGlobalDataRequest(request));
                            }
                            for (ComputationGetDataRequest request : chunk.getStateRequestList()) {
                                long requestId = chunk.getRequestId(i++);
                                errorCollector.checkThat(seenIds.add(requestId), Matchers.is(true));
                                sendResponse(requestId, processStateRequest(request));
                            }
                            flushResponse();
                        }
                    } catch (Exception e) {
                        errorCollector.addError(e);
                    }
                }

                @Override
                public void onError(Throwable throwable) {
                }

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

                private ByteString processGlobalDataRequest(GlobalDataRequest request) {
                    errorCollector.checkThat(request.getStateFamily(), Matchers.is("family"));
                    return GlobalData.newBuilder().setDataId(request.getDataId()).setStateFamily("family").setData(ByteString.copyFromUtf8(request.getDataId().getTag())).build().toByteString();
                }

                private ByteString processStateRequest(ComputationGetDataRequest compRequest) {
                    errorCollector.checkThat(compRequest.getRequestsCount(), Matchers.is(1));
                    errorCollector.checkThat(compRequest.getComputationId(), Matchers.is("computation"));
                    KeyedGetDataRequest request = compRequest.getRequests(0);
                    KeyedGetDataResponse response = makeGetDataResponse(request.getValuesToFetch(0).getTag().toStringUtf8());
                    return response.toByteString();
                }

                private void sendResponse(long id, ByteString serializedResponse) {
                    if (ThreadLocalRandom.current().nextInt(4) == 0) {
                        sendChunkedResponse(id, serializedResponse);
                    } else {
                        responseBuilder.addRequestId(id).addSerializedResponse(serializedResponse);
                        if (responseBuilder.getRequestIdCount() > 10) {
                            flushResponse();
                        }
                    }
                }

                private void sendChunkedResponse(long id, ByteString serializedResponse) {
                    LOG.info("Sending response with {} chunks", (serializedResponse.size() / 10) + 1);
                    for (int i = 0; i < serializedResponse.size(); i += 10) {
                        int end = Math.min(serializedResponse.size(), i + 10);
                        try {
                            responseObserver.onNext(StreamingGetDataResponse.newBuilder().addRequestId(id).addSerializedResponse(serializedResponse.substring(i, end)).setRemainingBytesForResponse(serializedResponse.size() - end).build());
                        } catch (IllegalStateException e) {
                        // Stream is already closed.
                        }
                    }
                }

                private void flushResponse() {
                    if (responseBuilder.getRequestIdCount() > 0) {
                        LOG.info("Sending batched response of {} ids", responseBuilder.getRequestIdCount());
                        try {
                            responseObserver.onNext(responseBuilder.build());
                        } catch (IllegalStateException e) {
                        // Stream is already closed.
                        }
                        responseBuilder.clear();
                    }
                }
            };
        }
    });
    GetDataStream stream = client.getDataStream();
    // Make requests of varying sizes to test chunking, and verify the responses.
    ExecutorService executor = Executors.newFixedThreadPool(50);
    final CountDownLatch done = new CountDownLatch(200);
    for (int i = 0; i < 100; ++i) {
        final String key = "key" + i;
        final String s = i % 5 == 0 ? largeString(i) : "tag";
        executor.submit(() -> {
            errorCollector.checkThat(stream.requestKeyedData("computation", makeGetDataRequest(key, s)), Matchers.equalTo(makeGetDataResponse(s)));
            done.countDown();
        });
        executor.execute(() -> {
            errorCollector.checkThat(stream.requestGlobalData(makeGlobalDataRequest(key)), Matchers.equalTo(makeGlobalDataResponse(key)));
            done.countDown();
        });
    }
    done.await();
    stream.close();
    assertTrue(stream.awaitTermination(60, TimeUnit.SECONDS));
    executor.shutdown();
}

Example 5 with CloudWindmillServiceV1Alpha1ImplBase

use of org.apache.beam.runners.dataflow.worker.windmill.CloudWindmillServiceV1Alpha1Grpc.CloudWindmillServiceV1Alpha1ImplBase in project beam by apache.

the class GrpcWindmillServerTest method testThrottleSignal.

@Test
public void testThrottleSignal() throws Exception {
    // This server responds with work items until the throttleMessage limit is hit at which point it
    // returns RESROUCE_EXHAUSTED errors for throttleTime msecs after which it resumes sending
    // work items.
    final int throttleTime = 2000;
    final int throttleMessage = 15;
    serviceRegistry.addService(new CloudWindmillServiceV1Alpha1ImplBase() {

        long throttleStartTime = -1;

        int messageCount = 0;

        @Override
        public StreamObserver<StreamingGetWorkRequest> getWorkStream(StreamObserver<StreamingGetWorkResponseChunk> responseObserver) {
            return new StreamObserver<StreamingGetWorkRequest>() {

                boolean sawHeader = false;

                @Override
                public void onNext(StreamingGetWorkRequest request) {
                    messageCount++;
                    // error.
                    if (messageCount == throttleMessage || throttleStartTime != -1) {
                        // If throttling has not started yet then start it.
                        if (throttleStartTime == -1) {
                            throttleStartTime = Instant.now().getMillis();
                        }
                        // throttling stop throttling.
                        if (throttleStartTime != -1 && ((Instant.now().getMillis() - throttleStartTime) > throttleTime)) {
                            throttleStartTime = -1;
                        }
                        StatusRuntimeException error = new StatusRuntimeException(Status.RESOURCE_EXHAUSTED);
                        responseObserver.onError(error);
                        return;
                    }
                    // We are not throttling this message so respond as normal.
                    try {
                        long maxItems;
                        if (!sawHeader) {
                            sawHeader = true;
                            maxItems = request.getRequest().getMaxItems();
                        } else {
                            maxItems = request.getRequestExtension().getMaxItems();
                        }
                        for (int item = 0; item < maxItems; item++) {
                            long id = ThreadLocalRandom.current().nextLong();
                            ByteString serializedResponse = WorkItem.newBuilder().setKey(ByteString.copyFromUtf8("somewhat_long_key")).setWorkToken(id).setShardingKey(id).build().toByteString();
                            StreamingGetWorkResponseChunk.Builder builder = StreamingGetWorkResponseChunk.newBuilder().setStreamId(id).setSerializedWorkItem(serializedResponse).setRemainingBytesForWorkItem(0);
                            try {
                                responseObserver.onNext(builder.build());
                            } catch (IllegalStateException e) {
                                // Client closed stream, we're done.
                                return;
                            }
                        }
                    } catch (Exception e) {
                        errorCollector.addError(e);
                    }
                }

                @Override
                public void onError(Throwable throwable) {
                }

                @Override
                public void onCompleted() {
                    responseObserver.onCompleted();
                }
            };
        }
    });
    // Read the stream of WorkItems until 100 of them are received.
    CountDownLatch latch = new CountDownLatch(100);
    GetWorkStream stream = client.getWorkStream(GetWorkRequest.newBuilder().setClientId(10).setMaxItems(3).setMaxBytes(10000).build(), (String computation, @Nullable Instant inputDataWatermark, Instant synchronizedProcessingTime, Windmill.WorkItem workItem) -> {
        latch.countDown();
    });
    // Wait for 100 items or 30 seconds.
    assertTrue(latch.await(30, TimeUnit.SECONDS));
    // Confirm that we report at least as much throttle time as our server sent errors for.  We will
    // actually report more due to backoff in restarting streams.
    assertTrue(this.client.getAndResetThrottleTime() > throttleTime);
    stream.close();
    assertTrue(stream.awaitTermination(30, TimeUnit.SECONDS));
}