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

Example 86 with Channel

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

the class ArtifactStagingServiceTest method setUp.

@Before
public void setUp() throws Exception {
    stagingDir = tempFolder.newFolder("staging").toPath();
    stagingService = new ArtifactStagingService(ArtifactStagingService.beamFilesystemArtifactDestinationProvider(stagingDir.toString()));
    retrievalService = new ArtifactRetrievalService(TEST_BUFFER_SIZE);
    grpcCleanup.register(InProcessServerBuilder.forName("server").directExecutor().addService(stagingService).addService(retrievalService).build().start());
    ManagedChannel channel = grpcCleanup.register(InProcessChannelBuilder.forName("server").build());
    stagingStub = ArtifactStagingServiceGrpc.newStub(channel);
    retrievalBlockingStub = ArtifactRetrievalServiceGrpc.newBlockingStub(channel);
}

Example 87 with Channel

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

the class GrpcContextHeaderAccessorProviderTest method testWorkerIdOnConnect.

@SuppressWarnings("unchecked")
@Test
public void testWorkerIdOnConnect() throws Exception {
    final String worker1 = "worker1";
    CompletableFuture<String> workerId = new CompletableFuture<>();
    Consumer<StreamObserver<Elements>> consumer = elementsStreamObserver -> {
        workerId.complete(GrpcContextHeaderAccessorProvider.getHeaderAccessor().getSdkWorkerId());
        elementsStreamObserver.onCompleted();
    };
    TestDataService testService = new TestDataService(Mockito.mock(StreamObserver.class), consumer);
    ApiServiceDescriptor serviceDescriptor = ApiServiceDescriptor.newBuilder().setUrl("testServer").build();
    cleanupRule.register(InProcessServerFactory.create().create(ImmutableList.of(testService), serviceDescriptor));
    final Metadata.Key<String> workerIdKey = Metadata.Key.of("worker_id", Metadata.ASCII_STRING_MARSHALLER);
    Channel channel = cleanupRule.register(InProcessChannelBuilder.forName(serviceDescriptor.getUrl()).intercept(new ClientInterceptor() {

        @Override
        public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(MethodDescriptor<ReqT, RespT> method, CallOptions callOptions, Channel next) {
            ClientCall<ReqT, RespT> call = next.newCall(method, callOptions);
            return new SimpleForwardingClientCall<ReqT, RespT>(call) {

                @Override
                public void start(ClientCall.Listener<RespT> responseListener, Metadata headers) {
                    headers.put(workerIdKey, worker1);
                    super.start(responseListener, headers);
                }
            };
        }
    }).build());
    BeamFnDataGrpc.BeamFnDataStub stub = BeamFnDataGrpc.newStub(channel);
    stub.data(Mockito.mock(StreamObserver.class)).onCompleted();
    Assert.assertEquals(worker1, workerId.get());
}

Example 88 with Channel

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

the class ArtifactRetrievalServiceTest method setUp.

@Before
public void setUp() throws Exception {
    retrievalService = new ArtifactRetrievalService(TEST_BUFFER_SIZE);
    grpcCleanup.register(InProcessServerBuilder.forName("server").directExecutor().addService(retrievalService).build().start());
    ManagedChannel channel = grpcCleanup.register(InProcessChannelBuilder.forName("server").build());
    retrievalBlockingStub = ArtifactRetrievalServiceGrpc.newBlockingStub(channel);
    stagingDir = tempFolder.newFolder("staging").toPath();
}

Example 89 with Channel

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

the class GrpcLoggingServiceTest method testMultipleClientsFailingIsHandledGracefullyByServer.

@Test
public void testMultipleClientsFailingIsHandledGracefullyByServer() throws Exception {
    ConcurrentLinkedQueue<BeamFnApi.LogEntry> logs = new ConcurrentLinkedQueue<>();
    GrpcLoggingService service = GrpcLoggingService.forWriter(new CollectionAppendingLogWriter(logs));
    try (GrpcFnServer<GrpcLoggingService> server = GrpcFnServer.allocatePortAndCreateFor(service, InProcessServerFactory.create())) {
        CountDownLatch waitForTermination = new CountDownLatch(3);
        final BlockingQueue<StreamObserver<LogEntry.List>> outboundObservers = new LinkedBlockingQueue<>();
        Collection<Callable<Void>> tasks = new ArrayList<>();
        for (int i = 1; i <= 3; ++i) {
            final int instructionId = i;
            tasks.add(() -> {
                ManagedChannel channel = InProcessChannelBuilder.forName(server.getApiServiceDescriptor().getUrl()).build();
                StreamObserver<LogEntry.List> outboundObserver = BeamFnLoggingGrpc.newStub(channel).logging(TestStreams.withOnNext(messageDiscarder).withOnError(new CountDown(waitForTermination)).build());
                outboundObserver.onNext(createLogsWithIds(instructionId, -instructionId));
                outboundObservers.add(outboundObserver);
                return null;
            });
        }
        ExecutorService executorService = Executors.newCachedThreadPool();
        executorService.invokeAll(tasks);
        for (int i = 1; i <= 3; ++i) {
            outboundObservers.take().onError(new RuntimeException("Client " + i));
        }
        waitForTermination.await();
    }
}

Example 90 with Channel

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

the class GrpcDataServiceTest method testMessageReceivedBySingleClientWhenThereAreMultipleClients.

@Test
public void testMessageReceivedBySingleClientWhenThereAreMultipleClients() throws Exception {
    final LinkedBlockingQueue<Elements> clientInboundElements = new LinkedBlockingQueue<>();
    ExecutorService executorService = Executors.newCachedThreadPool();
    final CountDownLatch waitForInboundElements = new CountDownLatch(1);
    GrpcDataService service = GrpcDataService.create(PipelineOptionsFactory.create(), Executors.newCachedThreadPool(), OutboundObserverFactory.serverDirect());
    try (GrpcFnServer<GrpcDataService> server = GrpcFnServer.allocatePortAndCreateFor(service, InProcessServerFactory.create())) {
        Collection<Future<Void>> clientFutures = new ArrayList<>();
        for (int i = 0; i < 3; ++i) {
            clientFutures.add(executorService.submit(() -> {
                ManagedChannel channel = InProcessChannelBuilder.forName(server.getApiServiceDescriptor().getUrl()).directExecutor().build();
                StreamObserver<Elements> outboundObserver = BeamFnDataGrpc.newStub(channel).data(TestStreams.withOnNext(clientInboundElements::add).build());
                waitForInboundElements.await();
                outboundObserver.onCompleted();
                return null;
            }));
        }
        for (int i = 0; i < 3; ++i) {
            CloseableFnDataReceiver<WindowedValue<String>> consumer = service.send(LogicalEndpoint.data(Integer.toString(i), TRANSFORM_ID), CODER);
            consumer.accept(WindowedValue.valueInGlobalWindow("A" + i));
            consumer.accept(WindowedValue.valueInGlobalWindow("B" + i));
            consumer.accept(WindowedValue.valueInGlobalWindow("C" + i));
            consumer.close();
        }
        waitForInboundElements.countDown();
        for (Future<Void> clientFuture : clientFutures) {
            clientFuture.get();
        }
        assertThat(clientInboundElements, containsInAnyOrder(elementsWithData("0"), elementsWithData("1"), elementsWithData("2")));
    }
}