Search in sources :

Example 36 with ClientInterceptor

use of org.apache.beam.vendor.grpc.v1p43p2.io.grpc.ClientInterceptor 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());
}
Also used : StreamObserver(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.stub.StreamObserver) ClientInterceptor(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.ClientInterceptor) Elements(org.apache.beam.model.fnexecution.v1.BeamFnApi.Elements) RunWith(org.junit.runner.RunWith) Channel(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.Channel) CompletableFuture(java.util.concurrent.CompletableFuture) BeamFnDataGrpc(org.apache.beam.model.fnexecution.v1.BeamFnDataGrpc) GrpcContextHeaderAccessorProvider(org.apache.beam.sdk.fn.server.GrpcContextHeaderAccessorProvider) InProcessChannelBuilder(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.inprocess.InProcessChannelBuilder) MethodDescriptor(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.MethodDescriptor) ClientCall(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.ClientCall) ApiServiceDescriptor(org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor) InProcessServerFactory(org.apache.beam.sdk.fn.server.InProcessServerFactory) Metadata(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.Metadata) Test(org.junit.Test) JUnit4(org.junit.runners.JUnit4) BeamFnApi(org.apache.beam.model.fnexecution.v1.BeamFnApi) TimeUnit(java.util.concurrent.TimeUnit) Consumer(java.util.function.Consumer) CallOptions(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.CallOptions) GrpcCleanupRule(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.testing.GrpcCleanupRule) Mockito(org.mockito.Mockito) Rule(org.junit.Rule) StreamObserver(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.stub.StreamObserver) SimpleForwardingClientCall(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.ForwardingClientCall.SimpleForwardingClientCall) ImmutableList(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList) Assert(org.junit.Assert) ApiServiceDescriptor(org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor) Channel(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.Channel) Metadata(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.Metadata) CallOptions(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.CallOptions) CompletableFuture(java.util.concurrent.CompletableFuture) ClientCall(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.ClientCall) SimpleForwardingClientCall(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.ForwardingClientCall.SimpleForwardingClientCall) ClientInterceptor(org.apache.beam.vendor.grpc.v1p43p2.io.grpc.ClientInterceptor) BeamFnDataGrpc(org.apache.beam.model.fnexecution.v1.BeamFnDataGrpc) Test(org.junit.Test)

Aggregations

ClientInterceptor (io.grpc.ClientInterceptor)34 CallOptions (io.grpc.CallOptions)23 Metadata (io.grpc.Metadata)22 Channel (io.grpc.Channel)18 MethodDescriptor (io.grpc.MethodDescriptor)15 ManagedChannel (io.grpc.ManagedChannel)14 Test (org.junit.Test)13 ClientCall (io.grpc.ClientCall)11 SimpleForwardingClientCall (io.grpc.ForwardingClientCall.SimpleForwardingClientCall)8 ChannelFactoryBuilder (com.navercorp.pinpoint.grpc.client.ChannelFactoryBuilder)5 DefaultChannelFactoryBuilder (com.navercorp.pinpoint.grpc.client.DefaultChannelFactoryBuilder)5 UnaryCallDeadlineInterceptor (com.navercorp.pinpoint.grpc.client.UnaryCallDeadlineInterceptor)5 ClientOption (com.navercorp.pinpoint.grpc.client.config.ClientOption)5 Status (io.grpc.Status)5 SslOption (com.navercorp.pinpoint.grpc.client.config.SslOption)4 ForwardingClientCall (io.grpc.ForwardingClientCall)4 SimpleForwardingClientCallListener (io.grpc.ForwardingClientCallListener.SimpleForwardingClientCallListener)4 PickSubchannelArgs (io.grpc.LoadBalancer.PickSubchannelArgs)4 SocketAddress (java.net.SocketAddress)4 AtomicLong (java.util.concurrent.atomic.AtomicLong)4