io.grpc.ClientInterceptor Java Examples

private static ClientInterceptor metadataInterceptor(Map<String, Object> metaDataMap) {
    return new ClientInterceptor() {
        @Override
        public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(
                final MethodDescriptor<ReqT, RespT> method, CallOptions callOptions, final Channel next) {

            return new CheckedForwardingClientCall<ReqT, RespT>(next.newCall(method, callOptions)) {
                @Override
                protected void checkedStart(Listener<RespT> responseListener, Metadata headers) {
                    metaDataMap.forEach((k, v) -> {
                        Key<String> mKey = Key.of(k, ASCII_STRING_MARSHALLER);
                        headers.put(mKey, String.valueOf(v));
                    });
                    delegate().start(responseListener, headers);
                }
            };
        }
    };
}
 

/**
 * Gets or creates the {@link ClientInterceptor}s that are referenced in the given annotation.
 *
 * <p>
 * <b>Note:</b> This methods return value does not contain the global client interceptors because they are handled
 * by the {@link GrpcChannelFactory}.
 * </p>
 *
 * @param annotation The annotation to get the interceptors for.
 * @return A list containing the interceptors for the given annotation.
 * @throws BeansException If the referenced interceptors weren't found or could not be created.
 */
protected List<ClientInterceptor> interceptorsFromAnnotation(final GrpcClient annotation) throws BeansException {
    final List<ClientInterceptor> list = Lists.newArrayList();
    for (final Class<? extends ClientInterceptor> interceptorClass : annotation.interceptors()) {
        final ClientInterceptor clientInterceptor;
        if (this.applicationContext.getBeanNamesForType(ClientInterceptor.class).length > 0) {
            clientInterceptor = this.applicationContext.getBean(interceptorClass);
        } else {
            try {
                clientInterceptor = interceptorClass.getConstructor().newInstance();
            } catch (final Exception e) {
                throw new BeanCreationException("Failed to create interceptor instance", e);
            }
        }
        list.add(clientInterceptor);
    }
    for (final String interceptorName : annotation.interceptorNames()) {
        list.add(this.applicationContext.getBean(interceptorName, ClientInterceptor.class));
    }
    return list;
}
 

/**
 * Must be called by the subclass setup method if overridden.
 */
@Before
public void setUp() {
  startServer();
  channel = createChannel();

  blockingStub =
      TestServiceGrpc.newBlockingStub(channel).withInterceptors(tracerSetupInterceptor);
  asyncStub = TestServiceGrpc.newStub(channel).withInterceptors(tracerSetupInterceptor);

  ClientInterceptor[] additionalInterceptors = getAdditionalInterceptors();
  if (additionalInterceptors != null) {
    blockingStub = blockingStub.withInterceptors(additionalInterceptors);
    asyncStub = asyncStub.withInterceptors(additionalInterceptors);
  }

  requestHeadersCapture.set(null);
}
 

@Override
public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(
    MethodDescriptor<ReqT, RespT> method,
    CallOptions callOptions,
    Channel next) {
  ClientInterceptor binlogInterceptor = getClientInterceptor(
      method.getFullMethodName(), callOptions);
  if (binlogInterceptor == null) {
    return next.newCall(method, callOptions);
  } else {
    return InternalClientInterceptors
        .wrapClientInterceptor(
            binlogInterceptor,
            BYTEARRAY_MARSHALLER,
            BYTEARRAY_MARSHALLER)
        .interceptCall(method, callOptions, next);
  }
}
 

@VisibleForTesting
final List<ClientInterceptor> getEffectiveInterceptors() {
  List<ClientInterceptor> effectiveInterceptors =
      new ArrayList<>(this.interceptors);
  temporarilyDisableRetry = false;
  if (statsEnabled) {
    temporarilyDisableRetry = true;
    CensusStatsModule censusStats = this.censusStatsOverride;
    if (censusStats == null) {
      censusStats = new CensusStatsModule(GrpcUtil.STOPWATCH_SUPPLIER, true);
    }
    // First interceptor runs last (see ClientInterceptors.intercept()), so that no
    // other interceptor can override the tracer factory we set in CallOptions.
    effectiveInterceptors.add(
        0, censusStats.getClientInterceptor(recordStartedRpcs, recordFinishedRpcs));
  }
  if (tracingEnabled) {
    temporarilyDisableRetry = true;
    CensusTracingModule censusTracing =
        new CensusTracingModule(Tracing.getTracer(),
            Tracing.getPropagationComponent().getBinaryFormat());
    effectiveInterceptors.add(0, censusTracing.getClientInterceptor());
  }
  return effectiveInterceptors;
}
 

