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

Example 16 with Data

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

the class BeamFnDataGrpcServiceTest method testMessageReceivedByProperClientWhenThereAreMultipleClients.

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

Example 17 with Data

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

the class GrpcDataServiceTest method testMessageReceivedBySingleClientWhenThereAreMultipleClients.

@Test
public void testMessageReceivedBySingleClientWhenThereAreMultipleClients() throws Exception {
    final LinkedBlockingQueue<Elements> clientInboundElements = new LinkedBlockingQueue<>();
    ExecutorService executorService = Executors.newCachedThreadPool();
    final CountDownLatch waitForInboundElements = new CountDownLatch(1);
    GrpcDataService service = GrpcDataService.create(PipelineOptionsFactory.create(), Executors.newCachedThreadPool(), OutboundObserverFactory.serverDirect());
    try (GrpcFnServer<GrpcDataService> server = GrpcFnServer.allocatePortAndCreateFor(service, InProcessServerFactory.create())) {
        Collection<Future<Void>> clientFutures = new ArrayList<>();
        for (int i = 0; i < 3; ++i) {
            clientFutures.add(executorService.submit(() -> {
                ManagedChannel channel = InProcessChannelBuilder.forName(server.getApiServiceDescriptor().getUrl()).directExecutor().build();
                StreamObserver<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.send(LogicalEndpoint.data(Integer.toString(i), TRANSFORM_ID), CODER);
            consumer.accept(WindowedValue.valueInGlobalWindow("A" + i));
            consumer.accept(WindowedValue.valueInGlobalWindow("B" + i));
            consumer.accept(WindowedValue.valueInGlobalWindow("C" + i));
            consumer.close();
        }
        waitForInboundElements.countDown();
        for (Future<Void> clientFuture : clientFutures) {
            clientFuture.get();
        }
        assertThat(clientInboundElements, containsInAnyOrder(elementsWithData("0"), elementsWithData("1"), elementsWithData("2")));
    }
}

Example 18 with Data

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

the class GrpcDataServiceTest method testMultipleClientsSendMessagesAreDirectedToProperConsumers.

@Test
public void testMultipleClientsSendMessagesAreDirectedToProperConsumers() throws Exception {
    final LinkedBlockingQueue<BeamFnApi.Elements> clientInboundElements = new LinkedBlockingQueue<>();
    ExecutorService executorService = Executors.newCachedThreadPool();
    final CountDownLatch waitForInboundElements = new CountDownLatch(1);
    GrpcDataService service = GrpcDataService.create(PipelineOptionsFactory.create(), Executors.newCachedThreadPool(), OutboundObserverFactory.serverDirect());
    try (GrpcFnServer<GrpcDataService> server = GrpcFnServer.allocatePortAndCreateFor(service, InProcessServerFactory.create())) {
        Collection<Future<Void>> clientFutures = new ArrayList<>();
        for (int i = 0; i < 3; ++i) {
            final String instructionId = Integer.toString(i);
            clientFutures.add(executorService.submit(() -> {
                ManagedChannel channel = InProcessChannelBuilder.forName(server.getApiServiceDescriptor().getUrl()).build();
                StreamObserver<Elements> outboundObserver = BeamFnDataGrpc.newStub(channel).data(TestStreams.withOnNext(clientInboundElements::add).build());
                outboundObserver.onNext(elementsWithData(instructionId));
                waitForInboundElements.await();
                outboundObserver.onCompleted();
                return null;
            }));
        }
        List<Collection<WindowedValue<String>>> serverInboundValues = new ArrayList<>();
        Collection<InboundDataClient> readFutures = new ArrayList<>();
        for (int i = 0; i < 3; ++i) {
            final Collection<WindowedValue<String>> serverInboundValue = new ArrayList<>();
            serverInboundValues.add(serverInboundValue);
            readFutures.add(service.receive(LogicalEndpoint.data(Integer.toString(i), TRANSFORM_ID), CODER, serverInboundValue::add));
        }
        for (InboundDataClient readFuture : readFutures) {
            readFuture.awaitCompletion();
        }
        waitForInboundElements.countDown();
        for (Future<Void> clientFuture : clientFutures) {
            clientFuture.get();
        }
        for (int i = 0; i < 3; ++i) {
            assertThat(serverInboundValues.get(i), contains(WindowedValue.valueInGlobalWindow("A" + i), WindowedValue.valueInGlobalWindow("B" + i), WindowedValue.valueInGlobalWindow("C" + i)));
        }
        assertThat(clientInboundElements, empty());
    }
}