Support setting onReadyThreshold through AbstractStub

Add copy of the onReadyThreshold property when copying CallOptions(fix bug)
grpc · Jun 27, 2024 · 25a8b7c · 25a8b7c
1 parent e7c3803
commit 25a8b7c
Show file tree

Hide file tree

Showing 4 changed files with 37 additions and 0 deletions.
diff --git a/api/src/main/java/io/grpc/CallOptions.java b/api/src/main/java/io/grpc/CallOptions.java
@@ -512,6 +512,7 @@ private static Builder toBuilder(CallOptions other) {
     builder.waitForReady = other.waitForReady;
     builder.maxInboundMessageSize = other.maxInboundMessageSize;
     builder.maxOutboundMessageSize = other.maxOutboundMessageSize;
+    builder.onReadyThreshold = other.onReadyThreshold;
     return builder;
   }
 
@@ -527,6 +528,7 @@ public String toString() {
         .add("waitForReady", isWaitForReady())
         .add("maxInboundMessageSize", maxInboundMessageSize)
         .add("maxOutboundMessageSize", maxOutboundMessageSize)
+        .add("onReadyThreshold", onReadyThreshold)
         .add("streamTracerFactories", streamTracerFactories)
         .toString();
   }

diff --git a/api/src/test/java/io/grpc/CallOptionsTest.java b/api/src/test/java/io/grpc/CallOptionsTest.java
@@ -81,6 +81,16 @@ public void withAndWithoutWaitForReady() {
         .isFalse();
   }
 
+  @Test
+  public void withOnReadyThreshold() {
+    int onReadyThreshold = 1024;
+    CallOptions callOptions = CallOptions.DEFAULT.withOnReadyThreshold(onReadyThreshold);
+    callOptions = callOptions.withWaitForReady();
+    assertThat(callOptions.getOnReadyThreshold()).isEqualTo(onReadyThreshold);
+    callOptions = callOptions.clearOnReadyThreshold();
+    assertThat(callOptions.getOnReadyThreshold()).isNull();
+  }
+
   @Test
   public void allWiths() {
     assertThat(allSet.getAuthority()).isSameInstanceAs(sampleAuthority);
@@ -148,6 +158,7 @@ public void toStringMatches_noDeadline_default() {
         .withCallCredentials(null)
         .withMaxInboundMessageSize(44)
         .withMaxOutboundMessageSize(55)
+        .withOnReadyThreshold(1024)
         .toString();
 
     assertThat(actual).contains("deadline=null");
@@ -159,6 +170,7 @@ public void toStringMatches_noDeadline_default() {
     assertThat(actual).contains("waitForReady=true");
     assertThat(actual).contains("maxInboundMessageSize=44");
     assertThat(actual).contains("maxOutboundMessageSize=55");
+    assertThat(actual).contains("onReadyThreshold=1024");
     assertThat(actual).contains("streamTracerFactories=[tracerFactory1, tracerFactory2]");
   }
 

diff --git a/stub/src/main/java/io/grpc/stub/AbstractStub.java b/stub/src/main/java/io/grpc/stub/AbstractStub.java
@@ -252,6 +252,16 @@ public final S withMaxOutboundMessageSize(int maxSize) {
     return build(channel, callOptions.withMaxOutboundMessageSize(maxSize));
   }
 
+  /**
+   * Returns a new stub that limits the maximum number of bytes per stream in the queue.
+   *
+   * @since 1.1.0
+   */
+  @ExperimentalApi("https://github.com/grpc/grpc-java/issues/11021")
+  public final S withOnReadyThreshold(int numBytes) {
+    return build(channel, callOptions.withOnReadyThreshold(numBytes));
+  }
+
   /**
    * A factory class for stub.
    *

diff --git a/stub/src/test/java/io/grpc/stub/BaseAbstractStubTest.java b/stub/src/test/java/io/grpc/stub/BaseAbstractStubTest.java
@@ -16,6 +16,7 @@
 
 package io.grpc.stub;
 
+import static com.google.common.truth.Truth.assertThat;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNull;
@@ -90,4 +91,16 @@ public void withExecutor() {
 
     assertEquals(callOptions.getExecutor(), executor);
   }
+
+  @Test
+  public void withOnReadyThreshold() {
+    T stub = create(channel);
+    CallOptions callOptions = stub.getCallOptions();
+    assertNull(callOptions.getOnReadyThreshold());
+
+    int onReadyThreshold = 1024;
+    stub = stub.withOnReadyThreshold(onReadyThreshold);
+    callOptions = stub.getCallOptions();
+    assertThat(callOptions.getOnReadyThreshold()).isEqualTo(onReadyThreshold);
+  }
 }