/**
 * Constructor a managed channel build for the given target name and interceptors.
 * @param target The target name
 * @param interceptors The interceptors
 * @return The channel builder
 */
@Bean
@Prototype
protected NettyChannelBuilder managedChannelBuilder(@Parameter String target, List<ClientInterceptor> interceptors) {
    GrpcManagedChannelConfiguration config = beanContext.findBean(GrpcManagedChannelConfiguration.class, Qualifiers.byName(target)).orElseGet(() -> {
                final GrpcDefaultManagedChannelConfiguration mcc = new GrpcDefaultManagedChannelConfiguration(
                        target,
                        beanContext.getEnvironment(),
                        executorService
                );
                beanContext.inject(mcc);
                return mcc;
            }
    );
    final NettyChannelBuilder channelBuilder = config.getChannelBuilder();
    if (CollectionUtils.isNotEmpty(interceptors)) {
        channelBuilder.intercept(interceptors);
    }
    return channelBuilder;
}
 

/**
 * Constructs a managed server channel.
 * @param server The server
 * @param executorService The executor service
 * @param clientInterceptors The client interceptors
 * @return The channel
 */
@Singleton
@Named(NAME)
@Requires(beans = GrpcEmbeddedServer.class)
@Bean(preDestroy = "shutdown")
protected ManagedChannel serverChannel(
        GrpcEmbeddedServer server,
        @javax.inject.Named(TaskExecutors.IO) ExecutorService executorService,
        List<ClientInterceptor> clientInterceptors) {
    final ManagedChannelBuilder<?> builder = ManagedChannelBuilder.forAddress(
            server.getHost(),
            server.getPort()
    ).executor(executorService);
    if (!server.getServerConfiguration().isSecure()) {
        builder.usePlaintext();
    }
    if (CollectionUtils.isNotEmpty(clientInterceptors)) {
        builder.intercept(clientInterceptors);
    }
    return builder.build();
}
 

private static ClientInterceptor metadataInterceptor(Map<String, Object> metaDataMap) {
    return new ClientInterceptor() {
        @Override
        public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(
                final MethodDescriptor<ReqT, RespT> method, CallOptions callOptions, final Channel next) {

            return new CheckedForwardingClientCall<ReqT, RespT>(next.newCall(method, callOptions)) {
                @Override
                protected void checkedStart(Listener<RespT> responseListener, Metadata headers) {
                    metaDataMap.forEach((k, v) -> {
                        Key<String> mKey = Key.of(k, ASCII_STRING_MARSHALLER);
                        headers.put(mKey, String.valueOf(v));
                    });
                    delegate().start(responseListener, headers);
                }
            };
        }
    };
}
 

/**
 * Constructor a managed channel build for the given target name and interceptors.
 * @param target The target name
 * @param interceptors The interceptors
 * @return The channel builder
 */
@Bean
@Prototype
protected NettyChannelBuilder managedChannelBuilder(@Parameter String target, List<ClientInterceptor> interceptors) {
    GrpcManagedChannelConfiguration config = beanContext.findBean(GrpcManagedChannelConfiguration.class, Qualifiers.byName(target)).orElseGet(() -> {
                final GrpcDefaultManagedChannelConfiguration mcc = new GrpcDefaultManagedChannelConfiguration(
                        target,
                        beanContext.getEnvironment(),
                        executorService
                );
                beanContext.inject(mcc);
                return mcc;
            }
    );
    final NettyChannelBuilder channelBuilder = config.getChannelBuilder();
    if (CollectionUtils.isNotEmpty(interceptors)) {
        channelBuilder.intercept(interceptors);
    }
    return channelBuilder;
}
 

