feat: optimize unary callables to not wait for trailers

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,6 +55,11 @@ public RequestParamsExtractor<RequestT> getParamsExtractor() {
     return paramsExtractor;
   }
 
+  @BetaApi
+  public boolean shouldAwaitTrailers() {
+    return alwaysAwaitTrailers;
+  }
+
   public static <RequestT, ResponseT> Builder<RequestT, ResponseT> newBuilder() {
     return new Builder<>();
   }
@@ -73,11 +78,16 @@ 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() {
+      shouldAwaitTrailers = true;
+    }
 
     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 +103,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 and
+      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.SEVERE, "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;
 
@@ -45,15 +46,16 @@
 class GrpcDirectCallable<RequestT, ResponseT> extends UnaryCallable<RequestT, ResponseT> {
   private final MethodDescriptor<RequestT, ResponseT> descriptor;
 
-  GrpcDirectCallable(MethodDescriptor<RequestT, ResponseT> descriptor) {
+  public GrpcDirectCallable(MethodDescriptor<RequestT, ResponseT> descriptor) {
+
     this.descriptor = Preconditions.checkNotNull(descriptor);
   }
 
   @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));
   }

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

@@ -0,0 +1,71 @@
+/*
+ * Copyright 2021 Google LLC
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are
+ * met:
+ *
+ *     * Redistributions of source code must retain the above copyright
+ * notice, this list of conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above
+ * copyright notice, this list of conditions and the following disclaimer
+ * in the documentation and/or other materials provided with the
+ * distribution.
+ *     * Neither the name of Google LLC nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+ * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+ * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+ * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+ * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+ * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+ * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+ * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+ * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+package com.google.api.gax.grpc;
+
+import com.google.api.core.ApiFuture;
+import com.google.api.core.ListenableFutureToApiFuture;
+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;
+
+/**
+ * {@code GrpcEagerDirectCallable} creates gRPC calls.
+ *
+ * <p>Unlike {@link GrpcDirectCallable}, this variant won't wait for trailers before resolving the
+ * future.
+ *
+ * <p>Package-private for internal use.
+ */
+class GrpcEagerDirectCallable<RequestT, ResponseT> extends UnaryCallable<RequestT, ResponseT> {
+  private final MethodDescriptor<RequestT, ResponseT> descriptor;
+
+  public GrpcEagerDirectCallable(
+      MethodDescriptor<RequestT, ResponseT> descriptor) {
+
+    this.descriptor = Preconditions.checkNotNull(descriptor);
+  }
+
+  @Override
+  public ApiFuture<ResponseT> futureCall(RequestT request, ApiCallContext inputContext) {
+    Preconditions.checkNotNull(request);
+    Preconditions.checkNotNull(inputContext);
+
+    ClientCall<RequestT, ResponseT> clientCall = GrpcClientCalls.newCall(descriptor, inputContext);
+    return GrpcClientCalls.eagerFutureUnaryCall(clientCall, request);
+  }
+
+  @Override
+  public String toString() {
+    return String.format("direct(%s)", descriptor);
+  }
+}

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

@@ -51,8 +51,12 @@ 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());
+    UnaryCallable<RequestT, ResponseT> callable;
+    if (grpcCallSettings.shouldAwaitTrailers()) {
+      callable = new GrpcDirectCallable<>(grpcCallSettings.getMethodDescriptor());
+    } else {
+      callable = new GrpcEagerDirectCallable<>(grpcCallSettings.getMethodDescriptor());
+    }
     if (grpcCallSettings.getParamsExtractor() != null) {
       callable =
           new GrpcUnaryRequestParamCallable<>(callable, grpcCallSettings.getParamsExtractor());