Search in sources :

Example 46 with StreamObserver

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

the class BeamFnLoggingClientTest method testWhenServerHangsUpEarlyThatClientIsAbleCleanup.

@Test
public void testWhenServerHangsUpEarlyThatClientIsAbleCleanup() 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.onCompleted();
            return inboundServerObserver;
        }
    }).build();
    server.start();
    ManagedChannel channel = InProcessChannelBuilder.forName(apiServiceDescriptor.getUrl()).build();
    try {
        BeamFnLoggingClient client = new BeamFnLoggingClient(PipelineOptionsFactory.fromArgs(new String[] { "--defaultSdkHarnessLogLevel=OFF", "--sdkHarnessLogLevelOverrides={\"ConfiguredLogger\": \"DEBUG\"}" }).create(), apiServiceDescriptor, (Endpoints.ApiServiceDescriptor descriptor) -> channel);
        // Keep a strong reference to the loggers in this block. Otherwise the call to client.close()
        // removes the only reference and the logger may get GC'd before the assertions (BEAM-4136).
        Logger rootLogger = LogManager.getLogManager().getLogger("");
        Logger configuredLogger = LogManager.getLogManager().getLogger("ConfiguredLogger");
        client.close();
        // Verify that after close, log levels are reset.
        assertEquals(Level.INFO, rootLogger.getLevel());
        assertNull(configuredLogger.getLevel());
    } finally {
        assertTrue(channel.isShutdown());
        server.shutdownNow();
    }
}
Also used : CallStreamObserver(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.stub.CallStreamObserver) StreamObserver(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.stub.StreamObserver) Server(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.Server) BeamFnApi(org.apache.beam.model.fnexecution.v1.BeamFnApi) AtomicReference(java.util.concurrent.atomic.AtomicReference) Logger(java.util.logging.Logger) Endpoints(org.apache.beam.model.pipeline.v1.Endpoints) ManagedChannel(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.ManagedChannel) ConcurrentLinkedQueue(java.util.concurrent.ConcurrentLinkedQueue) Test(org.junit.Test)

Example 47 with StreamObserver

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

the class FnHarnessTest method testLaunchFnHarnessAndTeardownCleanly.

@Test
// failure will cause test to timeout.
@SuppressWarnings("FutureReturnValueIgnored")
public void testLaunchFnHarnessAndTeardownCleanly() throws Exception {
    Function<String, String> environmentVariableMock = mock(Function.class);
    PipelineOptions options = PipelineOptionsFactory.create();
    when(environmentVariableMock.apply("HARNESS_ID")).thenReturn("id");
    when(environmentVariableMock.apply("PIPELINE_OPTIONS")).thenReturn(PipelineOptionsTranslation.toJson(options));
    List<BeamFnApi.LogEntry> logEntries = new ArrayList<>();
    List<BeamFnApi.InstructionResponse> instructionResponses = mock(List.class);
    BeamFnLoggingGrpc.BeamFnLoggingImplBase loggingService = new BeamFnLoggingGrpc.BeamFnLoggingImplBase() {

        @Override
        public StreamObserver<BeamFnApi.LogEntry.List> logging(StreamObserver<LogControl> responseObserver) {
            return TestStreams.withOnNext((BeamFnApi.LogEntry.List entries) -> logEntries.addAll(entries.getLogEntriesList())).withOnCompleted(responseObserver::onCompleted).build();
        }
    };
    BeamFnControlGrpc.BeamFnControlImplBase controlService = new BeamFnControlGrpc.BeamFnControlImplBase() {

        @Override
        public StreamObserver<InstructionResponse> control(StreamObserver<InstructionRequest> responseObserver) {
            CountDownLatch waitForResponses = new CountDownLatch(1);
            options.as(GcsOptions.class).getExecutorService().submit(() -> {
                responseObserver.onNext(INSTRUCTION_REQUEST);
                Uninterruptibles.awaitUninterruptibly(waitForResponses);
                responseObserver.onCompleted();
            });
            return TestStreams.withOnNext((InstructionResponse t) -> {
                instructionResponses.add(t);
                waitForResponses.countDown();
            }).withOnCompleted(waitForResponses::countDown).build();
        }
    };
    Server loggingServer = ServerBuilder.forPort(0).addService(loggingService).build();
    loggingServer.start();
    try {
        Server controlServer = ServerBuilder.forPort(0).addService(controlService).build();
        controlServer.start();
        try {
            Endpoints.ApiServiceDescriptor loggingDescriptor = Endpoints.ApiServiceDescriptor.newBuilder().setUrl("localhost:" + loggingServer.getPort()).build();
            Endpoints.ApiServiceDescriptor controlDescriptor = Endpoints.ApiServiceDescriptor.newBuilder().setUrl("localhost:" + controlServer.getPort()).build();
            when(environmentVariableMock.apply("LOGGING_API_SERVICE_DESCRIPTOR")).thenReturn(TextFormat.printToString(loggingDescriptor));
            when(environmentVariableMock.apply("CONTROL_API_SERVICE_DESCRIPTOR")).thenReturn(TextFormat.printToString(controlDescriptor));
            FnHarness.main(environmentVariableMock);
        } finally {
            controlServer.shutdownNow();
        }
    } finally {
        loggingServer.shutdownNow();
    }
    // Verify that we first run onStartup functions before even reading the environment, and that
    // we then call beforeProcessing functions before executing instructions.
    InOrder inOrder = inOrder(onStartupMock, beforeProcessingMock, environmentVariableMock, instructionResponses);
    inOrder.verify(onStartupMock).run();
    inOrder.verify(environmentVariableMock, atLeastOnce()).apply(any());
    inOrder.verify(beforeProcessingMock).accept(any());
    inOrder.verify(instructionResponses).add(INSTRUCTION_RESPONSE);
}
Also used : StreamObserver(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.stub.StreamObserver) InOrder(org.mockito.InOrder) BeamFnLoggingGrpc(org.apache.beam.model.fnexecution.v1.BeamFnLoggingGrpc) Server(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.Server) BeamFnApi(org.apache.beam.model.fnexecution.v1.BeamFnApi) ArrayList(java.util.ArrayList) InstructionResponse(org.apache.beam.model.fnexecution.v1.BeamFnApi.InstructionResponse) CountDownLatch(java.util.concurrent.CountDownLatch) Endpoints(org.apache.beam.model.pipeline.v1.Endpoints) PipelineOptions(org.apache.beam.sdk.options.PipelineOptions) ArrayList(java.util.ArrayList) List(java.util.List) BeamFnControlGrpc(org.apache.beam.model.fnexecution.v1.BeamFnControlGrpc) Test(org.junit.Test)