/**
 * Constructs a managed server channel.
 * @param server The server
 * @param executorService The executor service
 * @param clientInterceptors The client interceptors
 * @return The channel
 */
@Singleton
@Named(NAME)
@Requires(beans = GrpcEmbeddedServer.class)
@Bean(preDestroy = "shutdown")
protected ManagedChannel serverChannel(
        GrpcEmbeddedServer server,
        @javax.inject.Named(TaskExecutors.IO) ExecutorService executorService,
        List<ClientInterceptor> clientInterceptors) {
    final ManagedChannelBuilder<?> builder = ManagedChannelBuilder.forAddress(
            server.getHost(),
            server.getPort()
    ).executor(executorService);
    if (!server.getServerConfiguration().isSecure()) {
        builder.usePlaintext();
    }
    if (CollectionUtils.isNotEmpty(clientInterceptors)) {
        builder.intercept(clientInterceptors);
    }
    return builder.build();
}
 

private static List<ClientInterceptor> getSortedInterceptors(Instance<ClientInterceptor> interceptors) {
    if (interceptors.isUnsatisfied()) {
        return Collections.emptyList();
    }

    return interceptors.stream().sorted(new Comparator<ClientInterceptor>() { // NOSONAR
        @Override
        public int compare(ClientInterceptor si1, ClientInterceptor si2) {
            int p1 = 0;
            int p2 = 0;
            if (si1 instanceof Prioritized) {
                p1 = ((Prioritized) si1).getPriority();
            }
            if (si2 instanceof Prioritized) {
                p2 = ((Prioritized) si2).getPriority();
            }
            if (si1.equals(si2)) {
                return 0;
            }
            return Integer.compare(p1, p2);
        }
    }).collect(Collectors.toList());
}
 

@Override
public Channel createChannel(final String name, final List<ClientInterceptor> customInterceptors,
        final boolean sortInterceptors) {
    final Channel channel;
    synchronized (this) {
        if (this.shutdown) {
            throw new IllegalStateException("GrpcChannelFactory is already closed!");
        }
        channel = this.channels.computeIfAbsent(name, this::newManagedChannel);
    }
    final List<ClientInterceptor> interceptors =
            Lists.newArrayList(this.globalClientInterceptorRegistry.getClientInterceptors());
    interceptors.addAll(customInterceptors);
    if (sortInterceptors) {
        this.globalClientInterceptorRegistry.sortInterceptors(interceptors);
    }
    return ClientInterceptors.interceptForward(channel, interceptors);
}
 

@Test
public void test() throws InterruptedException, ExecutionException {
    final CountDownLatch latch = new CountDownLatch(1);

    final ClientBuilder builder = Client.builder().endpoints(cluster.getClientEndpoints())
        .header("MyHeader1", "MyHeaderVal1").header("MyHeader2", "MyHeaderVal2").interceptor(new ClientInterceptor() {
            @Override
            public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(MethodDescriptor<ReqT, RespT> method,
                CallOptions callOptions, Channel next) {
                return new ForwardingClientCall.SimpleForwardingClientCall<ReqT, RespT>(
                    next.newCall(method, callOptions)) {
                    @Override
                    public void start(Listener<RespT> responseListener, Metadata headers) {
                        super.start(responseListener, headers);
                        assertThat(headers.get(Metadata.Key.of("MyHeader1", Metadata.ASCII_STRING_MARSHALLER)))
                            .isEqualTo("MyHeaderVal1");
                        assertThat(headers.get(Metadata.Key.of("MyHeader2", Metadata.ASCII_STRING_MARSHALLER)))
                            .isEqualTo("MyHeaderVal2");

                        latch.countDown();
                    }
                };
            }
        });

    try (Client client = builder.build()) {
        CompletableFuture<PutResponse> future = client.getKVClient().put(bytesOf("sample_key"), bytesOf("sample_key"));
        latch.await(1, TimeUnit.MINUTES);
        future.get();
    }
}
 

