Examples with CallStreamObserver org.apache.beam.vendor.grpc.v1p43p2.io.grpc.stub.CallStreamObserver used on opensource projects

Example 6 with CallStreamObserver

use of org.apache.beam.vendor.grpc.v1p43p2.io.grpc.stub.CallStreamObserver in project beam by apache.

the class BeamFnDataGrpcClientTest method testForInboundConsumer.

@Test
public void testForInboundConsumer() throws Exception {
    CountDownLatch waitForClientToConnect = new CountDownLatch(1);
    Collection<WindowedValue<String>> inboundValuesA = new ConcurrentLinkedQueue<>();
    Collection<WindowedValue<String>> inboundValuesB = new ConcurrentLinkedQueue<>();
    Collection<BeamFnApi.Elements> inboundServerValues = new ConcurrentLinkedQueue<>();
    AtomicReference<StreamObserver<BeamFnApi.Elements>> outboundServerObserver = new AtomicReference<>();
    CallStreamObserver<BeamFnApi.Elements> inboundServerObserver = TestStreams.withOnNext(inboundServerValues::add).build();
    Endpoints.ApiServiceDescriptor apiServiceDescriptor = Endpoints.ApiServiceDescriptor.newBuilder().setUrl(this.getClass().getName() + "-" + UUID.randomUUID()).build();
    Server server = InProcessServerBuilder.forName(apiServiceDescriptor.getUrl()).addService(new BeamFnDataGrpc.BeamFnDataImplBase() {

        @Override
        public StreamObserver<BeamFnApi.Elements> data(StreamObserver<BeamFnApi.Elements> outboundObserver) {
            outboundServerObserver.set(outboundObserver);
            waitForClientToConnect.countDown();
            return inboundServerObserver;
        }
    }).build();
    server.start();
    try {
        ManagedChannel channel = InProcessChannelBuilder.forName(apiServiceDescriptor.getUrl()).build();
        BeamFnDataGrpcClient clientFactory = new BeamFnDataGrpcClient(PipelineOptionsFactory.create(), (Endpoints.ApiServiceDescriptor descriptor) -> channel, OutboundObserverFactory.trivial());
        BeamFnDataInboundObserver2 observerA = BeamFnDataInboundObserver2.forConsumers(Arrays.asList(DataEndpoint.create(TRANSFORM_ID_A, CODER, inboundValuesA::add)), Collections.emptyList());
        BeamFnDataInboundObserver2 observerB = BeamFnDataInboundObserver2.forConsumers(Arrays.asList(DataEndpoint.create(TRANSFORM_ID_B, CODER, inboundValuesB::add)), Collections.emptyList());
        clientFactory.registerReceiver(INSTRUCTION_ID_A, Arrays.asList(apiServiceDescriptor), observerA);
        waitForClientToConnect.await();
        outboundServerObserver.get().onNext(ELEMENTS_A_1);
        // Purposefully transmit some data before the consumer for B is bound showing that
        // data is not lost
        outboundServerObserver.get().onNext(ELEMENTS_B_1);
        Thread.sleep(100);
        clientFactory.registerReceiver(INSTRUCTION_ID_B, Arrays.asList(apiServiceDescriptor), observerB);
        // Show that out of order stream completion can occur.
        observerB.awaitCompletion();
        assertThat(inboundValuesB, contains(valueInGlobalWindow("JKL"), valueInGlobalWindow("MNO")));
        outboundServerObserver.get().onNext(ELEMENTS_A_2);
        observerA.awaitCompletion();
        assertThat(inboundValuesA, contains(valueInGlobalWindow("ABC"), valueInGlobalWindow("DEF"), valueInGlobalWindow("GHI")));
    } finally {
        server.shutdownNow();
    }
}

Example 7 with CallStreamObserver

use of org.apache.beam.vendor.grpc.v1p43p2.io.grpc.stub.CallStreamObserver in project beam by apache.

the class BeamFnLoggingClientTest method testWhenServerFailsThatClientIsAbleToCleanup.

