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

Example 1 with Elements

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

the class ServerFactoryTest method runTestUsing.

private Endpoints.ApiServiceDescriptor runTestUsing(ServerFactory serverFactory, ManagedChannelFactory channelFactory) throws Exception {
    Endpoints.ApiServiceDescriptor.Builder apiServiceDescriptorBuilder = Endpoints.ApiServiceDescriptor.newBuilder();
    final Collection<Elements> serverElements = new ArrayList<>();
    final CountDownLatch clientHangedUp = new CountDownLatch(1);
    CallStreamObserver<Elements> serverInboundObserver = TestStreams.withOnNext(serverElements::add).withOnCompleted(clientHangedUp::countDown).build();
    TestDataService service = new TestDataService(serverInboundObserver);
    Server server = serverFactory.allocateAddressAndCreate(ImmutableList.of(service), apiServiceDescriptorBuilder);
    assertFalse(server.isShutdown());
    ManagedChannel channel = channelFactory.forDescriptor(apiServiceDescriptorBuilder.build());
    BeamFnDataGrpc.BeamFnDataStub stub = BeamFnDataGrpc.newStub(channel);
    final Collection<BeamFnApi.Elements> clientElements = new ArrayList<>();
    final CountDownLatch serverHangedUp = new CountDownLatch(1);
    CallStreamObserver<BeamFnApi.Elements> clientInboundObserver = TestStreams.withOnNext(clientElements::add).withOnCompleted(serverHangedUp::countDown).build();
    StreamObserver<Elements> clientOutboundObserver = stub.data(clientInboundObserver);
    StreamObserver<BeamFnApi.Elements> serverOutboundObserver = service.outboundObservers.take();
    clientOutboundObserver.onNext(CLIENT_DATA);
    serverOutboundObserver.onNext(SERVER_DATA);
    clientOutboundObserver.onCompleted();
    clientHangedUp.await();
    serverOutboundObserver.onCompleted();
    serverHangedUp.await();
    assertThat(clientElements, contains(SERVER_DATA));
    assertThat(serverElements, contains(CLIENT_DATA));
    return apiServiceDescriptorBuilder.build();
}

Example 2 with Elements

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

the class ServerFactoryTest method usesUrlFactory.

@Test
public void usesUrlFactory() throws Exception {
    ServerFactory serverFactory = ServerFactory.createWithUrlFactory((host, port) -> "foo");
    CallStreamObserver<Elements> observer = TestStreams.withOnNext((Elements unused) -> {
    }).withOnCompleted(() -> {
    }).build();
    TestDataService service = new TestDataService(observer);
    ApiServiceDescriptor.Builder descriptorBuilder = ApiServiceDescriptor.newBuilder();
    grpcCleanupRule.register(serverFactory.allocateAddressAndCreate(ImmutableList.of(service), descriptorBuilder));
    assertThat(descriptorBuilder.getUrl(), is("foo"));
}

Example 3 with Elements

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

the class ProcessBundleHandler method embedOutboundElementsIfApplicable.

private void embedOutboundElementsIfApplicable(ProcessBundleResponse.Builder response, BundleProcessor bundleProcessor) {
    if (bundleProcessor.getOutboundAggregators().isEmpty()) {
        return;
    }
    List<Elements> collectedElements = new ArrayList<>(bundleProcessor.getOutboundAggregators().size());
    boolean hasFlushedAggregator = false;
    for (BeamFnDataOutboundAggregator aggregator : bundleProcessor.getOutboundAggregators().values()) {
        Elements elements = aggregator.sendOrCollectBufferedDataAndFinishOutboundStreams();
        if (elements == null) {
            hasFlushedAggregator = true;
        }
        collectedElements.add(elements);
    }
    if (!hasFlushedAggregator) {
        Elements.Builder elementsToEmbed = Elements.newBuilder();
        for (Elements collectedElement : collectedElements) {
            elementsToEmbed.mergeFrom(collectedElement);
        }
        response.setElements(elementsToEmbed.build());
    } else {
        // Since there was at least one flushed aggregator, we have to use the aggregators that were
        // able to successfully collect their elements to emit them and can not send them as part of
        // the ProcessBundleResponse.
        int i = 0;
        for (BeamFnDataOutboundAggregator aggregator : bundleProcessor.getOutboundAggregators().values()) {
            Elements elements = collectedElements.get(i++);
            if (elements != null) {
                aggregator.sendElements(elements);
            }
        }
    }
}

Example 4 with Elements

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

the class BeamFnDataGrpcMultiplexerTest method testInboundObserverBlocksTillConsumerConnects.