Example 48 with StreamObserver

use of org.apache.beam.vendor.grpc.v1p43p2.io.grpc.stub.StreamObserver 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();
    Endpoints.ApiServiceDescriptor apiServiceDescriptor = Endpoints.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 {
        EnumMap<BeamFnApi.InstructionRequest.RequestCase, ThrowingFunction<BeamFnApi.InstructionRequest, BeamFnApi.InstructionResponse.Builder>> handlers = new EnumMap<>(BeamFnApi.InstructionRequest.RequestCase.class);
        handlers.put(BeamFnApi.InstructionRequest.RequestCase.PROCESS_BUNDLE, value -> {
            assertEquals(value.getInstructionId(), BeamFnLoggingMDC.getInstructionId());
            return BeamFnApi.InstructionResponse.newBuilder().setProcessBundle(BeamFnApi.ProcessBundleResponse.getDefaultInstance());
        });
        handlers.put(BeamFnApi.InstructionRequest.RequestCase.REGISTER, value -> {
            assertEquals(value.getInstructionId(), BeamFnLoggingMDC.getInstructionId());
            throw FAILURE;
        });
        ExecutorService executor = Executors.newCachedThreadPool();
        BeamFnControlClient client = new BeamFnControlClient(apiServiceDescriptor, ManagedChannelFactory.createInProcess(), OutboundObserverFactory.trivial(), executor, handlers);
        // Get the connected client and attempt to send and receive an instruction
        StreamObserver<BeamFnApi.InstructionRequest> outboundServerObserver = outboundServerObservers.take();
        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();
        client.waitForTermination();
    } finally {
        server.shutdownNow();
    }
}
Also used : CallStreamObserver(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.stub.CallStreamObserver) StreamObserver(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.stub.StreamObserver) ThrowingFunction(org.apache.beam.sdk.function.ThrowingFunction) Server(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.Server) BeamFnApi(org.apache.beam.model.fnexecution.v1.BeamFnApi) LinkedBlockingQueue(java.util.concurrent.LinkedBlockingQueue) Endpoints(org.apache.beam.model.pipeline.v1.Endpoints) AtomicBoolean(java.util.concurrent.atomic.AtomicBoolean) InstructionRequest(org.apache.beam.model.fnexecution.v1.BeamFnApi.InstructionRequest) ExecutorService(java.util.concurrent.ExecutorService) EnumMap(java.util.EnumMap) Test(org.junit.Test)

Example 49 with StreamObserver

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

the class BeamFnControlClientTest method testJavaErrorResponse.

