diff --git a/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/Publisher.java b/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/Publisher.java index 1cb5f43133c9..569421adcae3 100644 --- a/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/Publisher.java +++ b/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 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 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 getRetryableCodes() { - return retryableCodes; - } - public void publishAllOutstanding() { OutstandingBatch unorderedOutstandingBatch = null; messagesBatchLock.lock(); @@ -722,7 +717,7 @@ boolean tryAcquireHedgeToken() { } private ApiFuture 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 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 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 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 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(); } diff --git a/java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/FakePublisherServiceImpl.java b/java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/FakePublisherServiceImpl.java index a247c240d650..5b98d118a3a3 100644 --- a/java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/FakePublisherServiceImpl.java +++ b/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 requests = new LinkedBlockingQueue<>(); + private final AtomicInteger cancelledCalls = new AtomicInteger(); private final LinkedBlockingQueue capturedHeaders = new LinkedBlockingQueue<>(); private final LinkedBlockingQueue 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 responseObserver) { + if (responseObserver instanceof ServerCallStreamObserver) { + ((ServerCallStreamObserver) 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 responseObserver) { if (response.isError()) { responseObserver.onError(response.getError()); @@ -175,5 +185,6 @@ public List getCapturedHeaders() { public void clearRequests() { requests.clear(); capturedHeaders.clear(); + cancelledCalls.set(0); } } diff --git a/java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/PublisherImplTest.java b/java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/PublisherImplTest.java index 1571e162da8a..175694ae1573 100644 --- a/java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/PublisherImplTest.java +++ b/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 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 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))