Java Code Examples for io.grpc.ClientCall#Listener

@Override
public void exitClientListenerMethod(@Nullable Throwable thrown,
                                     ClientCall.Listener<?> listener,
                                     @Nullable Span span,
                                     boolean isLastMethod) {

    if (span != null) {
        span.captureException(thrown)
            .deactivate();
    }

    if (span != null && (isLastMethod || thrown != null)) {
        // span needs to be ended when last listener method is called or on the 1st thrown exception
        span.end();
    }

    if (isLastMethod) {
        clientCallListenerSpans.remove(listener);
    }

}
 

@Override
public Span clientCallStartEnter(ClientCall<?, ?> clientCall,
                                 ClientCall.Listener<?> listener,
                                 Metadata headers) {

    // span should already have been registered
    // no other lookup by client call is required, thus removing entry
    Span span = clientCallSpans.remove(clientCall);
    if (span == null) {
        return null;
    }

    clientCallListenerSpans.put(listener, span);

    span.propagateTraceContext(headers, headerSetter);
    return span;
}
 

@Advice.OnMethodEnter(suppress = Throwable.class)
private static void onEnter(@Advice.This ClientCall<?, ?> clientCall,
                            @Advice.Argument(0) ClientCall.Listener<?> listener,
                            @Advice.Argument(1) Metadata headers,
                            @Advice.Local("span") Span span) {

    if (tracer == null || grpcHelperManager == null) {
        return;
    }

    ElasticApmAgent.ensureInstrumented(listener.getClass(), RESPONSE_LISTENER_INSTRUMENTATIONS);

    GrpcHelper helper = grpcHelperManager.getForClassLoaderOfClass(ClientCall.class);
    if (helper != null) {
        span = helper.clientCallStartEnter(clientCall, listener, headers);
    }

}
 

@Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
private static void onExit(@Advice.Argument(0) ClientCall.Listener<?> listener,
                           @Advice.Thrown @Nullable Throwable thrown,
                           @Advice.Local("span") @Nullable Span span) {

    if (span == null) {
        return;
    }
    GrpcHelper helper = grpcHelperManager.getForClassLoaderOfClass(ClientCall.class);
    if (helper != null) {
        helper.clientCallStartExit(listener, thrown);
    }
}
 

private void verifyCallListenerClosed(
    ClientCall.Listener<Integer> listener, Status.Code code, Throwable cause) {
  ArgumentCaptor<Status> captor = ArgumentCaptor.forClass(null);
  verify(listener).onClose(captor.capture(), any(Metadata.class));
  Status rpcStatus = captor.getValue();
  assertEquals(code, rpcStatus.getCode());
  assertSame(cause, rpcStatus.getCause());
  verifyNoMoreInteractions(listener);
}
 

/**
 * Creates a new delegating ClientCallListener that will wrap the given client call listener to collect metrics.
 *
 * @param delegate The original call to wrap.
 * @param registry The registry to save the metrics to.
 * @param responseCounter The counter for incoming responses.
 * @param timerFunction A function that will return a timer for a given status code.
 */
public MetricCollectingClientCallListener(
        final ClientCall.Listener<A> delegate,
        final MeterRegistry registry,
        final Counter responseCounter,
        final Function<Code, Timer> timerFunction) {
    super(delegate);
    this.responseCounter = responseCounter;
    this.timerFunction = timerFunction;
    this.timerSample = Timer.start(registry);
}
 

@Advice.OnMethodEnter(suppress = Throwable.class)
private static void onEnter(@Advice.This ClientCall.Listener<?> listener,
                            @Advice.Local("span") Span span) {

    if (tracer == null || grpcHelperManager == null) {
        return;
    }

    GrpcHelper helper = grpcHelperManager.getForClassLoaderOfClass(ClientCall.Listener.class);
    if (helper != null) {
        span = helper.enterClientListenerMethod(listener);
    }
}
 

static void startCall(ClientCall<Message, Message> call,
    ClientCall.Listener<Message> responseListener, boolean streamingResponse) {
  call.start(responseListener, new Metadata());
  if (streamingResponse) {
    call.request(1);
  } else {
    call.request(2);
  }
}
 

@Advice.OnMethodEnter(suppress = Throwable.class)
private static void onEnter(@Advice.This ClientCall.Listener<?> listener,
                            @Advice.Local("span") Span span) {

    if (tracer == null || grpcHelperManager == null) {
        return;
    }

    GrpcHelper helper = grpcHelperManager.getForClassLoaderOfClass(ClientCall.Listener.class);
    if (helper != null) {
        span = helper.enterClientListenerMethod(listener);
    }
}
 

