Search in sources :

Example 11 with ApiServiceDescriptor

use of org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor 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 12 with ApiServiceDescriptor

use of org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor 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)

Example 13 with ApiServiceDescriptor

use of org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor in project beam by apache.

the class BeamFnStateGrpcClientCacheTest method testCachingOfClient.

@Test
public void testCachingOfClient() throws Exception {
    Endpoints.ApiServiceDescriptor otherApiServiceDescriptor = Endpoints.ApiServiceDescriptor.newBuilder().setUrl(apiServiceDescriptor.getUrl() + "-other").build();
    Server testServer2 = InProcessServerBuilder.forName(otherApiServiceDescriptor.getUrl()).addService(new BeamFnStateGrpc.BeamFnStateImplBase() {

        @Override
        public StreamObserver<StateRequest> state(StreamObserver<StateResponse> outboundObserver) {
            throw new IllegalStateException("Unexpected in test.");
        }
    }).build();
    testServer2.start();
    try {
        assertSame(clientCache.forApiServiceDescriptor(apiServiceDescriptor), clientCache.forApiServiceDescriptor(apiServiceDescriptor));
        assertNotSame(clientCache.forApiServiceDescriptor(apiServiceDescriptor), clientCache.forApiServiceDescriptor(otherApiServiceDescriptor));
    } finally {
        testServer2.shutdownNow();
    }
}
Also used : Endpoints(org.apache.beam.model.pipeline.v1.Endpoints) CallStreamObserver(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.stub.CallStreamObserver) StreamObserver(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.stub.StreamObserver) StateRequest(org.apache.beam.model.fnexecution.v1.BeamFnApi.StateRequest) Server(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.Server) Test(org.junit.Test)

Example 14 with ApiServiceDescriptor

use of org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor in project beam by apache.

the class ManagedChannelFactoryTest method testEpollDomainSocketChannel.

@Test
public void testEpollDomainSocketChannel() throws Exception {
    assumeTrue(SystemUtils.IS_OS_LINUX);
    assertTrue(org.apache.beam.vendor.grpc.v1p43p2.io.netty.channel.epoll.Epoll.isAvailable());
    Endpoints.ApiServiceDescriptor apiServiceDescriptor = Endpoints.ApiServiceDescriptor.newBuilder().setUrl("unix://" + tmpFolder.newFile().getAbsolutePath()).build();
    ManagedChannel channel = ManagedChannelFactory.createEpoll().forDescriptor(apiServiceDescriptor);
    assertEquals(apiServiceDescriptor.getUrl().substring("unix://".length()), channel.authority());
    channel.shutdownNow();
}
Also used : Endpoints(org.apache.beam.model.pipeline.v1.Endpoints) ManagedChannel(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.ManagedChannel) Test(org.junit.Test)

Example 15 with ApiServiceDescriptor

use of org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor in project beam by apache.

the class GrpcFnServer method allocatePortAndCreateFor.

/**
 * Create {@link GrpcFnServer}s for the provided {@link FnService}s running on an arbitrary port.
 */
public static List<GrpcFnServer<? extends FnService>> allocatePortAndCreateFor(List<? extends FnService> services, ServerFactory factory) throws IOException {
    ApiServiceDescriptor.Builder apiServiceDescriptor = ApiServiceDescriptor.newBuilder();
    Server server = factory.allocateAddressAndCreate(Collections.unmodifiableList(services), apiServiceDescriptor);
    AtomicInteger countdown = new AtomicInteger(services.size());
    return Lists.transform(services, service -> new SharedGrpcFnServer<>(server, service, apiServiceDescriptor.build(), countdown));
}
Also used : ApiServiceDescriptor(org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor) Server(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.Server) AtomicInteger(java.util.concurrent.atomic.AtomicInteger)

Aggregations

Test (org.junit.Test)18 Endpoints (org.apache.beam.model.pipeline.v1.Endpoints)16 Server (org.apache.beam.vendor.grpc.v1p43p2.io.grpc.Server)12 BeamFnApi (org.apache.beam.model.fnexecution.v1.BeamFnApi)11 ApiServiceDescriptor (org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor)11 ManagedChannel (org.apache.beam.vendor.grpc.v1p43p2.io.grpc.ManagedChannel)10 StreamObserver (org.apache.beam.vendor.grpc.v1p43p2.io.grpc.stub.StreamObserver)10 CallStreamObserver (org.apache.beam.vendor.grpc.v1p43p2.io.grpc.stub.CallStreamObserver)9 ConcurrentLinkedQueue (java.util.concurrent.ConcurrentLinkedQueue)6 AtomicReference (java.util.concurrent.atomic.AtomicReference)6 WindowedValue (org.apache.beam.sdk.util.WindowedValue)6 Coder (org.apache.beam.sdk.coders.Coder)5 PTransform (org.apache.beam.model.pipeline.v1.RunnerApi.PTransform)4 ArrayList (java.util.ArrayList)3 LinkedHashMap (java.util.LinkedHashMap)3 Map (java.util.Map)3 CountDownLatch (java.util.concurrent.CountDownLatch)3 Logger (java.util.logging.Logger)3 RunnerApi (org.apache.beam.model.pipeline.v1.RunnerApi)3 ByteStringCoder (org.apache.beam.runners.fnexecution.wire.ByteStringCoder)3