@Test
public void testJavaErrorResponse() throws Exception {
    BlockingQueue<StreamObserver<BeamFnApi.InstructionRequest>> outboundServerObservers = new LinkedBlockingQueue<>();
    BlockingQueue<Throwable> error = new LinkedBlockingQueue<>();
    CallStreamObserver<BeamFnApi.InstructionResponse> inboundServerObserver = TestStreams.<BeamFnApi.InstructionResponse>withOnNext(response -> fail(String.format("Unexpected Response %s", response))).withOnError(error::add).build();
    Endpoints.ApiServiceDescriptor apiServiceDescriptor = Endpoints.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 {
        EnumMap<BeamFnApi.InstructionRequest.RequestCase, ThrowingFunction<BeamFnApi.InstructionRequest, BeamFnApi.InstructionResponse.Builder>> handlers = new EnumMap<>(BeamFnApi.InstructionRequest.RequestCase.class);
        handlers.put(BeamFnApi.InstructionRequest.RequestCase.REGISTER, value -> {
            assertEquals(value.getInstructionId(), BeamFnLoggingMDC.getInstructionId());
            throw new Error("Test Error");
        });
        ExecutorService executor = Executors.newCachedThreadPool();
        BeamFnControlClient client = new BeamFnControlClient(apiServiceDescriptor, ManagedChannelFactory.createInProcess(), OutboundObserverFactory.trivial(), executor, handlers);
        // Get the connected client and attempt to send and receive an instruction
        StreamObserver<BeamFnApi.InstructionRequest> outboundServerObserver = outboundServerObservers.take();
        // Ensure that all exceptions are caught and translated to failures
        outboundServerObserver.onNext(InstructionRequest.newBuilder().setInstructionId("0").setRegister(RegisterRequest.getDefaultInstance()).build());
        // There should be an error reported to the StreamObserver.
        assertThat(error.take(), not(nullValue()));
        // Ensure that the client shuts down when an Error is thrown from the harness
        try {
            client.waitForTermination();
            throw new IllegalStateException("The future should have terminated with an error");
        } catch (ExecutionException errorWrapper) {
            assertThat(errorWrapper.getCause().getMessage(), containsString("Test Error"));
        }
    } finally {
        server.shutdownNow();
    }
}
Also used : CallStreamObserver(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.stub.CallStreamObserver) StreamObserver(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.stub.StreamObserver) ThrowingFunction(org.apache.beam.sdk.function.ThrowingFunction) Server(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.Server) BeamFnApi(org.apache.beam.model.fnexecution.v1.BeamFnApi) LinkedBlockingQueue(java.util.concurrent.LinkedBlockingQueue) Endpoints(org.apache.beam.model.pipeline.v1.Endpoints) InstructionRequest(org.apache.beam.model.fnexecution.v1.BeamFnApi.InstructionRequest) ExecutorService(java.util.concurrent.ExecutorService) ExecutionException(java.util.concurrent.ExecutionException) EnumMap(java.util.EnumMap) Test(org.junit.Test)

Example 50 with StreamObserver

use of org.apache.beam.vendor.grpc.v1p43p2.io.grpc.stub.StreamObserver 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();
    }
}
Also used : CallStreamObserver(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.stub.CallStreamObserver) StreamObserver(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.stub.StreamObserver) Server(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.Server) BeamFnApi(org.apache.beam.model.fnexecution.v1.BeamFnApi) AtomicReference(java.util.concurrent.atomic.AtomicReference) CountDownLatch(java.util.concurrent.CountDownLatch) BeamFnDataInboundObserver2(org.apache.beam.sdk.fn.data.BeamFnDataInboundObserver2) Endpoints(org.apache.beam.model.pipeline.v1.Endpoints) WindowedValue(org.apache.beam.sdk.util.WindowedValue) ManagedChannel(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.ManagedChannel) ConcurrentLinkedQueue(java.util.concurrent.ConcurrentLinkedQueue) Test(org.junit.Test)

Aggregations

StreamObserver (io.grpc.stub.StreamObserver)133 Test (org.junit.Test)95 CountDownLatch (java.util.concurrent.CountDownLatch)50 ArrayList (java.util.ArrayList)47 AtomicReference (java.util.concurrent.atomic.AtomicReference)38 StreamObserver (org.apache.beam.vendor.grpc.v1p43p2.io.grpc.stub.StreamObserver)27 StatusRuntimeException (io.grpc.StatusRuntimeException)26 Status (io.grpc.Status)20 List (java.util.List)18 BeamFnApi (org.apache.beam.model.fnexecution.v1.BeamFnApi)18 ManagedChannel (org.apache.beam.vendor.grpc.v1p43p2.io.grpc.ManagedChannel)18 CompletableFuture (java.util.concurrent.CompletableFuture)17 ExecutorService (java.util.concurrent.ExecutorService)16 SegmentId (io.pravega.controller.stream.api.grpc.v1.Controller.SegmentId)14 ServerRequest (io.pravega.controller.stream.api.grpc.v1.Controller.ServerRequest)14 VisibleForTesting (com.google.common.annotations.VisibleForTesting)12 Strings (com.google.common.base.Strings)12 Throwables (com.google.common.base.Throwables)12 ImmutableMap (com.google.common.collect.ImmutableMap)12 AuthHandler (io.pravega.auth.AuthHandler)12