public StreamResponseListener(final ClientCall.Listener<RespT> responseListener) {
    super(responseListener);
}
 

@Override
public void start(ClientCall.Listener<RespT> listener, Metadata headers) {}
 

@Test
@SuppressWarnings("unchecked")
public void testInterceptors() {
    String requestDescriptor = "createStream-myScope-myStream";
    long requestId = 1234L;
    ClientInterceptor clientInterceptor = RPCTracingHelpers.getClientInterceptor();
    RequestTracker requestTracker = new RequestTracker(true);

    // Mocking RPC elements.
    MethodDescriptor.Marshaller<Object> mockMarshaller = Mockito.mock(MethodDescriptor.Marshaller.class);
    ClientCall.Listener<Object> listener = Mockito.mock(ClientCall.Listener.class);
    ServerCall serverCall = Mockito.mock(ServerCall.class);
    ServerCallHandler serverCallHandler = Mockito.mock(ServerCallHandler.class);
    ManagedChannel channel = NettyChannelBuilder.forTarget("localhost").build();
    MethodDescriptor method = MethodDescriptor.newBuilder()
                                              .setFullMethodName("createStream")
                                              .setType(MethodDescriptor.MethodType.UNARY)
                                              .setRequestMarshaller(mockMarshaller)
                                              .setResponseMarshaller(mockMarshaller)
                                              .build();
    Mockito.when(serverCall.getMethodDescriptor()).thenReturn(method);

    // Actual elements to work with.
    CallOptions callOptions = CallOptions.DEFAULT;
    Metadata headers = new Metadata();

    // Test that headers do not contain tracing-related key/values, as call options are not set.
    clientInterceptor.interceptCall(method, callOptions, channel).start(listener, headers);
    assertFalse(headers.containsKey(RPCTracingHelpers.DESCRIPTOR_HEADER));
    assertFalse(headers.containsKey(RPCTracingHelpers.ID_HEADER));

    // Check that the server interceptor handles clients not sending tracing headers and that the cache remains clean.
    ServerInterceptor serverInterceptor = RPCTracingHelpers.getServerInterceptor(requestTracker);
    serverInterceptor.interceptCall(serverCall, headers, serverCallHandler);
    assertEquals(0, requestTracker.getNumDescriptors());

    // Add call options and check that headers are correctly set.
    callOptions = callOptions.withOption(RPCTracingHelpers.REQUEST_DESCRIPTOR_CALL_OPTION, requestDescriptor)
                             .withOption(RPCTracingHelpers.REQUEST_ID_CALL_OPTION, String.valueOf(requestId));
    clientInterceptor.interceptCall(method, callOptions, channel).start(listener, headers);
    assertEquals(requestDescriptor, headers.get(RPCTracingHelpers.DESCRIPTOR_HEADER));
    assertEquals(requestId, Long.parseLong(headers.get(RPCTracingHelpers.ID_HEADER)));

    // Test that the server interceptor correctly sets these headers in the cache for further tracking.
    serverInterceptor.interceptCall(serverCall, headers, serverCallHandler);
    assertEquals(1, requestTracker.getNumDescriptors());
    assertEquals(requestId, requestTracker.getRequestIdFor(requestDescriptor));
}
 

@SuppressWarnings({"rawtypes", "unchecked"})
private static void onClientMessageHelper(ClientCall.Listener listener, Object request) {
  listener.onMessage(request);
}
 

@SuppressWarnings({"rawtypes", "unchecked"})
private static void onClientMessageHelper(ClientCall.Listener listener, Object request) {
  listener.onMessage(request);
}
 

