Examples with WorkItem org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkItem used on opensource projects

Example 1 with WorkItem

use of org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkItem 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 WorkItem

use of org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkItem in project beam by apache.

the class StreamingGroupAlsoByWindowsReshuffleDoFnTest method testFixedWindows.

@Test
public void testFixedWindows() throws Exception {
    TupleTag<KV<String, Iterable<String>>> outputTag = new TupleTag<>();
    ListOutputManager outputManager = new ListOutputManager();
    DoFnRunner<KeyedWorkItem<String, String>, KV<String, Iterable<String>>> runner = makeRunner(outputTag, outputManager);
    runner.startBundle();
    WorkItem.Builder workItem = WorkItem.newBuilder();
    workItem.setKey(ByteString.copyFromUtf8(KEY));
    workItem.setWorkToken(WORK_TOKEN);
    InputMessageBundle.Builder messageBundle = workItem.addMessageBundlesBuilder();
    messageBundle.setSourceComputationId(SOURCE_COMPUTATION_ID);
    Coder<String> valueCoder = StringUtf8Coder.of();
    addElement(messageBundle, Arrays.asList(window(0, 10)), new Instant(1), valueCoder, "v1");
    addElement(messageBundle, Arrays.asList(window(0, 10)), new Instant(2), valueCoder, "v2");
    addElement(messageBundle, Arrays.asList(window(0, 10)), new Instant(0), valueCoder, "v0");
    addElement(messageBundle, Arrays.asList(window(10, 20)), new Instant(13), valueCoder, "v3");
    runner.processElement(createValue(workItem, valueCoder));
    runner.finishBundle();
    List<WindowedValue<KV<String, Iterable<String>>>> result = outputManager.getOutput(outputTag);
    assertEquals(4, result.size());
    WindowedValue<KV<String, Iterable<String>>> item0 = result.get(0);
    assertEquals(KEY, item0.getValue().getKey());
    assertThat(item0.getValue().getValue(), Matchers.containsInAnyOrder("v1"));
    assertEquals(new Instant(1), item0.getTimestamp());
    assertThat(item0.getWindows(), Matchers.<BoundedWindow>contains(window(0, 10)));
    WindowedValue<KV<String, Iterable<String>>> item1 = result.get(1);
    assertEquals(KEY, item1.getValue().getKey());
    assertThat(item1.getValue().getValue(), Matchers.containsInAnyOrder("v2"));
    assertEquals(new Instant(2), item1.getTimestamp());
    assertThat(item1.getWindows(), Matchers.<BoundedWindow>contains(window(0, 10)));
    WindowedValue<KV<String, Iterable<String>>> item2 = result.get(2);
    assertEquals(KEY, item2.getValue().getKey());
    assertThat(item2.getValue().getValue(), Matchers.containsInAnyOrder("v0"));
    assertEquals(new Instant(0), item2.getTimestamp());
    assertThat(item2.getWindows(), Matchers.<BoundedWindow>contains(window(0, 10)));
    WindowedValue<KV<String, Iterable<String>>> item3 = result.get(3);
    assertEquals(KEY, item3.getValue().getKey());
    assertThat(item3.getValue().getValue(), Matchers.containsInAnyOrder("v3"));
    assertEquals(new Instant(13), item3.getTimestamp());
    assertThat(item3.getWindows(), Matchers.<BoundedWindow>contains(window(10, 20)));
}

Example 3 with WorkItem

use of org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkItem in project beam by apache.

the class WindowingWindmillReader method iterator.

@Override
public NativeReaderIterator<WindowedValue<KeyedWorkItem<K, T>>> iterator() throws IOException {
    final K key = keyCoder.decode(context.getSerializedKey().newInput(), Coder.Context.OUTER);
    final WorkItem workItem = context.getWork();
    KeyedWorkItem<K, T> keyedWorkItem = new WindmillKeyedWorkItem<>(key, workItem, windowCoder, windowsCoder, valueCoder);
    final boolean isEmptyWorkItem = (Iterables.isEmpty(keyedWorkItem.timersIterable()) && Iterables.isEmpty(keyedWorkItem.elementsIterable()));
    final WindowedValue<KeyedWorkItem<K, T>> value = new ValueInEmptyWindows<>(keyedWorkItem);
    // Return a noop iterator when current workitem is an empty workitem.
    if (isEmptyWorkItem) {
        return new NativeReaderIterator<WindowedValue<KeyedWorkItem<K, T>>>() {

            @Override
            public boolean start() throws IOException {
                return false;
            }

            @Override
            public boolean advance() throws IOException {
                return false;
            }

            @Override
            public WindowedValue<KeyedWorkItem<K, T>> getCurrent() {
                throw new NoSuchElementException();
            }
        };
    } else {
        return new NativeReaderIterator<WindowedValue<KeyedWorkItem<K, T>>>() {

            private WindowedValue<KeyedWorkItem<K, T>> current;

            @Override
            public boolean start() throws IOException {
                current = value;
                return true;
            }

            @Override
            public boolean advance() throws IOException {
                current = null;
                return false;
            }

            @Override
            public WindowedValue<KeyedWorkItem<K, T>> getCurrent() {
                if (current == null) {
                    throw new NoSuchElementException();
                }
                return value;
            }
        };
    }
}

Example 4 with WorkItem

use of org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkItem 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));
}