Allow configuration of the queued byte threshold at which a Stream is considered not ready (#10977)

* Allow the queued byte threshold for a Stream to be ready to be configurable

- on clients this is exposed by setting a CallOption
- on servers this is configured by calling a method on ServerCall or ServerStreamListener

Author: jduo

api/src/main/java/io/grpc/CallOptions.java

@@ -79,6 +79,8 @@ public final class CallOptions {
   private final Integer maxInboundMessageSize;
   @Nullable
   private final Integer maxOutboundMessageSize;
+  @Nullable
+  private final Integer onReadyThreshold;
 
   private CallOptions(Builder builder) {
     this.deadline = builder.deadline;
@@ -91,6 +93,7 @@ private CallOptions(Builder builder) {
     this.waitForReady = builder.waitForReady;
     this.maxInboundMessageSize = builder.maxInboundMessageSize;
     this.maxOutboundMessageSize = builder.maxOutboundMessageSize;
+    this.onReadyThreshold = builder.onReadyThreshold;
   }
 
   static class Builder {
@@ -105,6 +108,7 @@ static class Builder {
     Boolean waitForReady;
     Integer maxInboundMessageSize;
     Integer maxOutboundMessageSize;
+    Integer onReadyThreshold;
 
     private CallOptions build() {
       return new CallOptions(this);
@@ -203,6 +207,46 @@ public CallOptions withoutWaitForReady() {
     return builder.build();
   }
 
+  /**
+   * Specifies how many bytes must be queued before the call is
+   * considered not ready to send more messages.
+   *
+   * @param numBytes The number of bytes that must be queued. Must be a
+   *                 positive integer.
+   */
+  @ExperimentalApi("https://github.com/grpc/grpc-java/issues/11021")
+  public CallOptions withOnReadyThreshold(int numBytes) {
+    checkArgument(numBytes > 0, "numBytes must be positive: %s", numBytes);
+    Builder builder = toBuilder(this);
+    builder.onReadyThreshold = numBytes;
+    return builder.build();
+  }
+
+  /**
+   * Resets to the default number of bytes that must be queued before the
+   * call will leave the <a href="https://github.com/grpc/grpc/blob/master/doc/wait-for-ready.md">
+   * 'wait for ready'</a> state.
+   */
+  @ExperimentalApi("https://github.com/grpc/grpc-java/issues/11021")
+  public CallOptions clearOnReadyThreshold() {
+    Builder builder = toBuilder(this);
+    builder.onReadyThreshold = null;
+    return builder.build();
+  }
+
+  /**
+   * Returns to the default number of bytes that must be queued before the
+   * call will leave the <a href="https://github.com/grpc/grpc/blob/master/doc/wait-for-ready.md">
+   * 'wait for ready'</a> state.
+   *
+   * @return null if the default threshold is used.
+   */
+  @Nullable
+  @ExperimentalApi("https://github.com/grpc/grpc-java/issues/11021")
+  public Integer getOnReadyThreshold() {
+    return onReadyThreshold;
+  }
+
   /**
    * Returns the compressor's name.
    */

api/src/main/java/io/grpc/PartialForwardingServerCall.java

@@ -58,6 +58,12 @@ public void setMessageCompression(boolean enabled) {
     delegate().setMessageCompression(enabled);
   }
 
+  @Override
+  @ExperimentalApi("https://github.com/grpc/grpc-java/issues/11021")
+  public void setOnReadyThreshold(int numBytes) {
+    delegate().setOnReadyThreshold(numBytes);
+  }
+
   @Override
   @ExperimentalApi("https://github.com/grpc/grpc-java/issues/1704")
   public void setCompression(String compressor) {

api/src/main/java/io/grpc/ServerCall.java

@@ -16,6 +16,8 @@
 
 package io.grpc;
 
+import static com.google.common.base.Preconditions.checkArgument;
+
 import javax.annotation.Nullable;
 
 /**
@@ -209,6 +211,19 @@ public void setCompression(String compressor) {
     // noop
   }
 
+  /**
+   * A hint to the call that specifies how many bytes must be queued before
+   * {@link #isReady()} will return false. A call may ignore this property if
+   * unsupported. This may only be set before any messages are sent.
+   *
+   * @param numBytes The number of bytes that must be queued. Must be a
+   *                 positive integer.
+   */
+  @ExperimentalApi("https://github.com/grpc/grpc-java/issues/11021")
+  public void setOnReadyThreshold(int numBytes) {
+    checkArgument(numBytes > 0, "numBytes must be positive: %s", numBytes);
+  }
+
   /**
    * Returns the level of security guarantee in communications
    *

binder/src/main/java/io/grpc/binder/internal/MultiMessageServerStream.java

@@ -64,6 +64,11 @@ public void setListener(ServerStreamListener listener) {
     }
   }
 
+  @Override
+  public void setOnReadyThreshold(int numBytes) {
+    // No-op
+  }
+
   @Override
   public boolean isReady() {
     return outbound.isReady();

binder/src/main/java/io/grpc/binder/internal/SingleMessageServerStream.java

@@ -67,6 +67,11 @@ public void setListener(ServerStreamListener listener) {
     }
   }
 
+  @Override
+  public void setOnReadyThreshold(int numBytes) {
+    // No-op
+  }
+
   @Override
   public boolean isReady() {
     return outbound.isReady();

core/src/main/java/io/grpc/internal/AbstractClientStream.java

@@ -243,9 +243,13 @@ protected abstract static class TransportState extends AbstractStream.TransportS
     protected TransportState(
         int maxMessageSize,
         StatsTraceContext statsTraceCtx,
-        TransportTracer transportTracer) {
+        TransportTracer transportTracer,
+        CallOptions options) {
       super(maxMessageSize, statsTraceCtx, transportTracer);
       this.statsTraceCtx = checkNotNull(statsTraceCtx, "statsTraceCtx");
+      if (options.getOnReadyThreshold() != null) {
+        this.setOnReadyThreshold(options.getOnReadyThreshold());
+      }
     }
 
     private void setFullStreamDecompression(boolean fullStreamDecompression) {

core/src/main/java/io/grpc/internal/AbstractServerStream.java

@@ -177,6 +177,19 @@ public StatsTraceContext statsTraceContext() {
     return statsTraceCtx;
   }
 
+  /**
+   * A hint to the stream that specifies how many bytes must be queued before
+   * {@link #isReady()} will return false. A stream may ignore this property
+   * if unsupported. This may only be set before any messages are sent.
+   *
+   * @param numBytes The number of bytes that must be queued. Must be a
+   *                 positive integer.
+   */
+  @Override
+  public void setOnReadyThreshold(int numBytes) {
+    super.setOnReadyThreshold(numBytes);
+  }
+
   /**
    * This should only be called from the transport thread (except for private interactions with
    * {@code AbstractServerStream}).
@@ -243,6 +256,8 @@ public void deframerClosed(boolean hasPartialMessage) {
       }
     }
 
+
+
     @Override
     protected ServerStreamListener listener() {
       return listener;

core/src/main/java/io/grpc/internal/AbstractStream.java

@@ -77,6 +77,19 @@ public final void flush() {
     }
   }
 
+  /**
+   * A hint to the stream that specifies how many bytes must be queued before
+   * {@link #isReady()} will return false. A stream may ignore this property if
+   * unsupported. This may only be set during stream initialization before
+   * any messages are set.
+   *
+   * @param numBytes The number of bytes that must be queued. Must be a
+   *                 positive integer.
+   */
+  protected void setOnReadyThreshold(int numBytes) {
+    transportState().setOnReadyThreshold(numBytes);
+  }
+
   /**
    * Closes the underlying framer. Should be called when the outgoing stream is gracefully closed
    * (half closure on client; closure on server).
@@ -143,6 +156,9 @@ public abstract static class TransportState
     @GuardedBy("onReadyLock")
     private boolean deallocated;
 
+    @GuardedBy("onReadyLock")
+    private int onReadyThreshold;
+
     protected TransportState(
         int maxMessageSize,
         StatsTraceContext statsTraceCtx,
@@ -157,6 +173,7 @@ protected TransportState(
           transportTracer);
       // TODO(#7168): use MigratingThreadDeframer when enabling retry doesn't break.
       deframer = rawDeframer;
+      onReadyThreshold = DEFAULT_ONREADY_THRESHOLD;
     }
 
     final void optimizeForDirectExecutor() {
@@ -178,6 +195,20 @@ final void setMaxInboundMessageSize(int maxSize) {
      */
     protected abstract StreamListener listener();
 
+    /**
+     * A hint to the stream that specifies how many bytes must be queued before
+     * {@link #isReady()} will return false. A stream may ignore this property if
+     * unsupported. This may only be set before any messages are sent.
+     *
+     * @param numBytes The number of bytes that must be queued. Must be a
+     *                 positive integer.
+     */
+    void setOnReadyThreshold(int numBytes) {
+      synchronized (onReadyLock) {
+        this.onReadyThreshold = numBytes;
+      }
+    }
+
     @Override
     public void messagesAvailable(StreamListener.MessageProducer producer) {
       listener().messagesAvailable(producer);
@@ -259,7 +290,7 @@ protected final void setDecompressor(Decompressor decompressor) {
 
     private boolean isReady() {
       synchronized (onReadyLock) {
-        return allocated && numSentBytesQueued < DEFAULT_ONREADY_THRESHOLD && !deallocated;
+        return allocated && numSentBytesQueued < onReadyThreshold && !deallocated;
       }
     }
 
@@ -316,9 +347,9 @@ public final void onSentBytes(int numBytes) {
       synchronized (onReadyLock) {
         checkState(allocated,
             "onStreamAllocated was not called, but it seems the stream is active");
-        boolean belowThresholdBefore = numSentBytesQueued < DEFAULT_ONREADY_THRESHOLD;
+        boolean belowThresholdBefore = numSentBytesQueued < onReadyThreshold;
         numSentBytesQueued -= numBytes;
-        boolean belowThresholdAfter = numSentBytesQueued < DEFAULT_ONREADY_THRESHOLD;
+        boolean belowThresholdAfter = numSentBytesQueued < onReadyThreshold;
         doNotify = !belowThresholdBefore && belowThresholdAfter;
       }
       if (doNotify) {

core/src/main/java/io/grpc/internal/Http2ClientStreamTransportState.java

@@ -18,6 +18,7 @@
 
 import com.google.common.base.Charsets;
 import com.google.common.base.Preconditions;
+import io.grpc.CallOptions;
 import io.grpc.InternalMetadata;
 import io.grpc.InternalStatus;
 import io.grpc.Metadata;
@@ -67,8 +68,9 @@ public Integer parseAsciiString(byte[] serialized) {
   protected Http2ClientStreamTransportState(
       int maxMessageSize,
       StatsTraceContext statsTraceCtx,
-      TransportTracer transportTracer) {
-    super(maxMessageSize, statsTraceCtx, transportTracer);
+      TransportTracer transportTracer,
+      CallOptions options) {
+    super(maxMessageSize, statsTraceCtx, transportTracer, options);
   }
 
   /**

core/src/main/java/io/grpc/internal/ServerCallImpl.java

@@ -184,6 +184,11 @@ public void setMessageCompression(boolean enable) {
     stream.setMessageCompression(enable);
   }
 
+  @Override
+  public void setOnReadyThreshold(int numBytes) {
+    stream.setOnReadyThreshold(numBytes);
+  }
+
   @Override
   public void setCompression(String compressorName) {
     // Added here to give a better error message.