@Test
@SuppressWarnings({"unchecked"})
public void clientDeadlineLogged_deadlineSetViaContext() throws Exception {
  // important: deadline is read from the ctx where call was created
  final SettableFuture<ClientCall<byte[], byte[]>> callFuture = SettableFuture.create();
  Context.current()
      .withDeadline(
          Deadline.after(1, TimeUnit.SECONDS), Executors.newSingleThreadScheduledExecutor())
      .run(new Runnable() {
        @Override
        public void run() {
          MethodDescriptor<byte[], byte[]> method =
              MethodDescriptor.<byte[], byte[]>newBuilder()
                  .setType(MethodType.UNKNOWN)
                  .setFullMethodName("service/method")
                  .setRequestMarshaller(BYTEARRAY_MARSHALLER)
                  .setResponseMarshaller(BYTEARRAY_MARSHALLER)
                  .build();

          callFuture.set(new BinlogHelper(mockSinkWriter)
              .getClientInterceptor(CALL_ID)
              .interceptCall(
                  method,
                  CallOptions.DEFAULT.withDeadlineAfter(1, TimeUnit.SECONDS),
                  new Channel() {
                    @Override
                    public <RequestT, ResponseT> ClientCall<RequestT, ResponseT> newCall(
                        MethodDescriptor<RequestT, ResponseT> methodDescriptor,
                        CallOptions callOptions) {
                      return new NoopClientCall<RequestT, ResponseT>();
                    }

                    @Override
                    public String authority() {
                      return null;
                    }
                  }));
        }
      });
  ClientCall.Listener<byte[]> mockListener = mock(ClientCall.Listener.class);
  callFuture.get().start(mockListener, new Metadata());
  ArgumentCaptor<Duration> callOptTimeoutCaptor = ArgumentCaptor.forClass(Duration.class);
  verify(mockSinkWriter).logClientHeader(
      any(Integer.class),
      any(String.class),
      any(String.class),
      callOptTimeoutCaptor.capture(),
      any(Metadata.class),
      any(GrpcLogEntry.Logger.class),
      any(Long.class),
      any(SocketAddress.class));
  Duration timeout = callOptTimeoutCaptor.getValue();
  assertThat(TimeUnit.SECONDS.toNanos(1) - Durations.toNanos(timeout))
      .isAtMost(TimeUnit.MILLISECONDS.toNanos(250));
}
 

@Override
public void start(ClientCall.Listener<RespT> listener, Metadata headers) {
  listener.onClose(error, new Metadata());
}
 

@Test
public void disablingInboundAutoFlowControlSuppressesRequestsForMoreMessages()
    throws Exception {
  final AtomicReference<ClientCall.Listener<String>> listener =
      new AtomicReference<ClientCall.Listener<String>>();
  final List<Integer> requests = new ArrayList<>();
  NoopClientCall<Integer, String> call = new NoopClientCall<Integer, String>() {
    @Override
    public void start(io.grpc.ClientCall.Listener<String> responseListener, Metadata headers) {
      listener.set(responseListener);
    }

    @Override
    public void request(int numMessages) {
      requests.add(numMessages);
    }
  };
  ClientCalls.asyncBidiStreamingCall(call, new ClientResponseObserver<Integer, String>() {
    @Override
    public void beforeStart(ClientCallStreamObserver<Integer> requestStream) {
      requestStream.disableAutoInboundFlowControl();
    }

    @Override
    public void onNext(String value) {

    }

    @Override
    public void onError(Throwable t) {

    }

    @Override
    public void onCompleted() {

    }
  });
  listener.get().onMessage("message");
  assertThat(requests).containsExactly(1);
}
 

@Test
public void callStreamObserverPropagatesFlowControlRequestsToCall()
    throws Exception {
  ClientResponseObserver<Integer, String> responseObserver =
      new ClientResponseObserver<Integer, String>() {
        @Override
        public void beforeStart(ClientCallStreamObserver<Integer> requestStream) {
          requestStream.disableAutoInboundFlowControl();
        }

        @Override
        public void onNext(String value) {
        }

        @Override
        public void onError(Throwable t) {
        }

        @Override
        public void onCompleted() {
        }
      };
  final AtomicReference<ClientCall.Listener<String>> listener =
      new AtomicReference<ClientCall.Listener<String>>();
  final List<Integer> requests = new ArrayList<>();
  NoopClientCall<Integer, String> call = new NoopClientCall<Integer, String>() {
    @Override
    public void start(io.grpc.ClientCall.Listener<String> responseListener, Metadata headers) {
      listener.set(responseListener);
    }

    @Override
    public void request(int numMessages) {
      requests.add(numMessages);
    }
  };
  CallStreamObserver<Integer> requestObserver =
      (CallStreamObserver<Integer>)
          ClientCalls.asyncBidiStreamingCall(call, responseObserver);
  listener.get().onMessage("message");
  requestObserver.request(5);
  assertThat(requests).contains(5);
}
 

@Override
public void start(ClientCall.Listener<Integer> responseListener, Metadata headers) {
  this.responseListener = responseListener;
  this.headers = headers;
}
 

/**
 * Starts client call and switch to client call listener instrumentation.
 * <br>
 * This is the 3rd method called during client call execution, the next in sequence is
 * {@link #clientCallStartExit(ClientCall.Listener, Throwable)}.
 *
 * @param clientCall client call
 * @param listener   client call listener
 * @param headers    headers
 * @return span, or {@literal null is there is none}
 */
@Nullable
Span clientCallStartEnter(ClientCall<?, ?> clientCall,
                          ClientCall.Listener<?> listener,
                          Metadata headers);