Support setting onReadyThreshold through AbstractStub

hlx502 · ejona86 · commit 25a8b7c50723 · 2024-06-27T09:36:49.000-07:00
Add copy of the onReadyThreshold property when copying CallOptions(fix bug)
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);
+  }
 }

Original file line number	Diff line number	Diff line change
`@@ -512,6 +512,7 @@ private static Builder toBuilder(CallOptions other) {`
`512`	`512`	`builder.waitForReady = other.waitForReady;`
`513`	`513`	`builder.maxInboundMessageSize = other.maxInboundMessageSize;`
`514`	`514`	`builder.maxOutboundMessageSize = other.maxOutboundMessageSize;`
	`515`	`+ builder.onReadyThreshold = other.onReadyThreshold;`
`515`	`516`	`return builder;`
`516`	`517`	`}`
`517`	`518`
`@@ -527,6 +528,7 @@ public String toString() {`
`527`	`528`	`.add("waitForReady", isWaitForReady())`
`528`	`529`	`.add("maxInboundMessageSize", maxInboundMessageSize)`
`529`	`530`	`.add("maxOutboundMessageSize", maxOutboundMessageSize)`
	`531`	`+ .add("onReadyThreshold", onReadyThreshold)`
`530`	`532`	`.add("streamTracerFactories", streamTracerFactories)`
`531`	`533`	`.toString();`
`532`	`534`	`}`