Modify scale factor and minimum refillRatio value. Add test annotation to publisher

Author: Tony Cui

java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/HedgeSettings.java

@@ -30,6 +30,12 @@ public final class HedgeSettings {
   /** Default refill rate (tokens per successful request). */
   private static final float DEFAULT_REFILL_RATIO = 0.1f;
 
+  /** Minimum refill rate. */
+  private static final float MIN_REFILL_RATIO = 0.001f;
+
+  /** Maximum refill rate. */
+  private static final float MAX_REFILL_RATIO = 0.2f;
+
   /** Hedging delay. */
   private final Duration hedgeDelay;
 
@@ -118,14 +124,17 @@ public Builder setMaxTokens(final int maxTokens) {
     /**
      * Allows the token bucket refill rate to be configurable.
      *
-     * @param refill the refill rate (tokens per successful request), must be 0 < RefillRatio <=
-     *     0.2.
+     * @param refillRatio the refill rate (tokens per successful request), must be 0.001 <=
+     *     RefillRatio <= 0.2.
      * @return this builder.
      */
     public Builder setRefillRatio(final float refillRatio) {
-      if (refillRatio <= 0.0f || refillRatio > 0.2f) {
+      if (refillRatio < MIN_REFILL_RATIO || refillRatio > MAX_REFILL_RATIO) {
         throw new IllegalArgumentException(
-            "refillRatio must be greater than 0.0 and less than or equal to 0.2");
+            "refillRatio must be greater than or equal to "
+                + MIN_REFILL_RATIO
+                + " and less than or equal to "
+                + MAX_REFILL_RATIO);
       }
       this.refillRatio = refillRatio;
       return this;

java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/Publisher.java

@@ -47,6 +47,7 @@
 import com.google.cloud.pubsub.v1.stub.GrpcPublisherStub;
 import com.google.cloud.pubsub.v1.stub.PublisherStub;
 import com.google.cloud.pubsub.v1.stub.PublisherStubSettings;
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
@@ -151,10 +152,10 @@ public class Publisher implements PublisherInterface {
 
   /**
    * Scale factor to represent decimal token values (e.g. 0.1 refill ratio) as integers inside the
-   * AtomicInteger token bucket. A scale of 100 allows representing decimal ratios down to 0.01
-   * (1%). For example, 1.0 logical token is represented as 100.
+   * AtomicInteger token bucket. A scale of 1000 allows representing decimal ratios down to 0.001.
+   * For example, 1.0 logical token is represented as 1000.
    */
-  private static final int HEDGE_TOKEN_SCALE = 100;
+  private static final int HEDGE_TOKEN_SCALE = 1000;
 
   private final AtomicInteger hedgeTokenBucket = new AtomicInteger();
   private int scaledMaxHedgeTokens;
@@ -303,6 +304,7 @@ public HedgeSettings getHedgeSettings() {
     return hedgeSettings;
   }
 
+  @VisibleForTesting
   Float getHedgeTokenBalance() {
     if (hedgeSettings == null) {
       return null;

java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/HedgeSettingsTest.java

@@ -109,4 +109,16 @@ public void testRefillTooLargeThrows() {
     assertThrows(
         IllegalArgumentException.class, () -> HedgeSettings.newBuilder().setRefillRatio(0.21f));
   }
+
+  @Test
+  public void testRefillTooSmallThrows() {
+    assertThrows(
+        IllegalArgumentException.class, () -> HedgeSettings.newBuilder().setRefillRatio(0.0009f));
+  }
+
+  @Test
+  public void testMinimumRefillValid() {
+    HedgeSettings settings = HedgeSettings.newBuilder().setRefillRatio(0.001f).build();
+    assertEquals(0.001f, settings.getRefillRatio(), 0.0001f);
+  }
 }

java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/PublisherImplTest.java

@@ -1475,6 +1475,21 @@ private void waitForRequests(FakePublisherServiceImpl service, int expectedCount
     }
   }
 
+  @Test
+  public void testTokenBucketRefillRate() throws Exception {
+    Publisher publisher = getPublisherWithHedge(Duration.ofMillis(100), 0.125f, 10);
+    // Starts at 0
+    assertThat(publisher.getHedgeTokenBalance()).isEqualTo(0.0f);
+
+    // Warm up 1 message (should succeed and refill by 0.125)
+    testPublisherServiceImpl.setAutoPublishResponse(true);
+    sendTestMessage(publisher, "refill-warmup").get();
+
+    // Balance should be exactly 0.125
+    assertThat(publisher.getHedgeTokenBalance()).isWithin(0.0001f).of(0.125f);
+    shutdownTestPublisher(publisher);
+  }
+
   @Test
   public void testHedgingNotTriggeredIfFast() throws Exception {
     Publisher publisher = getPublisherWithHedge(Duration.ofMillis(100));