fix(pubsub): make publish hedging robust to clock jumps, queue errors, and shutdown

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

@@ -25,6 +25,7 @@
 import com.google.api.core.ApiFutures;
 import com.google.api.core.BetaApi;
 import com.google.api.core.CurrentMillisClock;
+import com.google.api.core.NanoClock;
 import com.google.api.core.SettableApiFuture;
 import com.google.api.gax.batching.BatchingSettings;
 import com.google.api.gax.batching.FlowControlSettings;
@@ -51,7 +52,6 @@
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
-import com.google.common.collect.ImmutableSet;
 import com.google.protobuf.CodedOutputStream;
 import com.google.protobuf.util.Timestamps;
 import com.google.pubsub.v1.PublishRequest;
@@ -74,7 +74,6 @@
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
-import java.util.Set;
 import java.util.concurrent.Callable;
 import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.CountDownLatch;
@@ -155,7 +154,6 @@ public class Publisher implements PublisherInterface {
   private OpenTelemetryPubsubTracer tracer = new OpenTelemetryPubsubTracer(null, false);
 
   private final HedgingSettings hedgingSettings;
-  private final Set<StatusCode.Code> retryableCodes;
 
   /**
    * Scale factor to represent decimal token values (e.g. 0.1 refill ratio) as integers inside the
@@ -167,7 +165,8 @@ public class Publisher implements PublisherInterface {
   private final AtomicInteger hedgeTokenBucket = new AtomicInteger();
   private int scaledMaxHedgeTokens;
   private int scaledHedgeRefillAmount;
-  private final ApiClock clock;
+  private final ApiClock hedgeClock;
+  private final ApiClock wallClock;
 
   private final ConcurrentLinkedQueue<HedgedRequest> hedgingQueue;
   private final AtomicBoolean isQueueProcessingScheduled;
@@ -263,7 +262,6 @@ private Publisher(Builder builder) throws IOException {
         .setRetrySettings(this.retrySettings)
         .setBatchingSettings(BatchingSettings.newBuilder().setIsEnabled(false).build());
     this.publisherStub = GrpcPublisherStub.create(stubSettings.build());
-    this.retryableCodes = ImmutableSet.copyOf(stubSettings.publishSettings().getRetryableCodes());
     backgroundResourceList.add(publisherStub);
     backgroundResources = new BackgroundResourceAggregation(backgroundResourceList);
     shutdown = new AtomicBoolean(false);
@@ -275,7 +273,8 @@ private Publisher(Builder builder) throws IOException {
           (int) Math.round(this.hedgingSettings.getRefillRatio() * HEDGE_TOKEN_SCALE);
       this.hedgeTokenBucket.set(0);
     }
-    this.clock = builder.clock != null ? builder.clock : CurrentMillisClock.getDefaultClock();
+    this.hedgeClock = builder.clock != null ? builder.clock : NanoClock.getDefaultClock();
+    this.wallClock = builder.clock != null ? builder.clock : CurrentMillisClock.getDefaultClock();
     this.publishContext = GrpcCallContext.createDefault();
     this.publishContextWithCompression =
         GrpcCallContext.createDefault()
@@ -466,10 +465,6 @@ public void run() {
    * wait for the send operations to complete. To wait for messages to send, call {@code get} on the
    * futures returned from {@code publish}.
    */
-  Set<StatusCode.Code> getRetryableCodes() {
-    return retryableCodes;
-  }
-
   public void publishAllOutstanding() {
     OutstandingBatch unorderedOutstandingBatch = null;
     messagesBatchLock.lock();
@@ -722,7 +717,7 @@ boolean tryAcquireHedgeToken() {
   }
 
   private ApiFuture<PublishResponse> startHedgedCall(final OutstandingBatch outstandingBatch) {
-    long deadlineMs = clock.millisTime() + retrySettings.getTotalTimeoutDuration().toMillis();
+    long deadlineMs = hedgeClock.millisTime() + retrySettings.getTotalTimeoutDuration().toMillis();
     final CancellationSharer coordinator =
         new CancellationSharer(outstandingBatch, this, deadlineMs);
 
@@ -747,7 +742,7 @@ public void run() {
     ApiFuture<PublishResponse> firstAttemptFuture = publishCall(outstandingBatch);
     coordinator.addAttempt(0, firstAttemptFuture);
     long delayMs = hedgingSettings.getHedgeDelay().toMillis();
-    HedgedRequest item = new HedgedRequest(coordinator, 1, clock.millisTime() + delayMs);
+    HedgedRequest item = new HedgedRequest(coordinator, 1, hedgeClock.millisTime() + delayMs);
     hedgingQueue.add(item);
     scheduleQueueProcessing();
 
@@ -765,7 +760,7 @@ private void scheduleQueueProcessing() {
         return;
       }
 
-      long delay = Math.max(0, nextItem.getSendAfterMs() - clock.millisTime());
+      long delay = Math.max(0, nextItem.getSendAfterMs() - hedgeClock.millisTime());
 
       queueProcessingFuture =
           executor.schedule(
@@ -784,7 +779,7 @@ public void run() {
   private void processQueue() {
     queueLock.lock();
     try {
-      long now = clock.millisTime();
+      long now = hedgeClock.millisTime();
 
       HedgedRequest item;
       while ((item = hedgingQueue.peek()) != null && item.getSendAfterMs() <= now) {
@@ -796,23 +791,23 @@ private void processQueue() {
           continue;
         }
 
-        long remainingTimeoutMs = coordinator.getAbsoluteDeadlineMs() - clock.millisTime();
+        long remainingTimeoutMs = coordinator.getAbsoluteDeadlineMs() - hedgeClock.millisTime();
         if (remainingTimeoutMs <= 0) {
           continue;
         }
         long attemptTimeoutMs = Math.min(10000, remainingTimeoutMs);
 
         if (tryAcquireHedgeToken()) {
+          // Start Hedged Attempt
+          ApiFuture<PublishResponse> hedgedFuture =
+              publishCall(batch, item.getAttemptNumber(), Duration.ofMillis(attemptTimeoutMs));
+          coordinator.addAttempt(item.getAttemptNumber(), hedgedFuture);
+
           // Clone and schedule next attempt check (Attempt + 1)
           long delayMs = hedgingSettings.getHedgeDelay().toMillis();
           HedgedRequest nextItem =
               new HedgedRequest(coordinator, item.getAttemptNumber() + 1, now + delayMs);
           hedgingQueue.add(nextItem);
-
-          // Start Hedged Attempt
-          ApiFuture<PublishResponse> hedgedFuture =
-              publishCall(batch, item.getAttemptNumber(), Duration.ofMillis(attemptTimeoutMs));
-          coordinator.addAttempt(item.getAttemptNumber(), hedgedFuture);
         } else {
           loggingUtil.logPublisher(
               LoggingUtil.SubSystem.PUBLISH_HEDGED,
@@ -821,10 +816,15 @@ private void processQueue() {
               batch.getMessageWrappers().get(0));
         }
       }
-      isQueueProcessingScheduled.set(false);
-      scheduleQueueProcessing();
+    } catch (RuntimeException e) {
+      logger.log(Level.WARNING, "Unexpected error while processing the publish hedging queue", e);
     } finally {
-      queueLock.unlock();
+      try {
+        isQueueProcessingScheduled.set(false);
+        scheduleQueueProcessing();
+      } finally {
+        queueLock.unlock();
+      }
     }
   }
 
@@ -841,7 +841,7 @@ final class OutstandingBatch {
         List<OutstandingPublish> outstandingPublishes, int batchSizeBytes, String orderingKey) {
       this.outstandingPublishes = outstandingPublishes;
       attempt = 1;
-      creationTime = clock.millisTime();
+      creationTime = wallClock.millisTime();
       this.batchSizeBytes = batchSizeBytes;
       this.orderingKey = orderingKey;
     }
@@ -919,14 +919,18 @@ public void shutdown() {
     if (currentAlarmFuture != null && activeAlarm.getAndSet(false)) {
       currentAlarmFuture.cancel(false);
     }
-    if (queueProcessingFuture != null) {
-      queueProcessingFuture.cancel(false);
-    }
-    if (hedgingQueue != null) {
-      hedgingQueue.clear();
-    }
     publishAllOutstanding();
     messagesWaiter.waitComplete();
+    // Stop the hedging timer and clear the queue to safely shut down.
+    queueLock.lock();
+    try {
+      if (queueProcessingFuture != null) {
+        queueProcessingFuture.cancel(false);
+      }
+      hedgingQueue.clear();
+    } finally {
+      queueLock.unlock();
+    }
     backgroundResources.shutdown();
   }
 

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

@@ -21,6 +21,7 @@
 import com.google.pubsub.v1.PublishResponse;
 import com.google.pubsub.v1.PublisherGrpc.PublisherImplBase;
 import io.grpc.Metadata;
+import io.grpc.stub.ServerCallStreamObserver;
 import io.grpc.stub.StreamObserver;
 import java.time.Duration;
 import java.util.ArrayList;
@@ -37,6 +38,7 @@
 class FakePublisherServiceImpl extends PublisherImplBase {
 
   private final LinkedBlockingQueue<PublishRequest> requests = new LinkedBlockingQueue<>();
+  private final AtomicInteger cancelledCalls = new AtomicInteger();
   private final LinkedBlockingQueue<Metadata> capturedHeaders = new LinkedBlockingQueue<>();
   private final LinkedBlockingQueue<Response> publishResponses = new LinkedBlockingQueue<>();
   private final AtomicInteger nextMessageId = new AtomicInteger(1);
@@ -83,6 +85,10 @@ public String toString() {
   @Override
   public void publish(
       PublishRequest request, final StreamObserver<PublishResponse> responseObserver) {
+    if (responseObserver instanceof ServerCallStreamObserver) {
+      ((ServerCallStreamObserver<PublishResponse>) responseObserver)
+          .setOnCancelHandler(() -> cancelledCalls.incrementAndGet());
+    }
     Response response;
     try {
       if (autoPublishResponse) {
@@ -115,6 +121,10 @@ public void run() {
     }
   }
 
+  public int getCancelledCallCount() {
+    return cancelledCalls.get();
+  }
+
   private void sendResponse(Response response, StreamObserver<PublishResponse> responseObserver) {
     if (response.isError()) {
       responseObserver.onError(response.getError());
@@ -175,5 +185,6 @@ public List<Metadata> getCapturedHeaders() {
   public void clearRequests() {
     requests.clear();
     capturedHeaders.clear();
+    cancelledCalls.set(0);
   }
 }

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

@@ -1513,7 +1513,8 @@ public void testTokenBucketRefillRate() throws Exception {
 
   @Test
   public void testHedgingNotTriggeredIfFast() throws Exception {
-    Publisher publisher = getPublisherWithHedge(Duration.ofMillis(100));
+    Publisher publisher = getPublisherWithHedge(Duration.ofMillis(100), 0.2f, 20);
+    fillTokenBucket(publisher, 5);
 
     // Prepare fast response (10ms delay)
     testPublisherServiceImpl.setAutoPublishResponse(false);
@@ -1528,10 +1529,11 @@ public void testHedgingNotTriggeredIfFast() throws Exception {
 
     // Future should be completed
     assertEquals("1", future.get());
+    fakeExecutor.advanceTime(Duration.ofMillis(200));
 
     // Only 1 request should be received by server
     assertThat(testPublisherServiceImpl.getCapturedRequests()).hasSize(1);
-
+    assertThat(publisher.getHedgeTokenBalance()).isWithin(0.0001f).of(1.2f);
     shutdownTestPublisher(publisher);
   }
 
@@ -1610,12 +1612,15 @@ public void testMultipleHedging() throws Exception {
 
   @Test
   public void testHedgingBypassedIfNoTokens() throws Exception {
-    Publisher publisher = getPublisherWithHedge(Duration.ofMillis(100));
+    Publisher publisher = getPublisherWithHedge(Duration.ofMillis(100), 0.2f, 20);
 
-    // Drain the token bucket completely (since it starts full)
-    while (publisher.tryAcquireHedgeToken()) {}
+    // The bucket starts empty. Earn one token (5 successes x 0.2), then spend it.
+    fillTokenBucket(publisher, 5);
+    assertTrue(publisher.tryAcquireHedgeToken());
+    assertFalse(publisher.tryAcquireHedgeToken());
     assertThat(publisher.getHedgeTokenBalance()).isEqualTo(0.0f);
 
+    testPublisherServiceImpl.setAutoPublishResponse(false);
     testPublisherServiceImpl.setPublishResponseDelay(Duration.ofMillis(200));
     testPublisherServiceImpl.addPublishResponse(PublishResponse.newBuilder().addMessageIds("1"));
 
@@ -1728,6 +1733,94 @@ public void testPermanentErrorOnHedgedAttemptDiscarded() throws Exception {
     shutdownTestPublisher(publisher);
   }
 
+  @Test
+  public void testHedgedAttemptWins() throws Exception {
+    Publisher publisher = getPublisherWithHedge(Duration.ofMillis(100), 0.2f, 20);
+    fillTokenBucket(publisher, 5); // 5 successes x 0.2 refill = 1 token
+
+    // Original attempt: slow (500ms).
+    testPublisherServiceImpl.setAutoPublishResponse(false);
+    testPublisherServiceImpl.setPublishResponseDelay(Duration.ofMillis(500));
+    testPublisherServiceImpl.addPublishResponse(PublishResponse.newBuilder().addMessageIds("1"));
+
+    ApiFuture<String> future = sendTestMessage(publisher, "msg-hedge-wins");
+    waitForRequests(testPublisherServiceImpl, 1);
+
+    // Hedged attempt: answered immediately.
+    testPublisherServiceImpl.setPublishResponseDelay(Duration.ZERO);
+    testPublisherServiceImpl.addPublishResponse(PublishResponse.newBuilder().addMessageIds("2"));
+
+    // Past the 100ms hedge delay: the hedge is sent and wins well before the original's 500ms.
+    fakeExecutor.advanceTime(Duration.ofMillis(120));
+    waitForRequests(testPublisherServiceImpl, 2);
+    assertEquals("2", future.get(5, TimeUnit.SECONDS));
+
+    // The winning request was the hedge.
+    PubsubClientTelemetry hedgedTelemetry =
+        extractTelemetryHeader(testPublisherServiceImpl.getCapturedHeaders().get(1));
+    assertThat(hedgedTelemetry.getPublishOperation().getHedgedAttemptCount()).isEqualTo(1);
+
+    // The hedge spent the only token, and the win refilled 0.2.
+    assertThat(publisher.getHedgeTokenBalance()).isWithin(0.0001f).of(0.2f);
+
+    // The losing original RPC was actually cancelled on the server.
+    waitForCancelledCalls(testPublisherServiceImpl, 1);
+
+    shutdownTestPublisher(publisher);
+  }
+
+  private void waitForCancelledCalls(FakePublisherServiceImpl service, int expectedCount)
+      throws InterruptedException {
+    long timeout = System.currentTimeMillis() + 5000;
+    while (service.getCancelledCallCount() < expectedCount
+        && System.currentTimeMillis() < timeout) {
+      Thread.sleep(5);
+    }
+    assertThat(service.getCancelledCallCount()).isEqualTo(expectedCount);
+  }
+
+  @Test
+  public void testInFlightBatchStillHedgedDuringShutdown() throws Exception {
+    Publisher publisher = getPublisherWithHedge(Duration.ofMillis(100), 0.2f, 20);
+    // Fills 1 token and, like steady traffic, leaves a hedge timer pending.
+    fillTokenBucket(publisher, 5);
+
+    // Original attempt is slow; the hedge is answered immediately.
+    testPublisherServiceImpl.setAutoPublishResponse(false);
+    testPublisherServiceImpl.setPublishResponseDelay(Duration.ofMillis(500));
+    testPublisherServiceImpl.addPublishResponse(PublishResponse.newBuilder().addMessageIds("1"));
+    ApiFuture<String> future = sendTestMessage(publisher, "msg-during-shutdown");
+    waitForRequests(testPublisherServiceImpl, 1);
+    testPublisherServiceImpl.setPublishResponseDelay(Duration.ZERO);
+    testPublisherServiceImpl.addPublishResponse(PublishResponse.newBuilder().addMessageIds("2"));
+
+    // Start shutdown() and wait until it's parked in waitComplete().
+    Thread shutdownThread = new Thread(publisher::shutdown);
+    shutdownThread.setDaemon(true); // don't hang the JVM if the test fails
+    shutdownThread.start();
+    long timeout = System.currentTimeMillis() + 5000;
+    while (shutdownThread.getState() != Thread.State.WAITING
+        && shutdownThread.isAlive()
+        && System.currentTimeMillis() < timeout) {
+      Thread.sleep(5);
+    }
+    assertEquals(Thread.State.WAITING, shutdownThread.getState());
+
+    // The pending hedge timer still fires during shutdown, and the hedge completes the batch.
+    fakeExecutor.advanceTime(Duration.ofMillis(120));
+    waitForRequests(testPublisherServiceImpl, 2);
+    assertEquals("2", future.get(5, TimeUnit.SECONDS));
+    PubsubClientTelemetry hedgedTelemetry =
+        extractTelemetryHeader(testPublisherServiceImpl.getCapturedHeaders().get(1));
+    assertThat(hedgedTelemetry.getPublishOperation().getHedgedAttemptCount()).isEqualTo(1);
+
+    // shutdown() can now finish.
+    shutdownThread.join(5000);
+    assertFalse(shutdownThread.isAlive());
+    fakeExecutor.advanceTime(Duration.ofSeconds(10));
+    assertTrue(publisher.awaitTermination(1, TimeUnit.MINUTES));
+  }
+
   private Builder getTestPublisherBuilder() {
     return Publisher.newBuilder(TEST_TOPIC)
         .setExecutorProvider(FixedExecutorProvider.create(fakeExecutor))