@Override
public Channel createChannel(final String name, final List<ClientInterceptor> customInterceptors,
        final boolean sortInterceptors) {
    final Channel channel;
    synchronized (this) {
        if (this.shutdown) {
            throw new IllegalStateException("GrpcChannelFactory is already closed!");
        }
        channel = this.channels.computeIfAbsent(name, this::newManagedChannel);
    }
    final List<ClientInterceptor> interceptors =
            Lists.newArrayList(this.globalClientInterceptorRegistry.getClientInterceptors());
    interceptors.addAll(customInterceptors);
    if (sortInterceptors) {
        this.globalClientInterceptorRegistry.sortInterceptors(interceptors);
    }
    return ClientInterceptors.interceptForward(channel, interceptors);
}
 

@Test
public void serverHeaderDeliveredToClient() {
  class SpyingClientInterceptor implements ClientInterceptor {
    ClientCall.Listener<?> spyListener;

    @Override
    public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(
        MethodDescriptor<ReqT, RespT> method, CallOptions callOptions, Channel next) {
      return new SimpleForwardingClientCall<ReqT, RespT>(next.newCall(method, callOptions)) {
        @Override
        public void start(Listener<RespT> responseListener, Metadata headers) {
          spyListener = responseListener =
              mock(ClientCall.Listener.class, delegatesTo(responseListener));
          super.start(responseListener, headers);
        }
      };
    }
  }

  SpyingClientInterceptor clientInterceptor = new SpyingClientInterceptor();
  GreeterBlockingStub blockingStub = GreeterGrpc.newBlockingStub(channel)
      .withInterceptors(clientInterceptor);
  ArgumentCaptor<Metadata> metadataCaptor = ArgumentCaptor.forClass(Metadata.class);

  blockingStub.sayHello(HelloRequest.getDefaultInstance());

  assertNotNull(clientInterceptor.spyListener);
  verify(clientInterceptor.spyListener).onHeaders(metadataCaptor.capture());
  assertEquals(
      "customRespondValue",
      metadataCaptor.getValue().get(HeaderServerInterceptor.CUSTOM_HEADER_KEY));
}
 

private void initRoleServiceFutureStubChannel() {
  Metadata requestHeaders = getMetadataHeaders();
  LOGGER.trace("Header attaching with stub : {}", requestHeaders);
  ClientInterceptor clientInterceptor = MetadataUtils.newAttachHeadersInterceptor(requestHeaders);
  roleServiceFutureStub =
      RoleServiceGrpc.newFutureStub(authServiceChannel).withInterceptors(clientInterceptor);
  LOGGER.trace("Header attached with stub");
}
 

InteropTask(
    Listener listener,
    ManagedChannel channel,
    List<ClientInterceptor> interceptors,
    String testCase) {
  this.listenerReference = new WeakReference<Listener>(listener);
  this.testCase = testCase;
  this.tester = new Tester(channel, interceptors);
}
 

@Override
public Channel createChannel(final String name, final List<ClientInterceptor> interceptors,
        boolean sortInterceptors) {
    final URI address = this.properties.getChannel(name).getAddress();
    if (address != null && IN_PROCESS_SCHEME.equals(address.getScheme())) {
        return this.inProcessChannelFactory.createChannel(address.getSchemeSpecificPart(), interceptors,
                sortInterceptors);
    }
    return this.alternativeChannelFactory.createChannel(name, interceptors, sortInterceptors);
}
 

/**
 * Gets the immutable and sorted list of global server interceptors.
 *
 * @return The list of globally registered server interceptors.
 */
public ImmutableList<ClientInterceptor> getClientInterceptors() {
    if (this.sortedClientInterceptors == null) {
        List<ClientInterceptor> temp = Lists.newArrayList(this.clientInterceptors);
        sortInterceptors(temp);
        this.sortedClientInterceptors = ImmutableList.copyOf(temp);
    }
    return this.sortedClientInterceptors;
}
 

TracedClient(
    ManagedChannel channel,
    long deadline,
    String compression,
    ClientInterceptor... interceptors) {
  blockingStub =
      GreeterGrpc.newBlockingStub(ClientInterceptors.intercept(channel, interceptors))
          .withDeadlineAfter(deadline, TimeUnit.MILLISECONDS)
          .withCompression(compression);
}
 