@Test
public void testWhenServerFailsThatClientIsAbleToCleanup() throws Exception {
    BeamFnLoggingMDC.setInstructionId("instruction-1");
    Collection<BeamFnApi.LogEntry> values = new ConcurrentLinkedQueue<>();
    AtomicReference<StreamObserver<BeamFnApi.LogControl>> outboundServerObserver = new AtomicReference<>();
    CallStreamObserver<BeamFnApi.LogEntry.List> inboundServerObserver = TestStreams.withOnNext((BeamFnApi.LogEntry.List logEntries) -> values.addAll(logEntries.getLogEntriesList())).build();
    Endpoints.ApiServiceDescriptor apiServiceDescriptor = Endpoints.ApiServiceDescriptor.newBuilder().setUrl(this.getClass().getName() + "-" + UUID.randomUUID().toString()).build();
    Server server = InProcessServerBuilder.forName(apiServiceDescriptor.getUrl()).addService(new BeamFnLoggingGrpc.BeamFnLoggingImplBase() {

        @Override
        public StreamObserver<BeamFnApi.LogEntry.List> logging(StreamObserver<BeamFnApi.LogControl> outboundObserver) {
            outboundServerObserver.set(outboundObserver);
            outboundObserver.onError(Status.INTERNAL.withDescription("TEST ERROR").asException());
            return inboundServerObserver;
        }
    }).build();
    server.start();
    ManagedChannel channel = InProcessChannelBuilder.forName(apiServiceDescriptor.getUrl()).build();
    // Keep a strong reference to the loggers. Otherwise the call to client.close()
    // removes the only reference and the logger may get GC'd before the assertions (BEAM-4136).
    Logger rootLogger = null;
    Logger configuredLogger = null;
    try {
        BeamFnLoggingClient client = new BeamFnLoggingClient(PipelineOptionsFactory.fromArgs(new String[] { "--defaultSdkHarnessLogLevel=OFF", "--sdkHarnessLogLevelOverrides={\"ConfiguredLogger\": \"DEBUG\"}" }).create(), apiServiceDescriptor, (Endpoints.ApiServiceDescriptor descriptor) -> channel);
        rootLogger = LogManager.getLogManager().getLogger("");
        configuredLogger = LogManager.getLogManager().getLogger("ConfiguredLogger");
        thrown.expectMessage("TEST ERROR");
        client.close();
    } finally {
        assertNotNull("rootLogger should be initialized before exception", rootLogger);
        assertNotNull("configuredLogger should be initialized before exception", rootLogger);
        // Verify that after close, log levels are reset.
        assertEquals(Level.INFO, rootLogger.getLevel());
        assertNull(configuredLogger.getLevel());
        assertTrue(channel.isShutdown());
        server.shutdownNow();
    }
}

Example 8 with CallStreamObserver

use of org.apache.beam.vendor.grpc.v1p43p2.io.grpc.stub.CallStreamObserver in project beam by apache.

the class BeamFnDataGrpcClientTest method testForInboundConsumerThatThrows.

