fix(tracing): use ApiFutures chaining to guarantee synchronous callback execution before future resolution

Author: blakeli0

sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/tracing/TracedBatchingCallable.java

@@ -29,6 +29,8 @@
  */
 package com.google.api.gax.tracing;
 
+import com.google.api.core.ApiAsyncFunction;
+import com.google.api.core.ApiFunction;
 import com.google.api.core.ApiFuture;
 import com.google.api.core.ApiFutures;
 import com.google.api.core.BetaApi;
@@ -100,9 +102,30 @@ public ApiFuture<ResponseT> futureCall(RequestT request, ApiCallContext context)
 
       context = context.withTracer(tracer);
       ApiFuture<ResponseT> future = innerCallable.futureCall(request, context);
-      ApiFutures.addCallback(future, finisher, MoreExecutors.directExecutor());
 
-      return future;
+      ApiFuture<ResponseT> onSuccessFuture =
+          ApiFutures.transform(
+              future,
+              new ApiFunction<ResponseT, ResponseT>() {
+                @Override
+                public ResponseT apply(ResponseT response) {
+                  finisher.onSuccess(response);
+                  return response;
+                }
+              },
+              MoreExecutors.directExecutor());
+
+      return ApiFutures.catchingAsync(
+          onSuccessFuture,
+          Throwable.class,
+          new ApiAsyncFunction<Throwable, ResponseT>() {
+            @Override
+            public ApiFuture<ResponseT> apply(Throwable t) {
+              finisher.onFailure(t);
+              return ApiFutures.immediateFailedFuture(t);
+            }
+          },
+          MoreExecutors.directExecutor());
     } catch (RuntimeException e) {
       finisher.onFailure(e);
       throw e;

sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/tracing/TracedUnaryCallable.java

@@ -29,6 +29,8 @@
  */
 package com.google.api.gax.tracing;
 
+import com.google.api.core.ApiAsyncFunction;
+import com.google.api.core.ApiFunction;
 import com.google.api.core.ApiFuture;
 import com.google.api.core.ApiFutures;
 import com.google.api.core.BetaApi;
@@ -100,9 +102,30 @@ public ApiFuture<ResponseT> futureCall(RequestT request, ApiCallContext context)
     try {
       context = context.withTracer(tracer);
       ApiFuture<ResponseT> future = innerCallable.futureCall(request, context);
-      ApiFutures.addCallback(future, finisher, MoreExecutors.directExecutor());
 
-      return future;
+      ApiFuture<ResponseT> onSuccessFuture =
+          ApiFutures.transform(
+              future,
+              new ApiFunction<ResponseT, ResponseT>() {
+                @Override
+                public ResponseT apply(ResponseT response) {
+                  finisher.onSuccess(response);
+                  return response;
+                }
+              },
+              MoreExecutors.directExecutor());
+
+      return ApiFutures.catchingAsync(
+          onSuccessFuture,
+          Throwable.class,
+          new ApiAsyncFunction<Throwable, ResponseT>() {
+            @Override
+            public ApiFuture<ResponseT> apply(Throwable t) {
+              finisher.onFailure(t);
+              return ApiFutures.immediateFailedFuture(t);
+            }
+          },
+          MoreExecutors.directExecutor());
     } catch (RuntimeException e) {
       finisher.onFailure(e);
       throw e;