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

Example 6 with Elements

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

the class BeamFnDataOutboundAggregatorTest method testConfiguredBufferLimit.

@Test
public void testConfiguredBufferLimit() throws Exception {
    List<BeamFnApi.Elements> values = new ArrayList<>();
    AtomicBoolean onCompletedWasCalled = new AtomicBoolean();
    PipelineOptions options = PipelineOptionsFactory.create();
    options.as(ExperimentalOptions.class).setExperiments(Arrays.asList("data_buffer_size_limit=100"));
    BeamFnDataOutboundAggregator aggregator = new BeamFnDataOutboundAggregator(options, 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[51]);
    assertThat(values, empty());
    // Test that when we cross the buffer, we emit.
    dataReceiver.accept(new byte[49]);
    assertEquals(messageWithData(new byte[51], new byte[49]), values.get(0));
    Receiver<?> receiver;
    if (endpoint.isTimer()) {
        receiver = Iterables.getOnlyElement(aggregator.outputTimersReceivers.values());
    } else {
        receiver = Iterables.getOnlyElement(aggregator.outputDataReceivers.values());
    }
    assertEquals(0L, receiver.getOutput().size());
    assertEquals(102L, receiver.getByteCount());
    assertEquals(2L, receiver.getElementCount());
    // Test that when we close we empty the value, and then send the stream terminator as part
    // of the same message
    dataReceiver.accept(new byte[1]);
    aggregator.sendOrCollectBufferedDataAndFinishOutboundStreams();
    // Test that receiver stats have been reset after
    // sendOrCollectBufferedDataAndFinishOutboundStreams.
    assertEquals(0L, receiver.getOutput().size());
    assertEquals(0L, receiver.getByteCount());
    assertEquals(0L, receiver.getElementCount());
    BeamFnApi.Elements.Builder builder = messageWithDataBuilder(new byte[1]);
    if (endpoint.isTimer()) {
        builder.addTimers(BeamFnApi.Elements.Timers.newBuilder().setInstructionId(endpoint.getInstructionId()).setTransformId(endpoint.getTransformId()).setTimerFamilyId(endpoint.getTimerFamilyId()).setIsLast(true));
    } else {
        builder.addData(BeamFnApi.Elements.Data.newBuilder().setInstructionId(endpoint.getInstructionId()).setTransformId(endpoint.getTransformId()).setIsLast(true));
    }
    assertEquals(builder.build(), values.get(1));
}

Example 7 with Elements

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

the class BeamFnDataOutboundAggregatorTest method testConfiguredBufferLimitMultipleEndpoints.

