Search in sources :

Example 16 with Server

use of org.apache.beam.vendor.grpc.v1p43p2.io.grpc.Server 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 17 with Server

use of org.apache.beam.vendor.grpc.v1p43p2.io.grpc.Server 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 18 with Server

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

the class PubsubGrpcClientTest method pullOneMessage.

@Test
public void pullOneMessage() throws IOException {
    initializeClient(null, null);
    String expectedSubscription = SUBSCRIPTION.getPath();
    final PullRequest expectedRequest = PullRequest.newBuilder().setSubscription(expectedSubscription).setReturnImmediately(true).setMaxMessages(10).build();
    Timestamp timestamp = Timestamp.newBuilder().setSeconds(PUB_TIME_MS / 1000).setNanos((int) (PUB_TIME_MS % 1000) * 1000 * 1000).build();
    PubsubMessage expectedPubsubMessage = PubsubMessage.newBuilder().setMessageId(MESSAGE_ID).setData(ByteString.copyFrom(DATA.getBytes(StandardCharsets.UTF_8))).setPublishTime(timestamp).build();
    ReceivedMessage expectedReceivedMessage = ReceivedMessage.newBuilder().setMessage(expectedPubsubMessage).setAckId(ACK_ID).build();
    final PullResponse response = PullResponse.newBuilder().addAllReceivedMessages(ImmutableList.of(expectedReceivedMessage)).build();
    final List<PullRequest> requestsReceived = new ArrayList<>();
    SubscriberImplBase subscriberImplBase = new SubscriberImplBase() {

        @Override
        public void pull(PullRequest request, StreamObserver<PullResponse> responseObserver) {
            requestsReceived.add(request);
            responseObserver.onNext(response);
            responseObserver.onCompleted();
        }
    };
    Server server = InProcessServerBuilder.forName(channelName).addService(subscriberImplBase).build().start();
    try {
        List<IncomingMessage> actualMessages = client.pull(REQ_TIME_MS, SUBSCRIPTION, 10, true);
        assertEquals(1, actualMessages.size());
        IncomingMessage actualMessage = actualMessages.get(0);
        assertEquals(ACK_ID, actualMessage.ackId());
        assertEquals(DATA, actualMessage.message().getData().toStringUtf8());
        assertEquals(MESSAGE_ID, actualMessage.recordId());
        assertEquals(REQ_TIME_MS, actualMessage.requestTimeMsSinceEpoch());
        assertEquals(PUB_TIME_MS, actualMessage.timestampMsSinceEpoch());
        assertEquals(expectedRequest, Iterables.getOnlyElement(requestsReceived));
    } finally {
        server.shutdownNow();
    }
}
Also used : StreamObserver(io.grpc.stub.StreamObserver) SubscriberImplBase(com.google.pubsub.v1.SubscriberGrpc.SubscriberImplBase) Server(io.grpc.Server) PullRequest(com.google.pubsub.v1.PullRequest) ArrayList(java.util.ArrayList) ByteString(com.google.protobuf.ByteString) ReceivedMessage(com.google.pubsub.v1.ReceivedMessage) Timestamp(com.google.protobuf.Timestamp) PubsubMessage(com.google.pubsub.v1.PubsubMessage) PullResponse(com.google.pubsub.v1.PullResponse) IncomingMessage(org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.IncomingMessage) Test(org.junit.Test)

Example 19 with Server

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

the class PubsubGrpcClientTest method pullOneMessageUsingAttributes.