@Test
public void testInboundObserverBlocksTillConsumerConnects() throws Exception {
    final Collection<BeamFnApi.Elements> outboundValues = new ArrayList<>();
    final Collection<KV<ByteString, Boolean>> dataInboundValues = new ArrayList<>();
    final Collection<KV<ByteString, Boolean>> timerInboundValues = new ArrayList<>();
    final BeamFnDataGrpcMultiplexer multiplexer = new BeamFnDataGrpcMultiplexer(DESCRIPTOR, OutboundObserverFactory.clientDirect(), inboundObserver -> TestStreams.withOnNext(outboundValues::add).build());
    ExecutorService executorService = Executors.newCachedThreadPool();
    executorService.submit(() -> {
        // Purposefully sleep to simulate a delay in a consumer connecting.
        Uninterruptibles.sleepUninterruptibly(100, TimeUnit.MILLISECONDS);
        multiplexer.registerConsumer(DATA_LOCATION, (payload, isLast) -> dataInboundValues.add(KV.of(payload, isLast)));
        multiplexer.registerConsumer(TIMER_LOCATION, (payload, isLast) -> timerInboundValues.add(KV.of(payload, isLast)));
    }).get();
    multiplexer.getInboundObserver().onNext(ELEMENTS);
    assertTrue(multiplexer.hasConsumer(DATA_LOCATION));
    assertTrue(multiplexer.hasConsumer(TIMER_LOCATION));
    // Ensure that when we see a terminal Elements object, we remove the consumer
    multiplexer.getInboundObserver().onNext(TERMINAL_ELEMENTS);
    assertFalse(multiplexer.hasConsumer(DATA_LOCATION));
    assertFalse(multiplexer.hasConsumer(TIMER_LOCATION));
    // Assert that normal and terminal Elements are passed to the consumer
    assertThat(dataInboundValues, contains(KV.of(ELEMENTS.getData(0).getData(), false), KV.of(ByteString.EMPTY, true)));
    assertThat(timerInboundValues, contains(KV.of(ELEMENTS.getTimers(0).getTimers(), false), KV.of(ByteString.EMPTY, true)));
}

Example 5 with Elements

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

the class BeamFnDataReadRunner method trySplit.

public void trySplit(ProcessBundleSplitRequest request, ProcessBundleSplitResponse.Builder response) {
    DesiredSplit desiredSplit = request.getDesiredSplitsMap().get(pTransformId);
    if (desiredSplit == null) {
        return;
    }
    long totalBufferSize = desiredSplit.getEstimatedInputElements();
    List<Long> allowedSplitPoints = new ArrayList<>(desiredSplit.getAllowedSplitPointsList());
    HandlesSplits splittingConsumer = null;
    if (consumer instanceof HandlesSplits) {
        splittingConsumer = ((HandlesSplits) consumer);
    }
    synchronized (splittingLock) {
        // provide.
        if (index == stopIndex) {
            return;
        }
        // being released.
        if (!request.getInstructionId().equals(processBundleInstructionIdSupplier.get())) {
            return;
        }
        // split request is bounded incorrectly, use the stop index as the upper bound.
        if (totalBufferSize < index + 1) {
            totalBufferSize = index + 1;
        } else if (totalBufferSize > stopIndex) {
            totalBufferSize = stopIndex;
        }
        // In the case where we have yet to process an element, set the current element progress to 1.
        double currentElementProgress = 1;
        // progress defaulting to 0.5 if no progress was able to get fetched.
        if (index >= 0) {
            if (splittingConsumer != null) {
                currentElementProgress = splittingConsumer.getProgress();
            } else {
                currentElementProgress = 0.5;
            }
        }
        // Now figure out where to split.
        // 
        // The units here (except for keepOfElementRemainder) are all in terms of number or
        // (possibly fractional) elements.
        // Compute the amount of "remaining" work that we know of.
        double remainder = totalBufferSize - index - currentElementProgress;
        // Compute the number of elements (including fractional elements) that we should "keep".
        double keep = remainder * desiredSplit.getFractionOfRemainder();
        // splittable.
        if (currentElementProgress < 1) {
            // See if the amount we need to keep falls within the current element's remainder and if
            // so, attempt to split it.
            double keepOfElementRemainder = keep / (1 - currentElementProgress);
            // If both index and index are allowed split point, we can split at index.
            if (keepOfElementRemainder < 1 && isValidSplitPoint(allowedSplitPoints, index) && isValidSplitPoint(allowedSplitPoints, index + 1)) {
                SplitResult splitResult = splittingConsumer != null ? splittingConsumer.trySplit(keepOfElementRemainder) : null;
                if (splitResult != null) {
                    stopIndex = index + 1;
                    response.addAllPrimaryRoots(splitResult.getPrimaryRoots()).addAllResidualRoots(splitResult.getResidualRoots()).addChannelSplitsBuilder().setLastPrimaryElement(index - 1).setFirstResidualElement(stopIndex);
                    return;
                }
            }
        }
        // Otherwise, split at the closest allowed element boundary.
        long newStopIndex = index + Math.max(1, Math.round(currentElementProgress + keep));
        if (!isValidSplitPoint(allowedSplitPoints, newStopIndex)) {
            // Choose the closest allowed split point.
            Collections.sort(allowedSplitPoints);
            int closestSplitPointIndex = -(Collections.binarySearch(allowedSplitPoints, newStopIndex) + 1);
            if (closestSplitPointIndex == 0) {
                newStopIndex = allowedSplitPoints.get(0);
            } else if (closestSplitPointIndex == allowedSplitPoints.size()) {
                newStopIndex = allowedSplitPoints.get(closestSplitPointIndex - 1);
            } else {
                long prevPoint = allowedSplitPoints.get(closestSplitPointIndex - 1);
                long nextPoint = allowedSplitPoints.get(closestSplitPointIndex);
                if (index < prevPoint && newStopIndex - prevPoint < nextPoint - newStopIndex) {
                    newStopIndex = prevPoint;
                } else {
                    newStopIndex = nextPoint;
                }
            }
        }
        if (newStopIndex < stopIndex && newStopIndex > index) {
            stopIndex = newStopIndex;
            response.addChannelSplitsBuilder().setLastPrimaryElement(stopIndex - 1).setFirstResidualElement(stopIndex);
            return;
        }
    }
}