Examples with CallStreamObserver io.grpc.stub.CallStreamObserver used on opensource projects

Example 1 with CallStreamObserver

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

the class BeamFnLoggingClientTest method testLogging.

@Test
public void testLogging() throws Exception {
    AtomicBoolean clientClosedStream = new AtomicBoolean();
    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())).withOnCompleted(new Runnable() {

        @Override
        public void run() {
            // Remember that the client told us that this stream completed
            clientClosedStream.set(true);
            outboundServerObserver.get().onCompleted();
        }
    }).build();
    BeamFnApi.ApiServiceDescriptor apiServiceDescriptor = BeamFnApi.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);
            return inboundServerObserver;
        }
    }).build();
    server.start();
    try {
        ManagedChannel channel = InProcessChannelBuilder.forName(apiServiceDescriptor.getUrl()).build();
        BeamFnLoggingClient client = new BeamFnLoggingClient(PipelineOptionsFactory.fromArgs(new String[] { "--defaultWorkerLogLevel=OFF", "--workerLogLevelOverrides={\"ConfiguredLogger\": \"DEBUG\"}" }).create(), apiServiceDescriptor, (BeamFnApi.ApiServiceDescriptor descriptor) -> channel, this::createStreamForTest);
        // Ensure that log levels were correctly set.
        assertEquals(Level.OFF, LogManager.getLogManager().getLogger("").getLevel());
        assertEquals(Level.FINE, LogManager.getLogManager().getLogger("ConfiguredLogger").getLevel());
        // Should be filtered because the default log level override is OFF
        LogManager.getLogManager().getLogger("").log(FILTERED_RECORD);
        // Should not be filtered because the default log level override for ConfiguredLogger is DEBUG
        LogManager.getLogManager().getLogger("ConfiguredLogger").log(TEST_RECORD);
        LogManager.getLogManager().getLogger("ConfiguredLogger").log(TEST_RECORD_WITH_EXCEPTION);
        client.close();
        // Verify that after close, log levels are reset.
        assertEquals(Level.INFO, LogManager.getLogManager().getLogger("").getLevel());
        assertNull(LogManager.getLogManager().getLogger("ConfiguredLogger").getLevel());
        assertTrue(clientClosedStream.get());
        assertTrue(channel.isShutdown());
        assertThat(values, contains(TEST_ENTRY, TEST_ENTRY_WITH_EXCEPTION));
    } finally {
        server.shutdownNow();
    }
}

Example 2 with CallStreamObserver

use of 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();
    BeamFnApi.ApiServiceDescriptor apiServiceDescriptor = BeamFnApi.ApiServiceDescriptor.newBuilder().setUrl(this.getClass().getName() + "-" + UUID.randomUUID().toString()).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(), (BeamFnApi.ApiServiceDescriptor descriptor) -> channel, this::createStreamForTest);
        CompletableFuture<Void> readFutureA = clientFactory.forInboundConsumer(apiServiceDescriptor, KEY_A, CODER, inboundValuesA::add);
        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);
        CompletableFuture<Void> readFutureB = clientFactory.forInboundConsumer(apiServiceDescriptor, KEY_B, CODER, inboundValuesB::add);
        // Show that out of order stream completion can occur.
        readFutureB.get();
        assertThat(inboundValuesB, contains(valueInGlobalWindow("JKL"), valueInGlobalWindow("MNO")));
        outboundServerObserver.get().onNext(ELEMENTS_A_2);
        readFutureA.get();
        assertThat(inboundValuesA, contains(valueInGlobalWindow("ABC"), valueInGlobalWindow("DEF"), valueInGlobalWindow("GHI")));
    } finally {
        server.shutdownNow();
    }
}

Example 3 with CallStreamObserver

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

the class BeamFnControlClientTest method testDelegation.

