feat: pass request in ApiTracer (#1491)

* feat: pass request or response in ApiTracer

* add comments

* make change backward compatible

* revert attemptSucceeded changes

Co-authored-by: Igor Bernstein <igorbernstein@google.com>

Author: mutianf

gax/src/main/java/com/google/api/gax/rpc/AttemptCallable.java

@@ -82,7 +82,7 @@ public ResponseT call() {
 
       callContext
           .getTracer()
-          .attemptStarted(externalFuture.getAttemptSettings().getOverallAttemptCount());
+          .attemptStarted(request, externalFuture.getAttemptSettings().getOverallAttemptCount());
 
       ApiFuture<ResponseT> internalFuture = callable.futureCall(request, callContext);
       externalFuture.setAttemptFuture(internalFuture);

gax/src/main/java/com/google/api/gax/rpc/ServerStreamingAttemptCallable.java

@@ -229,7 +229,7 @@ public Void call() {
 
     attemptContext
         .getTracer()
-        .attemptStarted(outerRetryingFuture.getAttemptSettings().getOverallAttemptCount());
+        .attemptStarted(request, outerRetryingFuture.getAttemptSettings().getOverallAttemptCount());
 
     innerCallable.call(
         request,

gax/src/main/java/com/google/api/gax/tracing/ApiTracer.java

@@ -83,9 +83,21 @@ public interface ApiTracer {
    * start of the operation. The attemptNumber is zero based. So the initial attempt will be 0.
    *
    * @param attemptNumber the zero based sequential attempt number.
+   * @deprecated Please use {@link #attemptStarted(Object, int)} instead.
    */
+  @Deprecated
   void attemptStarted(int attemptNumber);
 
+  /**
+   * Adds an annotation that an attempt is about to start with additional information from the
+   * request. In general this should occur at the very start of the operation. The attemptNumber is
+   * zero based. So the initial attempt will be 0.
+   *
+   * @param attemptNumber the zero based sequential attempt number.
+   * @param request request of this attempt.
+   */
+  void attemptStarted(Object request, int attemptNumber);
+
   /** Adds an annotation that the attempt succeeded. */
   void attemptSucceeded();
 

gax/src/main/java/com/google/api/gax/tracing/BaseApiTracer.java

@@ -85,6 +85,11 @@ public void attemptStarted(int attemptNumber) {
     // noop
   }
 
+  @Override
+  public void attemptStarted(Object request, int attemptNumber) {
+    attemptStarted(attemptNumber);
+  }
+
   @Override
   public void attemptSucceeded() {
     // noop

gax/src/main/java/com/google/api/gax/tracing/OpencensusTracer.java

@@ -307,6 +307,12 @@ public void attemptStarted(int attemptNumber) {
     // This simply is used for state management.
   }
 
+  /** {@inheritDoc} */
+  @Override
+  public void attemptStarted(Object request, int attemptNumber) {
+    attemptStarted(attemptNumber);
+  }
+
   /** {@inheritDoc} */
   @Override
   public void attemptSucceeded() {

gax/src/test/java/com/google/api/gax/retrying/AbstractRetryingExecutorTest.java

@@ -37,6 +37,7 @@
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertSame;
 import static org.junit.Assert.assertTrue;
+import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.any;
 import static org.mockito.Mockito.anyInt;
 import static org.mockito.Mockito.times;
@@ -115,7 +116,7 @@ private RetrySettings getDefaultRetrySettings() {
 
   @Test
   public void testSuccess() throws Exception {
-    FailingCallable callable = new FailingCallable(0, "SUCCESS", tracer);
+    FailingCallable callable = new FailingCallable(0, "request", "SUCCESS", tracer);
     RetryingExecutorWithContext<String> executor =
         getExecutor(getAlgorithm(getDefaultRetrySettings(), 0, null));
     RetryingFuture<String> future = executor.createFuture(callable, retryingContext);
@@ -124,14 +125,14 @@ public void testSuccess() throws Exception {
     assertFutureSuccess(future);
     assertEquals(0, future.getAttemptSettings().getAttemptCount());
 
-    verify(tracer, times(1)).attemptStarted(0);
+    verify(tracer, times(1)).attemptStarted(eq("request"), eq(0));
     verify(tracer, times(1)).attemptSucceeded();
     verifyNoMoreInteractions(tracer);
   }
 
   @Test
   public void testSuccessWithFailures() throws Exception {
-    FailingCallable callable = new FailingCallable(5, "SUCCESS", tracer);
+    FailingCallable callable = new FailingCallable(5, "request", "SUCCESS", tracer);
     RetryingExecutorWithContext<String> executor =
         getExecutor(getAlgorithm(getDefaultRetrySettings(), 0, null));
     RetryingFuture<String> future = executor.createFuture(callable, retryingContext);
@@ -140,15 +141,15 @@ public void testSuccessWithFailures() throws Exception {
     assertFutureSuccess(future);
     assertEquals(5, future.getAttemptSettings().getAttemptCount());
 
-    verify(tracer, times(6)).attemptStarted(anyInt());
+    verify(tracer, times(6)).attemptStarted(eq("request"), anyInt());
     verify(tracer, times(5)).attemptFailed(any(Throwable.class), any(Duration.class));
     verify(tracer, times(1)).attemptSucceeded();
     verifyNoMoreInteractions(tracer);
   }
 
   @Test
   public void testSuccessWithFailuresPeekGetAttempt() throws Exception {
-    FailingCallable callable = new FailingCallable(5, "SUCCESS", tracer);
+    FailingCallable callable = new FailingCallable(5, "request", "SUCCESS", tracer);
     RetryingExecutorWithContext<String> executor =
         getExecutor(getAlgorithm(getDefaultRetrySettings(), 0, null));
     RetryingFuture<String> future = executor.createFuture(callable, retryingContext);
@@ -174,7 +175,7 @@ public void testSuccessWithFailuresPeekGetAttempt() throws Exception {
 
   @Test
   public void testMaxRetriesExceeded() throws Exception {
-    FailingCallable callable = new FailingCallable(6, "FAILURE", tracer);
+    FailingCallable callable = new FailingCallable(6, "request", "FAILURE", tracer);
     RetryingExecutorWithContext<String> executor =
         getExecutor(getAlgorithm(getDefaultRetrySettings(), 0, null));
     RetryingFuture<String> future = executor.createFuture(callable, retryingContext);
@@ -183,7 +184,7 @@ public void testMaxRetriesExceeded() throws Exception {
     assertFutureFail(future, CustomException.class);
     assertEquals(5, future.getAttemptSettings().getAttemptCount());
 
-    verify(tracer, times(6)).attemptStarted(anyInt());
+    verify(tracer, times(6)).attemptStarted(eq("request"), anyInt());
     verify(tracer, times(5)).attemptFailed(any(Throwable.class), any(Duration.class));
     verify(tracer, times(1)).attemptFailedRetriesExhausted(any(Throwable.class));
     verifyNoMoreInteractions(tracer);
@@ -202,7 +203,7 @@ public void testTotalTimeoutExceeded() throws Exception {
         getExecutor(
             getAlgorithm(
                 useContextRetrySettings ? getDefaultRetrySettings() : retrySettings, 0, null));
-    FailingCallable callable = new FailingCallable(6, "FAILURE", tracer);
+    FailingCallable callable = new FailingCallable(6, "request", "FAILURE", tracer);
     RetryingContext context;
     if (useContextRetrySettings) {
       context = FakeCallContext.createDefault().withTracer(tracer).withRetrySettings(retrySettings);
@@ -215,14 +216,14 @@ public void testTotalTimeoutExceeded() throws Exception {
     assertFutureFail(future, CustomException.class);
     assertTrue(future.getAttemptSettings().getAttemptCount() < 4);
 
-    verify(tracer, times(1)).attemptStarted(anyInt());
+    verify(tracer, times(1)).attemptStarted(eq("request"), anyInt());
     verify(tracer, times(1)).attemptFailedRetriesExhausted(any(Throwable.class));
     verifyNoMoreInteractions(tracer);
   }
 
   @Test
   public void testCancelOuterFutureBeforeStart() throws Exception {
-    FailingCallable callable = new FailingCallable(4, "SUCCESS", tracer);
+    FailingCallable callable = new FailingCallable(4, "request", "SUCCESS", tracer);
 
     RetrySettings retrySettings =
         FAST_RETRY_SETTINGS
@@ -248,7 +249,7 @@ public void testCancelOuterFutureBeforeStart() throws Exception {
 
   @Test
   public void testCancelByRetryingAlgorithm() throws Exception {
-    FailingCallable callable = new FailingCallable(6, "FAILURE", tracer);
+    FailingCallable callable = new FailingCallable(6, "request", "FAILURE", tracer);
     RetryingExecutorWithContext<String> executor =
         getExecutor(getAlgorithm(getDefaultRetrySettings(), 5, new CancellationException()));
     RetryingFuture<String> future = executor.createFuture(callable, retryingContext);
@@ -257,7 +258,7 @@ public void testCancelByRetryingAlgorithm() throws Exception {
     assertFutureCancel(future);
     assertEquals(4, future.getAttemptSettings().getAttemptCount());
 
-    verify(tracer, times(5)).attemptStarted(anyInt());
+    verify(tracer, times(5)).attemptStarted(eq("request"), anyInt());
     // Pre-apocalypse failures
     verify(tracer, times(4)).attemptFailed(any(Throwable.class), any(Duration.class));
     // Apocalypse failure
@@ -267,7 +268,7 @@ public void testCancelByRetryingAlgorithm() throws Exception {
 
   @Test
   public void testUnexpectedExceptionFromRetryAlgorithm() throws Exception {
-    FailingCallable callable = new FailingCallable(6, "FAILURE", tracer);
+    FailingCallable callable = new FailingCallable(6, "request", "FAILURE", tracer);
     RetryingExecutorWithContext<String> executor =
         getExecutor(getAlgorithm(getDefaultRetrySettings(), 5, new RuntimeException()));
     RetryingFuture<String> future = executor.createFuture(callable, retryingContext);
@@ -276,7 +277,7 @@ public void testUnexpectedExceptionFromRetryAlgorithm() throws Exception {
     assertFutureFail(future, RuntimeException.class);
     assertEquals(4, future.getAttemptSettings().getAttemptCount());
 
-    verify(tracer, times(5)).attemptStarted(anyInt());
+    verify(tracer, times(5)).attemptStarted(eq("request"), anyInt());
     // Pre-apocalypse failures
     verify(tracer, times(4)).attemptFailed(any(Throwable.class), any(Duration.class));
     // Apocalypse failure
@@ -302,7 +303,7 @@ public void testPollExceptionByPollAlgorithm() throws Exception {
                 NanoClock.getDefaultClock()));
 
     RetryingExecutorWithContext<String> executor = getExecutor(retryAlgorithm);
-    FailingCallable callable = new FailingCallable(6, "FAILURE", tracer);
+    FailingCallable callable = new FailingCallable(6, "request", "FAILURE", tracer);
     RetryingContext context;
     if (useContextRetrySettings) {
       context = FakeCallContext.createDefault().withTracer(tracer).withRetrySettings(retrySettings);
@@ -315,7 +316,7 @@ public void testPollExceptionByPollAlgorithm() throws Exception {
     assertFutureFail(future, PollException.class);
     assertTrue(future.getAttemptSettings().getAttemptCount() < 4);
 
-    verify(tracer, times(1)).attemptStarted(anyInt());
+    verify(tracer, times(1)).attemptStarted(eq("request"), anyInt());
     verify(tracer, times(1)).attemptPermanentFailure(any(PollException.class));
     verifyNoMoreInteractions(tracer);
   }

gax/src/test/java/com/google/api/gax/retrying/FailingCallable.java

@@ -62,10 +62,12 @@ class FailingCallable implements Callable<String> {
   private AtomicInteger attemptsCount = new AtomicInteger(0);
   private final ApiTracer tracer;
   private final int expectedFailuresCount;
+  private final String request;
   private final String result;
   private final CountDownLatch firstAttemptFinished = new CountDownLatch(1);
 
-  FailingCallable(int expectedFailuresCount, String result, ApiTracer tracer) {
+  FailingCallable(int expectedFailuresCount, String request, String result, ApiTracer tracer) {
+    this.request = request;
     this.tracer = tracer;
     this.expectedFailuresCount = expectedFailuresCount;
     this.result = result;
@@ -80,7 +82,7 @@ public String call() throws Exception {
     try {
       int attemptNumber = attemptsCount.getAndIncrement();
 
-      tracer.attemptStarted(attemptNumber);
+      tracer.attemptStarted(request, attemptNumber);
 
       if (attemptNumber < expectedFailuresCount) {
         throw new CustomException();

gax/src/test/java/com/google/api/gax/retrying/ScheduledRetryingExecutorTest.java

@@ -88,7 +88,7 @@ public void testSuccessWithFailuresPeekAttempt() throws Exception {
       final int maxRetries = 100;
 
       ScheduledExecutorService localExecutor = Executors.newSingleThreadScheduledExecutor();
-      FailingCallable callable = new FailingCallable(15, "SUCCESS", tracer);
+      FailingCallable callable = new FailingCallable(15, "request", "SUCCESS", tracer);
 
       RetrySettings retrySettings =
           FAST_RETRY_SETTINGS
@@ -139,7 +139,7 @@ public void testSuccessWithFailuresGetAttempt() throws Exception {
       final int maxRetries = 100;
 
       ScheduledExecutorService localExecutor = Executors.newSingleThreadScheduledExecutor();
-      FailingCallable callable = new FailingCallable(15, "SUCCESS", tracer);
+      FailingCallable callable = new FailingCallable(15, "request", "SUCCESS", tracer);
       RetrySettings retrySettings =
           FAST_RETRY_SETTINGS
               .toBuilder()
@@ -192,7 +192,7 @@ public void testCancelGetAttempt() throws Exception {
       ScheduledExecutorService localExecutor = Executors.newSingleThreadScheduledExecutor();
       final int maxRetries = 100;
 
-      FailingCallable callable = new FailingCallable(maxRetries - 1, "SUCCESS", tracer);
+      FailingCallable callable = new FailingCallable(maxRetries - 1, "request", "SUCCESS", tracer);
       RetrySettings retrySettings =
           FAST_RETRY_SETTINGS
               .toBuilder()
@@ -249,7 +249,7 @@ public void testCancelGetAttempt() throws Exception {
   public void testCancelOuterFutureAfterStart() throws Exception {
     for (int executionsCount = 0; executionsCount < EXECUTIONS_COUNT; executionsCount++) {
       ScheduledExecutorService localExecutor = Executors.newSingleThreadScheduledExecutor();
-      FailingCallable callable = new FailingCallable(4, "SUCCESS", tracer);
+      FailingCallable callable = new FailingCallable(4, "requset", "SUCCESS", tracer);
       RetrySettings retrySettings =
           FAST_RETRY_SETTINGS
               .toBuilder()
@@ -276,7 +276,7 @@ public void testCancelOuterFutureAfterStart() throws Exception {
   @Test
   public void testCancelIsTraced() throws Exception {
     ScheduledExecutorService localExecutor = Executors.newSingleThreadScheduledExecutor();
-    FailingCallable callable = new FailingCallable(4, "SUCCESS", tracer);
+    FailingCallable callable = new FailingCallable(4, "request", "SUCCESS", tracer);
     RetrySettings retrySettings =
         FAST_RETRY_SETTINGS
             .toBuilder()
@@ -305,7 +305,7 @@ public void testCancelProxiedFutureAfterStart() throws Exception {
     // this is a heavy test, which takes a lot of time, so only few executions.
     for (int executionsCount = 0; executionsCount < 2; executionsCount++) {
       ScheduledExecutorService localExecutor = Executors.newSingleThreadScheduledExecutor();
-      FailingCallable callable = new FailingCallable(5, "SUCCESS", tracer);
+      FailingCallable callable = new FailingCallable(5, "request", "SUCCESS", tracer);
       RetrySettings retrySettings =
           FAST_RETRY_SETTINGS
               .toBuilder()

gax/src/test/java/com/google/api/gax/tracing/TracedCallableTest.java

@@ -111,7 +111,7 @@ public void testNonRetriedCallable() throws Exception {
     ApiFuture<String> future = callable.futureCall("Is your refrigerator running?", callContext);
 
     verify(tracerFactory, times(1)).newTracer(parentTracer, SPAN_NAME, OperationType.Unary);
-    verify(tracer, times(1)).attemptStarted(anyInt());
+    verify(tracer, times(1)).attemptStarted(anyString(), anyInt());
     verify(tracer, times(1)).attemptSucceeded();
     verify(tracer, times(1)).operationSucceeded();
     verifyNoMoreInteractions(tracer);