@Test
public void testConfiguredBufferLimitMultipleEndpoints() throws Exception {
    List<BeamFnApi.Elements> values = new ArrayList<>();
    AtomicBoolean onCompletedWasCalled = new AtomicBoolean();
    PipelineOptions options = PipelineOptionsFactory.create();
    options.as(ExperimentalOptions.class).setExperiments(Arrays.asList("data_buffer_size_limit=100"));
    BeamFnDataOutboundAggregator aggregator = new BeamFnDataOutboundAggregator(options, 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.
    LogicalEndpoint additionalEndpoint = LogicalEndpoint.data(endpoint.getInstructionId(), "additional:" + endpoint.getTransformId());
    FnDataReceiver<byte[]> dataReceiver = registerOutputLocation(aggregator, endpoint, CODER);
    FnDataReceiver<byte[]> additionalDataReceiver = registerOutputLocation(aggregator, additionalEndpoint, CODER);
    aggregator.start();
    dataReceiver.accept(new byte[51]);
    assertThat(values, empty());
    // Test that when we cross the buffer, we emit.
    additionalDataReceiver.accept(new byte[49]);
    checkEqualInAnyOrder(messageWithDataBuilder(new byte[51]).mergeFrom(messageWithDataBuilder(additionalEndpoint, new byte[49]).build()).build(), values.get(0));
    // Test that when we close we empty the value, and then the stream terminator as part
    // of the same message
    dataReceiver.accept(new byte[1]);
    aggregator.sendOrCollectBufferedDataAndFinishOutboundStreams();
    BeamFnApi.Elements.Builder builder = messageWithDataBuilder(new byte[1]);
    if (endpoint.isTimer()) {
        builder.addTimers(BeamFnApi.Elements.Timers.newBuilder().setInstructionId(endpoint.getInstructionId()).setTransformId(endpoint.getTransformId()).setTimerFamilyId(endpoint.getTimerFamilyId()).setIsLast(true));
    } else {
        builder.addData(BeamFnApi.Elements.Data.newBuilder().setInstructionId(endpoint.getInstructionId()).setTransformId(endpoint.getTransformId()).setIsLast(true));
    }
    builder.addData(BeamFnApi.Elements.Data.newBuilder().setInstructionId(additionalEndpoint.getInstructionId()).setTransformId(additionalEndpoint.getTransformId()).setIsLast(true));
    checkEqualInAnyOrder(builder.build(), values.get(1));
}

Example 8 with Elements

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

the class BeamFnDataOutboundAggregator method sendOrCollectBufferedDataAndFinishOutboundStreams.

/**
 * Closes the streams for all registered outbound endpoints. Should be called at the end of each
 * bundle. Returns the buffered Elements if the BeamFnDataOutboundAggregator started with
 * collectElementsIfNoFlushes=true, and there was no previous flush in this bundle, otherwise
 * returns null.
 */
public Elements sendOrCollectBufferedDataAndFinishOutboundStreams() {
    if (outputTimersReceivers.isEmpty() && outputDataReceivers.isEmpty()) {
        return null;
    }
    Elements.Builder bufferedElements;
    if (timeLimit > 0) {
        synchronized (flushLock) {
            bufferedElements = convertBufferForTransmission();
        }
    } else {
        bufferedElements = convertBufferForTransmission();
    }
    LOG.debug("Closing streams for instruction {} and outbound data {} and timers {}.", processBundleRequestIdSupplier.get(), outputDataReceivers, outputTimersReceivers);
    for (Map.Entry<String, Receiver<?>> entry : outputDataReceivers.entrySet()) {
        String pTransformId = entry.getKey();
        bufferedElements.addDataBuilder().setInstructionId(processBundleRequestIdSupplier.get()).setTransformId(pTransformId).setIsLast(true);
        entry.getValue().resetStats();
    }
    for (Map.Entry<TimerEndpoint, Receiver<?>> entry : outputTimersReceivers.entrySet()) {
        TimerEndpoint timerKey = entry.getKey();
        bufferedElements.addTimersBuilder().setInstructionId(processBundleRequestIdSupplier.get()).setTransformId(timerKey.pTransformId).setTimerFamilyId(timerKey.timerFamilyId).setIsLast(true);
        entry.getValue().resetStats();
    }
    if (collectElementsIfNoFlushes && !hasFlushedForBundle) {
        return bufferedElements.build();
    }
    outboundObserver.onNext(bufferedElements.build());
    // This is now at the end of a bundle, so we reset hasFlushedForBundle to prepare for new
    // bundles.
    hasFlushedForBundle = false;
    return null;
}

Example 9 with Elements

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

the class BeamFnDataInboundObserver2 method multiplexElements.

/**
 * Dispatches the data and timers from the elements to corresponding receivers. Returns true if
 * all the endpoints are done after elements dispatching.
 */
public boolean multiplexElements(Elements elements) throws Exception {
    for (BeamFnApi.Elements.Data data : elements.getDataList()) {
        EndpointStatus<DataEndpoint<?>> endpoint = transformIdToDataEndpoint.get(data.getTransformId());
        if (endpoint == null) {
            throw new IllegalStateException(String.format("Unable to find inbound data receiver for instruction %s and transform %s.", data.getInstructionId(), data.getTransformId()));
        } else if (endpoint.isDone) {
            throw new IllegalStateException(String.format("Received data after inbound data receiver is done for instruction %s and transform %s.", data.getInstructionId(), data.getTransformId()));
        }
        InputStream inputStream = data.getData().newInput();
        Coder<Object> coder = (Coder<Object>) endpoint.endpoint.getCoder();
        FnDataReceiver<Object> receiver = (FnDataReceiver<Object>) endpoint.endpoint.getReceiver();
        while (inputStream.available() > 0) {
            receiver.accept(coder.decode(inputStream));
        }
        if (data.getIsLast()) {
            endpoint.isDone = true;
            numEndpointsThatAreIncomplete -= 1;
        }
    }
    for (BeamFnApi.Elements.Timers timers : elements.getTimersList()) {
        Map<String, EndpointStatus<TimerEndpoint<?>>> timerFamilyIdToEndpoints = transformIdToTimerFamilyIdToTimerEndpoint.get(timers.getTransformId());
        if (timerFamilyIdToEndpoints == null) {
            throw new IllegalStateException(String.format("Unable to find inbound timer receiver for instruction %s, transform %s, and timer family %s.", timers.getInstructionId(), timers.getTransformId(), timers.getTimerFamilyId()));
        }
        EndpointStatus<TimerEndpoint<?>> endpoint = timerFamilyIdToEndpoints.get(timers.getTimerFamilyId());
        if (endpoint == null) {
            throw new IllegalStateException(String.format("Unable to find inbound timer receiver for instruction %s, transform %s, and timer family %s.", timers.getInstructionId(), timers.getTransformId(), timers.getTimerFamilyId()));
        } else if (endpoint.isDone) {
            throw new IllegalStateException(String.format("Received timer after inbound timer receiver is done for instruction %s, transform %s, and timer family %s.", timers.getInstructionId(), timers.getTransformId(), timers.getTimerFamilyId()));
        }
        InputStream inputStream = timers.getTimers().newInput();
        Coder<Object> coder = (Coder<Object>) endpoint.endpoint.getCoder();
        FnDataReceiver<Object> receiver = (FnDataReceiver<Object>) endpoint.endpoint.getReceiver();
        while (inputStream.available() > 0) {
            receiver.accept(coder.decode(inputStream));
        }
        if (timers.getIsLast()) {
            endpoint.isDone = true;
            numEndpointsThatAreIncomplete -= 1;
        }
    }
    return numEndpointsThatAreIncomplete == 0;
}

Example 10 with Elements

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

the class BeamFnDataGrpcServiceTest method testMessageReceivedBySingleClientWhenThereAreMultipleClients.

@Test
public void testMessageReceivedBySingleClientWhenThereAreMultipleClients() throws Exception {
    BlockingQueue<Elements> clientInboundElements = new LinkedBlockingQueue<>();
    ExecutorService executorService = Executors.newCachedThreadPool();
    CountDownLatch waitForInboundElements = new CountDownLatch(1);
    int numberOfClients = 3;
    for (int client = 0; client < numberOfClients; ++client) {
        executorService.submit(() -> {
            ManagedChannel channel = ManagedChannelFactory.createDefault().withInterceptors(Arrays.asList(AddHarnessIdInterceptor.create(WORKER_ID))).forDescriptor(service.getApiServiceDescriptor());
            StreamObserver<BeamFnApi.Elements> outboundObserver = BeamFnDataGrpc.newStub(channel).data(TestStreams.withOnNext(clientInboundElements::add).build());
            waitForInboundElements.await();
            outboundObserver.onCompleted();
            return null;
        });
    }
    for (int i = 0; i < 3; ++i) {
        CloseableFnDataReceiver<WindowedValue<String>> consumer = service.getDataService(WORKER_ID).send(LogicalEndpoint.data(Integer.toString(i), TRANSFORM_ID), CODER);
        consumer.accept(valueInGlobalWindow("A" + i));
        consumer.accept(valueInGlobalWindow("B" + i));
        consumer.accept(valueInGlobalWindow("C" + i));
        consumer.close();
    }
    // Specifically copy the elements to a new list so we perform blocking calls on the queue
    // to ensure the elements arrive.
    List<Elements> copy = new ArrayList<>();
    for (int i = 0; i < numberOfClients; ++i) {
        copy.add(clientInboundElements.take());
    }
    assertThat(copy, containsInAnyOrder(elementsWithData("0"), elementsWithData("1"), elementsWithData("2")));
    waitForInboundElements.countDown();
}