Search in sources :

Example 91 with Channel

use of org.apache.beam.vendor.grpc.v1p43p2.io.grpc.Channel 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());
    }
}
Also used : StreamObserver(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.stub.StreamObserver) ArrayList(java.util.ArrayList) ByteString(org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.ByteString) Elements(org.apache.beam.model.fnexecution.v1.BeamFnApi.Elements) LinkedBlockingQueue(java.util.concurrent.LinkedBlockingQueue) CountDownLatch(java.util.concurrent.CountDownLatch) LogicalEndpoint(org.apache.beam.sdk.fn.data.LogicalEndpoint) InboundDataClient(org.apache.beam.sdk.fn.data.InboundDataClient) WindowedValue(org.apache.beam.sdk.util.WindowedValue) ExecutorService(java.util.concurrent.ExecutorService) Future(java.util.concurrent.Future) ManagedChannel(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.ManagedChannel) Collection(java.util.Collection) Test(org.junit.Test)

Aggregations

Channel (io.grpc.Channel)63 Test (org.junit.Test)55 CallOptions (io.grpc.CallOptions)38 Metadata (io.grpc.Metadata)26 ManagedChannel (org.apache.beam.vendor.grpc.v1p43p2.io.grpc.ManagedChannel)25 MethodDescriptor (io.grpc.MethodDescriptor)21 ClientInterceptor (io.grpc.ClientInterceptor)20 ManagedChannel (io.grpc.ManagedChannel)20 CountDownLatch (java.util.concurrent.CountDownLatch)16 ClientCall (io.grpc.ClientCall)14 ArrayList (java.util.ArrayList)14 ExecutorService (java.util.concurrent.ExecutorService)12 StreamObserver (org.apache.beam.vendor.grpc.v1p43p2.io.grpc.stub.StreamObserver)12 AtomicReference (java.util.concurrent.atomic.AtomicReference)11 ConcurrentLinkedQueue (java.util.concurrent.ConcurrentLinkedQueue)10 BeamFnApi (org.apache.beam.model.fnexecution.v1.BeamFnApi)10 LinkedBlockingQueue (java.util.concurrent.LinkedBlockingQueue)9 Endpoints (org.apache.beam.model.pipeline.v1.Endpoints)9 SimpleForwardingClientCall (io.grpc.ForwardingClientCall.SimpleForwardingClientCall)8 SocketAddress (java.net.SocketAddress)8