@Test
public void pullOneMessageUsingAttributes() throws IOException {
    initializeClient(TIMESTAMP_ATTRIBUTE, ID_ATTRIBUTE);
    String expectedSubscription = SUBSCRIPTION.getPath();
    final PullRequest expectedRequest = PullRequest.newBuilder().setSubscription(expectedSubscription).setReturnImmediately(true).setMaxMessages(10).build();
    Timestamp timestamp = Timestamp.newBuilder().setSeconds(PUB_TIME_MS / 1000).setNanos((int) (PUB_TIME_MS % 1000) * 1000 * 1000).build();
    PubsubMessage expectedPubsubMessage = PubsubMessage.newBuilder().setMessageId(MESSAGE_ID).setData(ByteString.copyFrom(DATA.getBytes(StandardCharsets.UTF_8))).setPublishTime(timestamp).putAllAttributes(ATTRIBUTES).putAllAttributes(ImmutableMap.of(TIMESTAMP_ATTRIBUTE, String.valueOf(MESSAGE_TIME_MS), ID_ATTRIBUTE, RECORD_ID)).build();
    ReceivedMessage expectedReceivedMessage = ReceivedMessage.newBuilder().setMessage(expectedPubsubMessage).setAckId(ACK_ID).build();
    final PullResponse response = PullResponse.newBuilder().addAllReceivedMessages(ImmutableList.of(expectedReceivedMessage)).build();
    final List<PullRequest> requestsReceived = new ArrayList<>();
    SubscriberImplBase subscriberImplBase = new SubscriberImplBase() {

        @Override
        public void pull(PullRequest request, StreamObserver<PullResponse> responseObserver) {
            requestsReceived.add(request);
            responseObserver.onNext(response);
            responseObserver.onCompleted();
        }
    };
    Server server = InProcessServerBuilder.forName(channelName).addService(subscriberImplBase).build().start();
    try {
        List<IncomingMessage> actualMessages = client.pull(REQ_TIME_MS, SUBSCRIPTION, 10, true);
        assertEquals(1, actualMessages.size());
        IncomingMessage actualMessage = actualMessages.get(0);
        assertEquals(ACK_ID, actualMessage.ackId());
        assertEquals(DATA, actualMessage.message().getData().toStringUtf8());
        assertEquals(RECORD_ID, actualMessage.recordId());
        assertEquals(REQ_TIME_MS, actualMessage.requestTimeMsSinceEpoch());
        assertEquals(MESSAGE_TIME_MS, actualMessage.timestampMsSinceEpoch());
        assertEquals(expectedRequest, Iterables.getOnlyElement(requestsReceived));
    } finally {
        server.shutdownNow();
    }
}
Also used : StreamObserver(io.grpc.stub.StreamObserver) SubscriberImplBase(com.google.pubsub.v1.SubscriberGrpc.SubscriberImplBase) Server(io.grpc.Server) PullRequest(com.google.pubsub.v1.PullRequest) ArrayList(java.util.ArrayList) ByteString(com.google.protobuf.ByteString) ReceivedMessage(com.google.pubsub.v1.ReceivedMessage) Timestamp(com.google.protobuf.Timestamp) PubsubMessage(com.google.pubsub.v1.PubsubMessage) PullResponse(com.google.pubsub.v1.PullResponse) IncomingMessage(org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.IncomingMessage) Test(org.junit.Test)

Example 20 with Server

use of org.apache.beam.vendor.grpc.v1p43p2.io.grpc.Server 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)76 Server (io.grpc.Server)68 IOException (java.io.IOException)27 ByteString (org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.ByteString)24 ArrayList (java.util.ArrayList)21 StreamObserver (org.apache.beam.vendor.grpc.v1p43p2.io.grpc.stub.StreamObserver)21 ExecutionException (java.util.concurrent.ExecutionException)20 TimeoutException (java.util.concurrent.TimeoutException)20 CountDownLatch (java.util.concurrent.CountDownLatch)19 ManagedChannel (org.apache.beam.vendor.grpc.v1p43p2.io.grpc.ManagedChannel)19 Server (org.apache.beam.vendor.grpc.v1p43p2.io.grpc.Server)18 StatusException (io.grpc.StatusException)17 Metadata (io.grpc.Metadata)12 List (java.util.List)12 BeamFnApi (org.apache.beam.model.fnexecution.v1.BeamFnApi)12 FilterChain (io.grpc.xds.EnvoyServerProtoData.FilterChain)11 ExecutorService (java.util.concurrent.ExecutorService)11 ParallelInstruction (com.google.api.services.dataflow.model.ParallelInstruction)10 ServerCall (io.grpc.ServerCall)10 StreamObserver (io.grpc.stub.StreamObserver)10