Examples with Elements org.apache.beam.model.fnexecution.v1.BeamFnApi.Elements used on opensource projects

Example 16 with Elements

use of org.apache.beam.model.fnexecution.v1.BeamFnApi.Elements in project beam by apache.

the class BeamFnDataWriteRunnerTest method createRecordingAggregator.

private BeamFnDataOutboundAggregator createRecordingAggregator(Map<String, List<WindowedValue<String>>> output, Supplier<String> bundleId) {
    PipelineOptions options = PipelineOptionsFactory.create();
    options.as(ExperimentalOptions.class).setExperiments(Arrays.asList("data_buffer_size_limit=0"));
    return new BeamFnDataOutboundAggregator(options, bundleId, new StreamObserver<Elements>() {

        @Override
        public void onNext(Elements elements) {
            for (Data data : elements.getDataList()) {
                try {
                    output.get(bundleId.get()).add(WIRE_CODER.decode(data.getData().newInput()));
                } catch (IOException e) {
                    throw new RuntimeException("Failed to decode output.");
                }
            }
        }

        @Override
        public void onError(Throwable throwable) {
        }

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

Example 17 with Elements

use of org.apache.beam.model.fnexecution.v1.BeamFnApi.Elements in project beam by apache.

the class BeamFnDataOutboundAggregatorTest method testConfiguredTimeLimit.

@Test
public void testConfiguredTimeLimit() throws Exception {
    List<Elements> values = new ArrayList<>();
    PipelineOptions options = PipelineOptionsFactory.create();
    options.as(ExperimentalOptions.class).setExperiments(Arrays.asList("data_buffer_time_limit_ms=1"));
    final CountDownLatch waitForFlush = new CountDownLatch(1);
    BeamFnDataOutboundAggregator aggregator = new BeamFnDataOutboundAggregator(options, endpoint::getInstructionId, TestStreams.withOnNext((Consumer<Elements>) e -> {
        values.add(e);
        waitForFlush.countDown();
    }).build(), false);
    // Test that it emits when time passed the time limit
    FnDataReceiver<byte[]> dataReceiver = registerOutputLocation(aggregator, endpoint, CODER);
    aggregator.start();
    dataReceiver.accept(new byte[1]);
    // wait the flush thread to flush the buffer
    waitForFlush.await();
    assertEquals(messageWithData(new byte[1]), values.get(0));
}

Example 18 with Elements

use of org.apache.beam.model.fnexecution.v1.BeamFnApi.Elements in project beam by apache.

the class ProcessBundleHandler method processBundle.

/**
 * Processes a bundle, running the start(), process(), and finish() functions. This function is
 * required to be reentrant.
 */
public BeamFnApi.InstructionResponse.Builder processBundle(BeamFnApi.InstructionRequest request) throws Exception {
    BeamFnApi.ProcessBundleResponse.Builder response = BeamFnApi.ProcessBundleResponse.newBuilder();
    BundleProcessor bundleProcessor = bundleProcessorCache.get(request, () -> {
        try {
            return createBundleProcessor(request.getProcessBundle().getProcessBundleDescriptorId(), request.getProcessBundle());
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    });
    try {
        PTransformFunctionRegistry startFunctionRegistry = bundleProcessor.getStartFunctionRegistry();
        PTransformFunctionRegistry finishFunctionRegistry = bundleProcessor.getFinishFunctionRegistry();
        ExecutionStateTracker stateTracker = bundleProcessor.getStateTracker();
        try (HandleStateCallsForBundle beamFnStateClient = bundleProcessor.getBeamFnStateClient()) {
            try (Closeable closeTracker = stateTracker.activate()) {
                // Already in reverse topological order so we don't need to do anything.
                for (ThrowingRunnable startFunction : startFunctionRegistry.getFunctions()) {
                    LOG.debug("Starting function {}", startFunction);
                    startFunction.run();
                }
                if (request.getProcessBundle().hasElements()) {
                    boolean inputFinished = bundleProcessor.getInboundObserver().multiplexElements(request.getProcessBundle().getElements());
                    if (!inputFinished) {
                        throw new RuntimeException("Elements embedded in ProcessBundleRequest do not contain stream terminators for " + "all data and timer inputs. Unterminated endpoints: " + bundleProcessor.getInboundObserver().getUnfinishedEndpoints());
                    }
                } else if (!bundleProcessor.getInboundEndpointApiServiceDescriptors().isEmpty()) {
                    BeamFnDataInboundObserver2 observer = bundleProcessor.getInboundObserver();
                    beamFnDataClient.registerReceiver(request.getInstructionId(), bundleProcessor.getInboundEndpointApiServiceDescriptors(), observer);
                    observer.awaitCompletion();
                    beamFnDataClient.unregisterReceiver(request.getInstructionId(), bundleProcessor.getInboundEndpointApiServiceDescriptors());
                }
                // Need to reverse this since we want to call finish in topological order.
                for (ThrowingRunnable finishFunction : Lists.reverse(finishFunctionRegistry.getFunctions())) {
                    LOG.debug("Finishing function {}", finishFunction);
                    finishFunction.run();
                }
            }
            // If bundleProcessor has not flushed any elements, embed them in response.
            embedOutboundElementsIfApplicable(response, bundleProcessor);
            // Add all checkpointed residuals to the response.
            response.addAllResidualRoots(bundleProcessor.getSplitListener().getResidualRoots());
            // Add all metrics to the response.
            Map<String, ByteString> monitoringData = monitoringData(bundleProcessor);
            if (runnerAcceptsShortIds) {
                response.putAllMonitoringData(monitoringData);
            } else {
                for (Map.Entry<String, ByteString> metric : monitoringData.entrySet()) {
                    response.addMonitoringInfos(shortIds.get(metric.getKey()).toBuilder().setPayload(metric.getValue()));
                }
            }
            if (!bundleProcessor.getBundleFinalizationCallbackRegistrations().isEmpty()) {
                finalizeBundleHandler.registerCallbacks(bundleProcessor.getInstructionId(), ImmutableList.copyOf(bundleProcessor.getBundleFinalizationCallbackRegistrations()));
                response.setRequiresFinalization(true);
            }
        }
        // Mark the bundle processor as re-usable.
        bundleProcessorCache.release(request.getProcessBundle().getProcessBundleDescriptorId(), bundleProcessor);
        return BeamFnApi.InstructionResponse.newBuilder().setProcessBundle(response);
    } catch (Exception e) {
        // Make sure we clean-up from the active set of bundle processors.
        bundleProcessorCache.discard(bundleProcessor);
        throw e;
    }
}

Example 19 with Elements

use of org.apache.beam.model.fnexecution.v1.BeamFnApi.Elements in project beam by apache.

the class BeamFnDataOutboundAggregatorTest method testConfiguredTimeLimitExceptionPropagation.

@Test
public void testConfiguredTimeLimitExceptionPropagation() throws Exception {
    PipelineOptions options = PipelineOptionsFactory.create();
    options.as(ExperimentalOptions.class).setExperiments(Arrays.asList("data_buffer_time_limit_ms=1"));
    BeamFnDataOutboundAggregator aggregator = new BeamFnDataOutboundAggregator(options, endpoint::getInstructionId, TestStreams.withOnNext((Consumer<Elements>) e -> {
        throw new RuntimeException("");
    }).build(), false);
    // Test that it emits when time passed the time limit
    FnDataReceiver<byte[]> dataReceiver = registerOutputLocation(aggregator, endpoint, CODER);
    aggregator.start();
    dataReceiver.accept(new byte[1]);
    // wait the flush thread to flush the buffer
    while (!aggregator.flushFuture.isDone()) {
        Thread.sleep(1);
    }
    try {
        // Test that the exception caught in the flush thread is propagated to
        // the main thread when processing the next element
        dataReceiver.accept(new byte[1]);
        fail();
    } catch (Exception e) {
    // expected
    }
    aggregator = new BeamFnDataOutboundAggregator(options, endpoint::getInstructionId, TestStreams.withOnNext((Consumer<Elements>) e -> {
        throw new RuntimeException("");
    }).build(), false);
    dataReceiver = registerOutputLocation(aggregator, endpoint, CODER);
    aggregator.start();
    dataReceiver.accept(new byte[1]);
    // wait the flush thread to flush the buffer
    while (!aggregator.flushFuture.isDone()) {
        Thread.sleep(1);
    }
    try {
        // Test that the exception caught in the flush thread is propagated to
        // the main thread when closing
        aggregator.sendOrCollectBufferedDataAndFinishOutboundStreams();
        fail();
    } catch (Exception e) {
    // expected
    }
}

Example 20 with Elements

use of org.apache.beam.model.fnexecution.v1.BeamFnApi.Elements in project beam by apache.

the class BeamFnDataOutboundAggregatorTest method testWithDefaultBuffer.

@Test
public void testWithDefaultBuffer() throws Exception {
    final List<Elements> values = new ArrayList<>();
    final AtomicBoolean onCompletedWasCalled = new AtomicBoolean();
    BeamFnDataOutboundAggregator aggregator = new BeamFnDataOutboundAggregator(PipelineOptionsFactory.create(), endpoint::getInstructionId, TestStreams.<Elements>withOnNext(values::add).withOnCompleted(() -> onCompletedWasCalled.set(true)).build(), false);
    // Test that nothing is emitted till the default buffer size is surpassed.
    FnDataReceiver<byte[]> dataReceiver = registerOutputLocation(aggregator, endpoint, CODER);
    aggregator.start();
    dataReceiver.accept(new byte[BeamFnDataOutboundAggregator.DEFAULT_BUFFER_LIMIT_BYTES - 50]);
    assertThat(values, empty());
    // Test that when we cross the buffer, we emit.
    dataReceiver.accept(new byte[50]);
    assertEquals(messageWithData(new byte[BeamFnDataOutboundAggregator.DEFAULT_BUFFER_LIMIT_BYTES - 50], new byte[50]), values.get(0));
    // Test that nothing is emitted till the default buffer size is surpassed after a reset
    dataReceiver.accept(new byte[BeamFnDataOutboundAggregator.DEFAULT_BUFFER_LIMIT_BYTES - 50]);
    assertEquals(1, values.size());
    // Test that when we cross the buffer, we emit.
    dataReceiver.accept(new byte[50]);
    assertEquals(messageWithData(new byte[BeamFnDataOutboundAggregator.DEFAULT_BUFFER_LIMIT_BYTES - 50], new byte[50]), values.get(1));
    // Test that when we close with an empty buffer we only have one end of stream
    aggregator.sendOrCollectBufferedDataAndFinishOutboundStreams();
    assertEquals(endMessage(), values.get(2));
    // Test that we can close twice.
    aggregator.sendOrCollectBufferedDataAndFinishOutboundStreams();
    assertEquals(endMessage(), values.get(2));
}