private void initAuthzServiceStubChannel(Metadata requestHeaders) {
  if (requestHeaders == null) requestHeaders = getMetadataHeaders();
  LOGGER.trace("Header attaching with stub : {}", requestHeaders);
  ClientInterceptor clientInterceptor = MetadataUtils.newAttachHeadersInterceptor(requestHeaders);
  authzServiceBlockingStub =
      AuthzServiceGrpc.newBlockingStub(authServiceChannel).withInterceptors(clientInterceptor);
  LOGGER.trace("Header attached with stub");
}
 

@Override
public void addClientInterceptors(final GlobalClientInterceptorRegistry registry) {
    this.context.getBeansWithAnnotation(GrpcGlobalClientInterceptor.class)
            .forEach((name, bean) -> {
                ClientInterceptor interceptor = (ClientInterceptor) bean;
                log.debug("Registering GlobalClientInterceptor: {} ({})", name, interceptor);
                registry.addClientInterceptors(interceptor);
            });
}
 

/**
 * The client interceptor.
 * @param configuration The configuration
 * @return The client interceptor
 */
@Requires(beans = GrpcClientTracingInterceptorConfiguration.class)
@Singleton
@Bean
protected @Nonnull ClientInterceptor clientTracingInterceptor(@Nonnull GrpcClientTracingInterceptorConfiguration configuration) {
    return configuration.getBuilder().build();
}
 

/**
 * A custom client.
 */
private CustomHeaderClient(String host, int port) {
  String target = "zookeeper:///" + GreeterGrpc.SERVICE_NAME;
  originChannel = ManagedChannelBuilder
          //.forAddress(host, port)
          .forTarget(target)
          .usePlaintext()
          .build();
  ClientInterceptor interceptor = new HeaderClientInterceptor();
  Channel channel = ClientInterceptors.intercept(originChannel, interceptor);
  blockingStub = GreeterGrpc.newBlockingStub(channel);
}
 

private ClientInterceptor authInterceptor(Algorithm alg) {
    return new ClientInterceptor() {
        @Override
        public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(MethodDescriptor<ReqT, RespT> call, CallOptions headers, Channel next) {
            return next.newCall(call, headers.withCallCredentials(JwtCallCredentials.blocking(() -> JWT
                    .create()
                    .sign(alg)
            )));
        }
    };
}
 

private HelloClientImpl(ManagedChannel channel) {
    super(channel);
    ClientInterceptor interceptor = getClientInterceptor();
    this.blockingStub = HelloGrpc.newBlockingStub(channel).withInterceptors(interceptor);
    this.futureStub = HelloGrpc.newFutureStub(channel).withInterceptors(interceptor);
    this.stub = HelloGrpc.newStub(channel).withInterceptors(interceptor);
}
 

private HelloClientImpl(ManagedChannel channel) {
    super(channel);
    ClientInterceptor interceptor = getClientInterceptor();
    this.blockingStub = HelloGrpc.newBlockingStub(channel).withInterceptors(interceptor);
    this.futureStub = HelloGrpc.newFutureStub(channel).withInterceptors(interceptor);
    this.stub = HelloGrpc.newStub(channel).withInterceptors(interceptor);
}
 

protected ClientInterceptor getClientInterceptor() {
    return new ClientInterceptor() {
        @Override
        public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(MethodDescriptor<ReqT, RespT> method, CallOptions callOptions, Channel next) {
            return new TestClientCallImpl<ReqT, RespT>(next.newCall(method, callOptions), exeptionMethod);
        }
    };
}
 

@Nullable
@Override
public ClientInterceptor getClientInterceptor(
    String fullMethodName, CallOptions callOptions) {
  BinlogHelper helperForMethod = factory.getLog(fullMethodName);
  if (helperForMethod == null) {
    return null;
  }
  return helperForMethod.getClientInterceptor(counter.getAndIncrement());
}
 

/**
 * The client interceptor.
 * @param configuration The configuration
 * @return The client interceptor
 */
@Requires(beans = GrpcClientTracingInterceptorConfiguration.class)
@Singleton
@Bean
protected @Nonnull ClientInterceptor clientTracingInterceptor(@Nonnull GrpcClientTracingInterceptorConfiguration configuration) {
    return configuration.getBuilder().build();
}