@Test
public void testDelegation() throws Exception {
    AtomicBoolean clientClosedStream = new AtomicBoolean();
    BlockingQueue<BeamFnApi.InstructionResponse> values = new LinkedBlockingQueue<>();
    BlockingQueue<StreamObserver<BeamFnApi.InstructionRequest>> outboundServerObservers = new LinkedBlockingQueue<>();
    CallStreamObserver<BeamFnApi.InstructionResponse> inboundServerObserver = TestStreams.withOnNext(values::add).withOnCompleted(() -> clientClosedStream.set(true)).build();
    BeamFnApi.ApiServiceDescriptor apiServiceDescriptor = BeamFnApi.ApiServiceDescriptor.newBuilder().setUrl(this.getClass().getName() + "-" + UUID.randomUUID().toString()).build();
    Server server = InProcessServerBuilder.forName(apiServiceDescriptor.getUrl()).addService(new BeamFnControlGrpc.BeamFnControlImplBase() {

        @Override
        public StreamObserver<BeamFnApi.InstructionResponse> control(StreamObserver<BeamFnApi.InstructionRequest> outboundObserver) {
            Uninterruptibles.putUninterruptibly(outboundServerObservers, outboundObserver);
            return inboundServerObserver;
        }
    }).build();
    server.start();
    try {
        ManagedChannel channel = InProcessChannelBuilder.forName(apiServiceDescriptor.getUrl()).build();
        EnumMap<BeamFnApi.InstructionRequest.RequestCase, ThrowingFunction<BeamFnApi.InstructionRequest, BeamFnApi.InstructionResponse.Builder>> handlers = new EnumMap<>(BeamFnApi.InstructionRequest.RequestCase.class);
        handlers.put(BeamFnApi.InstructionRequest.RequestCase.PROCESS_BUNDLE, new ThrowingFunction<BeamFnApi.InstructionRequest, BeamFnApi.InstructionResponse.Builder>() {

            @Override
            public BeamFnApi.InstructionResponse.Builder apply(BeamFnApi.InstructionRequest value) throws Exception {
                return BeamFnApi.InstructionResponse.newBuilder().setProcessBundle(BeamFnApi.ProcessBundleResponse.getDefaultInstance());
            }
        });
        handlers.put(BeamFnApi.InstructionRequest.RequestCase.REGISTER, new ThrowingFunction<BeamFnApi.InstructionRequest, BeamFnApi.InstructionResponse.Builder>() {

            @Override
            public BeamFnApi.InstructionResponse.Builder apply(BeamFnApi.InstructionRequest value) throws Exception {
                throw FAILURE;
            }
        });
        BeamFnControlClient client = new BeamFnControlClient(apiServiceDescriptor, (BeamFnApi.ApiServiceDescriptor descriptor) -> channel, this::createStreamForTest, handlers);
        // Get the connected client and attempt to send and receive an instruction
        StreamObserver<BeamFnApi.InstructionRequest> outboundServerObserver = outboundServerObservers.take();
        ExecutorService executor = Executors.newCachedThreadPool();
        Future<Void> future = executor.submit(new Callable<Void>() {

            @Override
            public Void call() throws Exception {
                client.processInstructionRequests(executor);
                return null;
            }
        });
        outboundServerObserver.onNext(SUCCESSFUL_REQUEST);
        assertEquals(SUCCESSFUL_RESPONSE, values.take());
        // Ensure that conversion of an unknown request type is properly converted to a
        // failure response.
        outboundServerObserver.onNext(UNKNOWN_HANDLER_REQUEST);
        assertEquals(UNKNOWN_HANDLER_RESPONSE, values.take());
        // Ensure that all exceptions are caught and translated to failures
        outboundServerObserver.onNext(FAILURE_REQUEST);
        assertEquals(FAILURE_RESPONSE, values.take());
        // Ensure that the server completing the stream translates to the completable future
        // being completed allowing for a successful shutdown of the client.
        outboundServerObserver.onCompleted();
        future.get();
    } finally {
        server.shutdownNow();
    }
}

Example 4 with CallStreamObserver

use of 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();
    BeamFnApi.ApiServiceDescriptor apiServiceDescriptor = BeamFnApi.ApiServiceDescriptor.newBuilder().setUrl(this.getClass().getName() + "-" + UUID.randomUUID().toString()).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(), (BeamFnApi.ApiServiceDescriptor descriptor) -> channel, this::createStreamForTest);
        CompletableFuture<Void> readFuture = clientFactory.forInboundConsumer(apiServiceDescriptor, KEY_A, CODER, new ThrowingConsumer<WindowedValue<String>>() {

            @Override
            public void accept(WindowedValue<String> t) throws Exception {
                consumerInvoked.incrementAndGet();
                throw exceptionToThrow;
            }
        });
        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 {
            readFuture.get();
            fail("Expected channel to fail");
        } catch (ExecutionException e) {
            assertEquals(exceptionToThrow, e.getCause());
        }
        // 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 5 with CallStreamObserver

use of 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(new Consumer<BeamFnApi.Elements>() {

        @Override
        public void accept(BeamFnApi.Elements t) {
            inboundServerValues.add(t);
            waitForInboundServerValuesCompletion.countDown();
        }
    }).build();
    BeamFnApi.ApiServiceDescriptor apiServiceDescriptor = BeamFnApi.ApiServiceDescriptor.newBuilder().setUrl(this.getClass().getName() + "-" + UUID.randomUUID().toString()).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=beam_fn_api_data_buffer_limit=20" }).create(), (BeamFnApi.ApiServiceDescriptor descriptor) -> channel, this::createStreamForTest);
        try (CloseableThrowingConsumer<WindowedValue<String>> consumer = clientFactory.forOutboundConsumer(apiServiceDescriptor, KEY_A, CODER)) {
            consumer.accept(valueInGlobalWindow("ABC"));
            consumer.accept(valueInGlobalWindow("DEF"));
            consumer.accept(valueInGlobalWindow("GHI"));
        }
        waitForInboundServerValuesCompletion.await();
        assertThat(inboundServerValues, contains(ELEMENTS_A_1, ELEMENTS_A_2));
    } finally {
        server.shutdownNow();
    }
}