@Test
public void testForInboundConsumerThatThrows() throws Exception {
    CountDownLatch waitForClientToConnect = new CountDownLatch(1);
    AtomicInteger consumerInvoked = new AtomicInteger();
    Collection<BeamFnApi.Elements> inboundServerValues = new ConcurrentLinkedQueue<>();
    AtomicReference<StreamObserver<BeamFnApi.Elements>> outboundServerObserver = new AtomicReference<>();
    CallStreamObserver<BeamFnApi.Elements> inboundServerObserver = TestStreams.withOnNext(inboundServerValues::add).build();
    Endpoints.ApiServiceDescriptor apiServiceDescriptor = Endpoints.ApiServiceDescriptor.newBuilder().setUrl(this.getClass().getName() + "-" + UUID.randomUUID()).build();
    Server server = InProcessServerBuilder.forName(apiServiceDescriptor.getUrl()).addService(new BeamFnDataGrpc.BeamFnDataImplBase() {

        @Override
        public StreamObserver<BeamFnApi.Elements> data(StreamObserver<BeamFnApi.Elements> outboundObserver) {
            outboundServerObserver.set(outboundObserver);
            waitForClientToConnect.countDown();
            return inboundServerObserver;
        }
    }).build();
    server.start();
    RuntimeException exceptionToThrow = new RuntimeException("TestFailure");
    try {
        ManagedChannel channel = InProcessChannelBuilder.forName(apiServiceDescriptor.getUrl()).build();
        BeamFnDataGrpcClient clientFactory = new BeamFnDataGrpcClient(PipelineOptionsFactory.create(), (Endpoints.ApiServiceDescriptor descriptor) -> channel, OutboundObserverFactory.trivial());
        BeamFnDataInboundObserver2 observer = BeamFnDataInboundObserver2.forConsumers(Arrays.asList(DataEndpoint.create(TRANSFORM_ID_A, CODER, t -> {
            consumerInvoked.incrementAndGet();
            throw exceptionToThrow;
        })), Collections.emptyList());
        clientFactory.registerReceiver(INSTRUCTION_ID_A, Arrays.asList(apiServiceDescriptor), observer);
        waitForClientToConnect.await();
        // This first message should cause a failure afterwards all other messages are dropped.
        outboundServerObserver.get().onNext(ELEMENTS_A_1);
        outboundServerObserver.get().onNext(ELEMENTS_A_2);
        try {
            observer.awaitCompletion();
            fail("Expected channel to fail");
        } catch (Exception e) {
            assertEquals(exceptionToThrow, e);
        }
        // The server should not have received any values
        assertThat(inboundServerValues, empty());
        // The consumer should have only been invoked once
        assertEquals(1, consumerInvoked.get());
    } finally {
        server.shutdownNow();
    }
}

Example 9 with CallStreamObserver

use of org.apache.beam.vendor.grpc.v1p43p2.io.grpc.stub.CallStreamObserver in project beam by apache.

the class BeamFnDataGrpcClientTest method testForOutboundConsumer.

@Test
public void testForOutboundConsumer() throws Exception {
    CountDownLatch waitForInboundServerValuesCompletion = new CountDownLatch(2);
    Collection<BeamFnApi.Elements> inboundServerValues = new ConcurrentLinkedQueue<>();
    CallStreamObserver<BeamFnApi.Elements> inboundServerObserver = TestStreams.withOnNext((BeamFnApi.Elements t) -> {
        inboundServerValues.add(t);
        waitForInboundServerValuesCompletion.countDown();
    }).build();
    Endpoints.ApiServiceDescriptor apiServiceDescriptor = Endpoints.ApiServiceDescriptor.newBuilder().setUrl(this.getClass().getName() + "-" + UUID.randomUUID()).build();
    Server server = InProcessServerBuilder.forName(apiServiceDescriptor.getUrl()).addService(new BeamFnDataGrpc.BeamFnDataImplBase() {

        @Override
        public StreamObserver<BeamFnApi.Elements> data(StreamObserver<BeamFnApi.Elements> outboundObserver) {
            return inboundServerObserver;
        }
    }).build();
    server.start();
    try {
        ManagedChannel channel = InProcessChannelBuilder.forName(apiServiceDescriptor.getUrl()).build();
        BeamFnDataGrpcClient clientFactory = new BeamFnDataGrpcClient(PipelineOptionsFactory.fromArgs(new String[] { "--experiments=data_buffer_size_limit=20" }).create(), (Endpoints.ApiServiceDescriptor descriptor) -> channel, OutboundObserverFactory.trivial());
        BeamFnDataOutboundAggregator aggregator = clientFactory.createOutboundAggregator(apiServiceDescriptor, () -> INSTRUCTION_ID_A, false);
        FnDataReceiver<WindowedValue<String>> fnDataReceiver = aggregator.registerOutputDataLocation(TRANSFORM_ID_A, CODER);
        fnDataReceiver.accept(valueInGlobalWindow("ABC"));
        fnDataReceiver.accept(valueInGlobalWindow("DEF"));
        fnDataReceiver.accept(valueInGlobalWindow("GHI"));
        aggregator.sendOrCollectBufferedDataAndFinishOutboundStreams();
        waitForInboundServerValuesCompletion.await();
        assertThat(inboundServerValues, contains(ELEMENTS_A_1, ELEMENTS_A_2));
    } finally {
        server.shutdownNow();
    }
}