feat: optimize unary callables to not wait for trailers (#1356)

* feat: optimize unary callables to not wait for trailers [draft]

gRPC ClientCalls and thus gax currently wait for trailers to resolve unary call futures. I believe the original reason for this was to mitigate misconfigured servers where a server endpoint was changed to be server streaming, but the client still expects a unary method. We measured the cost of this safety net to be O(hundreds of millis). For low latency services like Bigtable, this is very high.

This PR is incomplete, but is meant to be a conversation starter. I would like to get gax's opinion on this and guidance how to proceed. Some initial proposals:
1. productionize this PR and roll it out
2. gate this behavior using a flag in UnaryCallSettings
3. expose a bit more surface in gax to allow cloud bigtable to build our callable chains (the current blocker is that GrpcUnaryRequestParamCallable & GrpcExceptionCallable are package private

* add an opt-in for skipping trailers

* oops

* address feedback

* remove separate callable

* oops

* format

gax-grpc/src/main/java/com/google/api/gax/grpc/GrpcCallSettings.java

@@ -38,12 +38,12 @@
 public class GrpcCallSettings<RequestT, ResponseT> {
   private final MethodDescriptor<RequestT, ResponseT> methodDescriptor;
   private final RequestParamsExtractor<RequestT> paramsExtractor;
+  private final boolean alwaysAwaitTrailers;
 
-  private GrpcCallSettings(
-      MethodDescriptor<RequestT, ResponseT> methodDescriptor,
-      RequestParamsExtractor<RequestT> paramsExtractor) {
-    this.methodDescriptor = methodDescriptor;
-    this.paramsExtractor = paramsExtractor;
+  private GrpcCallSettings(Builder builder) {
+    this.methodDescriptor = builder.methodDescriptor;
+    this.paramsExtractor = builder.paramsExtractor;
+    this.alwaysAwaitTrailers = builder.shouldAwaitTrailers;
   }
 
   public MethodDescriptor<RequestT, ResponseT> getMethodDescriptor() {
@@ -55,8 +55,13 @@ public RequestParamsExtractor<RequestT> getParamsExtractor() {
     return paramsExtractor;
   }
 
+  @BetaApi
+  public boolean shouldAwaitTrailers() {
+    return alwaysAwaitTrailers;
+  }
+
   public static <RequestT, ResponseT> Builder<RequestT, ResponseT> newBuilder() {
-    return new Builder<>();
+    return new Builder<RequestT, ResponseT>().setShouldAwaitTrailers(true);
   }
 
   public static <RequestT, ResponseT> GrpcCallSettings<RequestT, ResponseT> create(
@@ -73,11 +78,14 @@ public Builder toBuilder() {
   public static class Builder<RequestT, ResponseT> {
     private MethodDescriptor<RequestT, ResponseT> methodDescriptor;
     private RequestParamsExtractor<RequestT> paramsExtractor;
+    private boolean shouldAwaitTrailers;
 
     private Builder() {}
 
     private Builder(GrpcCallSettings<RequestT, ResponseT> settings) {
       this.methodDescriptor = settings.methodDescriptor;
+      this.paramsExtractor = settings.paramsExtractor;
+      this.shouldAwaitTrailers = settings.alwaysAwaitTrailers;
     }
 
     public Builder<RequestT, ResponseT> setMethodDescriptor(
@@ -93,8 +101,14 @@ public Builder<RequestT, ResponseT> setParamsExtractor(
       return this;
     }
 
+    @BetaApi
+    public Builder<RequestT, ResponseT> setShouldAwaitTrailers(boolean b) {
+      this.shouldAwaitTrailers = b;
+      return this;
+    }
+
     public GrpcCallSettings<RequestT, ResponseT> build() {
-      return new GrpcCallSettings<>(methodDescriptor, paramsExtractor);
+      return new GrpcCallSettings<>(this);
     }
   }
 }

gax-grpc/src/main/java/com/google/api/gax/grpc/GrpcClientCalls.java

@@ -30,6 +30,9 @@
 package com.google.api.gax.grpc;
 
 import com.google.api.client.util.Preconditions;
+import com.google.api.core.AbstractApiFuture;
+import com.google.api.core.ApiFuture;
+import com.google.api.core.BetaApi;
 import com.google.api.gax.rpc.ApiCallContext;
 import com.google.api.gax.tracing.ApiTracer.Scope;
 import io.grpc.CallOptions;
@@ -38,16 +41,22 @@
 import io.grpc.ClientInterceptor;
 import io.grpc.ClientInterceptors;
 import io.grpc.Deadline;
+import io.grpc.Metadata;
 import io.grpc.MethodDescriptor;
+import io.grpc.Status;
 import io.grpc.stub.MetadataUtils;
 import java.util.concurrent.TimeUnit;
+import java.util.logging.Level;
+import java.util.logging.Logger;
 
 /**
  * {@code GrpcClientCalls} creates a new {@code ClientCall} from the given call context.
  *
  * <p>Package-private for internal use.
  */
 class GrpcClientCalls {
+  private static final Logger LOGGER = Logger.getLogger(GrpcDirectCallable.class.getName());
+
   private GrpcClientCalls() {};
 
   public static <RequestT, ResponseT> ClientCall<RequestT, ResponseT> newCall(
@@ -90,4 +99,101 @@ public static <RequestT, ResponseT> ClientCall<RequestT, ResponseT> newCall(
       return channel.newCall(descriptor, callOptions);
     }
   }
+
+  /**
+   * A work-alike of {@link io.grpc.stub.ClientCalls#futureUnaryCall(ClientCall, Object)}.
+   *
+   * <p>The only difference is that unlike grpc-stub's implementation. This implementation doesn't
+   * wait for trailers to resolve a unary RPC. This can save milliseconds when the server is
+   * overloaded.
+   */
+  @BetaApi
+  static <RequestT, ResponseT> ApiFuture<ResponseT> eagerFutureUnaryCall(
+      ClientCall<RequestT, ResponseT> clientCall, RequestT request) {
+    // Start the call
+    GrpcFuture<ResponseT> future = new GrpcFuture<>(clientCall);
+    clientCall.start(new EagerFutureListener<>(future), new Metadata());
+
+    // Send the request
+    try {
+      clientCall.sendMessage(request);
+      clientCall.halfClose();
+      // Request an extra message to detect misconfigured servers
+      clientCall.request(2);
+    } catch (Throwable sendError) {
+      // Cancel if anything goes wrong
+      try {
+        clientCall.cancel(null, sendError);
+      } catch (Throwable cancelError) {
+        LOGGER.log(Level.SEVERE, "Error encountered while closing it", sendError);
+      }
+
+      throw sendError;
+    }
+
+    return future;
+  }
+
+  /** Thin wrapper around an ApiFuture that will cancel the underlying ClientCall. */
+  private static class GrpcFuture<T> extends AbstractApiFuture<T> {
+    private final ClientCall<?, T> call;
+
+    private GrpcFuture(ClientCall<?, T> call) {
+      this.call = call;
+    }
+
+    @Override
+    protected void interruptTask() {
+      call.cancel("GrpcFuture was cancelled", null);
+    }
+
+    @Override
+    public boolean set(T value) {
+      return super.set(value);
+    }
+
+    @Override
+    public boolean setException(Throwable throwable) {
+      return super.setException(throwable);
+    }
+  }
+
+  /**
+   * A bridge between gRPC's ClientCall.Listener to an ApiFuture.
+   *
+   * <p>The Listener will eagerly resolve the future when the first message is received and will not
+   * wait for the trailers. This should cut down on the latency at the expense of safety. If the
+   * server is misconfigured and sends a second response for a unary call, the error will be logged,
+   * but the future will still be successful.
+   */
+  private static class EagerFutureListener<T> extends ClientCall.Listener<T> {
+    private final GrpcFuture<T> future;
+
+    private EagerFutureListener(GrpcFuture<T> future) {
+      this.future = future;
+    }
+
+    @Override
+    public void onMessage(T message) {
+      if (!future.set(message)) {
+        throw Status.INTERNAL
+            .withDescription("More than one value received for unary call")
+            .asRuntimeException();
+      }
+    }
+
+    @Override
+    public void onClose(Status status, Metadata trailers) {
+      if (!future.isDone()) {
+        future.setException(
+            Status.INTERNAL
+                .withDescription("No value received for unary call")
+                .asException(trailers));
+      }
+      if (!status.isOk()) {
+        LOGGER.log(
+            Level.WARNING, "Received error for unary call after receiving a successful response");
+      }
+    }
+  }
 }

gax-grpc/src/main/java/com/google/api/gax/grpc/GrpcDirectCallable.java

@@ -34,6 +34,7 @@
 import com.google.api.gax.rpc.ApiCallContext;
 import com.google.api.gax.rpc.UnaryCallable;
 import com.google.common.base.Preconditions;
+import io.grpc.ClientCall;
 import io.grpc.MethodDescriptor;
 import io.grpc.stub.ClientCalls;
 
@@ -44,18 +45,25 @@
  */
 class GrpcDirectCallable<RequestT, ResponseT> extends UnaryCallable<RequestT, ResponseT> {
   private final MethodDescriptor<RequestT, ResponseT> descriptor;
+  private final boolean awaitTrailers;
 
-  GrpcDirectCallable(MethodDescriptor<RequestT, ResponseT> descriptor) {
+  GrpcDirectCallable(MethodDescriptor<RequestT, ResponseT> descriptor, boolean awaitTrailers) {
     this.descriptor = Preconditions.checkNotNull(descriptor);
+    this.awaitTrailers = awaitTrailers;
   }
 
   @Override
   public ApiFuture<ResponseT> futureCall(RequestT request, ApiCallContext inputContext) {
     Preconditions.checkNotNull(request);
     Preconditions.checkNotNull(inputContext);
 
-    return new ListenableFutureToApiFuture<>(
-        ClientCalls.futureUnaryCall(GrpcClientCalls.newCall(descriptor, inputContext), request));
+    ClientCall<RequestT, ResponseT> clientCall = GrpcClientCalls.newCall(descriptor, inputContext);
+
+    if (awaitTrailers) {
+      return new ListenableFutureToApiFuture<>(ClientCalls.futureUnaryCall(clientCall, request));
+    } else {
+      return GrpcClientCalls.eagerFutureUnaryCall(clientCall, request);
+    }
   }
 
   @Override

gax-grpc/src/main/java/com/google/api/gax/grpc/GrpcRawCallableFactory.java

@@ -52,7 +52,9 @@ private GrpcRawCallableFactory() {}
   public static <RequestT, ResponseT> UnaryCallable<RequestT, ResponseT> createUnaryCallable(
       GrpcCallSettings<RequestT, ResponseT> grpcCallSettings, Set<StatusCode.Code> retryableCodes) {
     UnaryCallable<RequestT, ResponseT> callable =
-        new GrpcDirectCallable<>(grpcCallSettings.getMethodDescriptor());
+        new GrpcDirectCallable<>(
+            grpcCallSettings.getMethodDescriptor(), grpcCallSettings.shouldAwaitTrailers());
+
     if (grpcCallSettings.getParamsExtractor() != null) {
       callable =
           new GrpcUnaryRequestParamCallable<>(callable, grpcCallSettings.getParamsExtractor());