diff --git a/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/CancellationSharer.java b/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/CancellationSharer.java new file mode 100644 index 000000000000..571c8a99d8ac --- /dev/null +++ b/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/CancellationSharer.java @@ -0,0 +1,143 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.pubsub.v1; + +import com.google.api.core.AbstractApiFuture; +import com.google.api.core.ApiFuture; +import com.google.api.core.ApiFutureCallback; +import com.google.api.core.ApiFutures; +import com.google.common.util.concurrent.MoreExecutors; +import com.google.pubsub.v1.PublishResponse; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; + +/** + * Coordinates multiple publish attempts for a single batch of messages. + * + *

Implements {@link ApiFuture} to act as the single future returned to the publisher's client. + * It manages the lifecycle of the original attempt and any subsequent hedged attempts. + */ +class CancellationSharer extends AbstractApiFuture { + private final Publisher.OutstandingBatch batch; + private final Publisher publisher; + private final Map> runningAttempts = + new ConcurrentHashMap<>(); + private final AtomicInteger totalAttemptsCount = new AtomicInteger(0); + private final AtomicBoolean done = new AtomicBoolean(false); + final AtomicBoolean isInQueue = new AtomicBoolean(false); + private final AtomicReference lastError = new AtomicReference<>(); + + CancellationSharer(Publisher.OutstandingBatch batch, Publisher publisher) { + this.batch = batch; + this.publisher = publisher; + } + + /** + * Adds an attempt to be tracked by this coordinator. + * + * @param attemptNumber the 1-based index of the attempt (1 is original, 2+ are hedged) + * @param future the future representing the gRPC call for this attempt + */ + void addAttempt(final int attemptNumber, ApiFuture future) { + runningAttempts.put(attemptNumber, future); + totalAttemptsCount.incrementAndGet(); + + if (done.get()) { + future.cancel(true); + runningAttempts.remove(attemptNumber); + return; + } + + ApiFutures.addCallback( + future, + new ApiFutureCallback() { + @Override + public void onSuccess(PublishResponse result) { + handleAttemptSuccess(attemptNumber, result); + } + + @Override + public void onFailure(Throwable t) { + handleAttemptFailure(attemptNumber, t); + } + }, + MoreExecutors.directExecutor()); + } + + private void handleAttemptSuccess(int attemptNumber, PublishResponse response) { + if (done.compareAndSet(false, true)) { + set(response); // Resolve parent future + cancelAllExcept(attemptNumber); + publisher.refillTokenBucket(); + } + } + + private void handleAttemptFailure(int attemptNumber, Throwable t) { + runningAttempts.remove(attemptNumber); + + if (!done.get()) { + lastError.set(t); + if (runningAttempts.isEmpty() && !isInQueue.get()) { + if (done.compareAndSet(false, true)) { + setException(lastError.get()); + } + } + } + } + + void checkCompletionOnQueueExit() { + if (!done.get() && runningAttempts.isEmpty() && !isInQueue.get()) { + if (done.compareAndSet(false, true)) { + Throwable error = lastError.get(); + setException( + error != null ? error : new RuntimeException("Hedging failed with no active attempts")); + } + } + } + + private void cancelAllExcept(int successfulAttempt) { + for (Map.Entry> entry : runningAttempts.entrySet()) { + if (entry.getKey() != successfulAttempt) { + entry.getValue().cancel(true); + } + } + } + + @Override + public boolean cancel(boolean mayInterruptIfRunning) { + if (super.cancel(mayInterruptIfRunning)) { + if (done.compareAndSet(false, true)) { + for (ApiFuture future : runningAttempts.values()) { + future.cancel(mayInterruptIfRunning); + } + return true; + } + } + return false; + } + + int getAttemptCount() { + return totalAttemptsCount.get(); + } + + Publisher.OutstandingBatch getBatch() { + return batch; + } +} diff --git a/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/HedgeSettings.java b/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/HedgeSettings.java new file mode 100644 index 000000000000..fdbfebbe9106 --- /dev/null +++ b/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/HedgeSettings.java @@ -0,0 +1,111 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.pubsub.v1; + +import com.google.common.base.Preconditions; +import java.time.Duration; + +/** Settings for configuring publish hedging. */ +public final class HedgeSettings { + /** Default hedging delay. */ + private static final Duration DEFAULT_DELAY = Duration.ofMillis(100); + + /** Default maximum number of tokens in the bucket. */ + private static final int DEFAULT_MAX_TOKENS = 100; + + /** Default refill rate (tokens per successful request). */ + private static final float DEFAULT_REFILL = 0.1f; + + /** Hedging delay. */ + private final Duration hedgeDelay; + + /** Maximum tokens. */ + private final int maxTokens; + + /** Refill rate. */ + private final float refill; + + private HedgeSettings(final Builder builder) { + this.hedgeDelay = builder.hedgeDelay; + this.maxTokens = builder.maxTokens; + this.refill = builder.refill; + } + + /** + * Returns the configured hedging delay. + * + * @return the hedging delay. + */ + Duration getHedgeDelay() { + return hedgeDelay; + } + + int getMaxTokens() { + return maxTokens; + } + + float getRefill() { + return refill; + } + + /** + * Returns a new builder for {@code HedgeSettings}. + * + * @return a new builder. + */ + public static Builder newBuilder() { + return new Builder(); + } + + /** Builder for {@code HedgeSettings}. */ + public static final class Builder { + /** Hedging delay. */ + private Duration hedgeDelay = DEFAULT_DELAY; + + /** Maximum tokens. */ + private int maxTokens = DEFAULT_MAX_TOKENS; + + /** Refill rate. */ + private float refill = DEFAULT_REFILL; + + private Builder() {} + + /** + * Allows hedging delay to be configurable. + * + * @param delay the hedging delay, must be positive. + * @return this builder. + */ + public Builder setHedgeDelay(final Duration delay) { + Preconditions.checkNotNull(delay); + if (delay.isNegative() || delay.isZero()) { + throw new IllegalArgumentException("delay must be positive"); + } + this.hedgeDelay = delay; + return this; + } + + /** + * Builds an instance of {@code HedgeSettings}. + * + * @return the built {@code HedgeSettings} instance. + */ + public HedgeSettings build() { + return new HedgeSettings(this); + } + } +} diff --git a/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/HedgeTokenBucket.java b/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/HedgeTokenBucket.java new file mode 100644 index 000000000000..5378478cb429 --- /dev/null +++ b/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/HedgeTokenBucket.java @@ -0,0 +1,65 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.pubsub.v1; + +import com.google.common.annotations.VisibleForTesting; +import javax.annotation.concurrent.GuardedBy; + +/** Token bucket for limiting hedged requests. Thread-safe. */ +final class HedgeTokenBucket { + private final int maxTokens; + private final float refillAmount; + + @GuardedBy("this") + private float tokens; + + HedgeTokenBucket(HedgeSettings settings) { + this.maxTokens = settings.getMaxTokens(); + this.refillAmount = settings.getRefill(); + this.tokens = maxTokens; + } + + @VisibleForTesting + HedgeTokenBucket(int maxTokens, float refillAmount) { + this.maxTokens = maxTokens; + this.refillAmount = refillAmount; + this.tokens = maxTokens; + } + + /** + * Attempts to acquire a token for a hedged request. + * + * @return {@code true} if a token was acquired, {@code false} otherwise. + */ + synchronized boolean tryAcquire() { + if (tokens >= 1.0f) { + tokens -= 1.0f; + return true; + } + return false; + } + + /** Refills the bucket by the configured refill amount, capped at max tokens. */ + synchronized void refill() { + tokens = Math.min(maxTokens, tokens + refillAmount); + } + + @VisibleForTesting + synchronized float getTokens() { + return tokens; + } +} diff --git a/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/HedgedRequest.java b/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/HedgedRequest.java new file mode 100644 index 000000000000..ad07de9a8a90 --- /dev/null +++ b/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/HedgedRequest.java @@ -0,0 +1,42 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.pubsub.v1; + +/** Represents a pending hedging check in the publisher's queue. */ +class HedgedRequest { + private final CancellationSharer coordinator; + private final int attemptNumber; + private final long sendAfterMs; + + HedgedRequest(CancellationSharer coordinator, int attemptNumber, long sendAfterMs) { + this.coordinator = coordinator; + this.attemptNumber = attemptNumber; + this.sendAfterMs = sendAfterMs; + } + + CancellationSharer getCoordinator() { + return coordinator; + } + + int getAttemptNumber() { + return attemptNumber; + } + + long getSendAfterMs() { + return sendAfterMs; + } +} 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 56c920bcfdc1..35ab918e5645 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 @@ -18,11 +18,13 @@ import static com.google.common.util.concurrent.MoreExecutors.directExecutor; +import com.google.api.core.ApiClock; import com.google.api.core.ApiFunction; import com.google.api.core.ApiFuture; import com.google.api.core.ApiFutureCallback; import com.google.api.core.ApiFutures; import com.google.api.core.BetaApi; +import com.google.api.core.CurrentMillisClock; import com.google.api.core.SettableApiFuture; import com.google.api.gax.batching.BatchingSettings; import com.google.api.gax.batching.FlowControlSettings; @@ -65,12 +67,14 @@ import java.util.List; import java.util.Map; import java.util.concurrent.Callable; +import java.util.concurrent.ConcurrentLinkedQueue; import java.util.concurrent.CountDownLatch; import java.util.concurrent.Executor; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.ScheduledFuture; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.locks.Lock; import java.util.concurrent.locks.ReentrantLock; import java.util.logging.Level; @@ -138,6 +142,15 @@ public class Publisher implements PublisherInterface { private final OpenTelemetry openTelemetry; private OpenTelemetryPubsubTracer tracer = new OpenTelemetryPubsubTracer(null, false); + private final HedgeSettings hedgeSettings; + private final HedgeTokenBucket hedgeTokenBucket; + private final ApiClock clock; + + private final ConcurrentLinkedQueue hedgingQueue; + private final AtomicBoolean isQueueProcessingScheduled; + private ScheduledFuture queueProcessingFuture; + private final Object queueLock; + /** The maximum number of messages in one request. Defined by the API. */ public static long getApiMaxRequestElementCount() { return 1000L; @@ -230,10 +243,18 @@ private Publisher(Builder builder) throws IOException { backgroundResources = new BackgroundResourceAggregation(backgroundResourceList); shutdown = new AtomicBoolean(false); messagesWaiter = new Waiter(); + this.hedgeSettings = builder.hedgeSettings; + this.hedgeTokenBucket = + this.hedgeSettings != null ? new HedgeTokenBucket(this.hedgeSettings) : null; + this.clock = builder.clock != null ? builder.clock : CurrentMillisClock.getDefaultClock(); this.publishContext = GrpcCallContext.createDefault(); this.publishContextWithCompression = GrpcCallContext.createDefault() .withCallOptions(CallOptions.DEFAULT.withCompression(GZIP_COMPRESSION)); + this.hedgingQueue = new ConcurrentLinkedQueue<>(); + this.isQueueProcessingScheduled = new AtomicBoolean(false); + this.queueLock = new Object(); + this.queueProcessingFuture = null; } /** Topic which the publisher publishes to. */ @@ -246,6 +267,15 @@ public String getTopicNameString() { return topicName; } + /** Returns the configured hedging settings, or null if hedging is disabled. */ + public HedgeSettings getHedgeSettings() { + return hedgeSettings; + } + + HedgeTokenBucket getHedgeTokenBucket() { + return hedgeTokenBucket; + } + /** * Schedules the publishing of a message. The publishing of the message may occur immediately or * be delayed based on the publisher batching options. @@ -572,7 +602,11 @@ public void onFailure(Throwable t) { ApiFuture future; Executor callbackExecutor = directExecutor(); if (outstandingBatch.orderingKey == null || outstandingBatch.orderingKey.isEmpty()) { - future = publishCall(outstandingBatch); + if (hedgeSettings != null) { + future = startHedgedCall(outstandingBatch); + } else { + future = publishCall(outstandingBatch); + } } else { // If ordering key is specified, publish the batch using the sequential executor. future = @@ -588,7 +622,13 @@ public ApiFuture call() { ApiFutures.addCallback(future, futureCallback, callbackExecutor); } - private final class OutstandingBatch { + void refillTokenBucket() { + if (hedgeTokenBucket != null) { + hedgeTokenBucket.refill(); + } + } + + final class OutstandingBatch { final List outstandingPublishes; final long creationTime; int attempt; @@ -600,7 +640,7 @@ private final class OutstandingBatch { List outstandingPublishes, int batchSizeBytes, String orderingKey) { this.outstandingPublishes = outstandingPublishes; attempt = 1; - creationTime = System.currentTimeMillis(); + creationTime = clock.millisTime(); this.batchSizeBytes = batchSizeBytes; this.orderingKey = orderingKey; } @@ -677,6 +717,9 @@ public void shutdown() { if (currentAlarmFuture != null && activeAlarm.getAndSet(false)) { currentAlarmFuture.cancel(false); } + if (queueProcessingFuture != null) { + queueProcessingFuture.cancel(false); + } publishAllOutstanding(); messagesWaiter.waitComplete(); backgroundResources.shutdown(); @@ -814,6 +857,8 @@ public PubsubMessage apply(PubsubMessage input) { private boolean enableOpenTelemetryTracing = false; private OpenTelemetry openTelemetry = null; + private HedgeSettings hedgeSettings = null; + ApiClock clock = null; private Builder(String topic) { this.topicName = Preconditions.checkNotNull(topic); @@ -966,12 +1011,26 @@ public Builder setOpenTelemetry(OpenTelemetry openTelemetry) { return this; } + /** Configures the Publisher's hedging parameters. */ + public Builder setHedgeSettings(HedgeSettings hedgeSettings) { + this.hedgeSettings = hedgeSettings; + return this; + } + + Builder setClock(ApiClock clock) { + this.clock = clock; + return this; + } + /** Returns the default BatchingSettings used by the client if settings are not provided. */ public static BatchingSettings getDefaultBatchingSettings() { return DEFAULT_BATCHING_SETTINGS; } public Publisher build() throws IOException { + Preconditions.checkState( + !(enableMessageOrdering && hedgeSettings != null), + "Publish hedging and message ordering cannot be enabled at the same time."); return new Publisher(this); } } @@ -1182,4 +1241,96 @@ && getBatchedBytes() + outstandingPublish.messageSize >= getMaxBatchBytes()) { return batchesToSend; } } + + private ApiFuture startHedgedCall(final OutstandingBatch outstandingBatch) { + final CancellationSharer coordinator = new CancellationSharer(outstandingBatch, this); + + // Register cancellation listeners on client futures to propagate cancel to coordinator + final AtomicInteger cancelledCount = new AtomicInteger(0); + final int batchSize = outstandingBatch.outstandingPublishes.size(); + for (final OutstandingPublish outstanding : outstandingBatch.outstandingPublishes) { + outstanding.publishResult.addListener( + new Runnable() { + @Override + public void run() { + if (outstanding.publishResult.isCancelled()) { + if (cancelledCount.incrementAndGet() == batchSize) { + coordinator.cancel(true); + } + } + } + }, + directExecutor()); + } + + ApiFuture firstAttemptFuture = publishCall(outstandingBatch); + coordinator.addAttempt(1, firstAttemptFuture); + long delayMs = hedgeSettings.getHedgeDelay().toMillis(); + HedgedRequest item = new HedgedRequest(coordinator, 2, clock.millisTime() + delayMs); + hedgingQueue.add(item); + coordinator.isInQueue.set(true); + scheduleQueueProcessing(); + + return coordinator; + } + + private void scheduleQueueProcessing() { + if (isQueueProcessingScheduled.compareAndSet(false, true)) { + HedgedRequest nextItem = hedgingQueue.peek(); + if (nextItem == null) { + isQueueProcessingScheduled.set(false); + return; + } + + long delay = nextItem.getSendAfterMs() - clock.millisTime(); + delay = Math.max(0, delay); + + queueProcessingFuture = + executor.schedule( + new Runnable() { + @Override + public void run() { + processQueue(); + } + }, + delay, + TimeUnit.MILLISECONDS); + } + } + + private void processQueue() { + synchronized (queueLock) { + isQueueProcessingScheduled.set(false); + long now = clock.millisTime(); + + HedgedRequest item; + while ((item = hedgingQueue.peek()) != null && item.getSendAfterMs() <= now) { + hedgingQueue.poll(); + + CancellationSharer coordinator = item.getCoordinator(); + if (coordinator.isDone()) { + coordinator.isInQueue.set(false); + continue; + } + + if (hedgeTokenBucket.tryAcquire()) { + // Clone and schedule next attempt check (Attempt + 1) + long delayMs = hedgeSettings.getHedgeDelay().toMillis(); + HedgedRequest nextItem = + new HedgedRequest(coordinator, item.getAttemptNumber() + 1, now + delayMs); + hedgingQueue.add(nextItem); + + // Start Hedged Attempt + ApiFuture hedgedFuture = publishCall(coordinator.getBatch()); + coordinator.addAttempt(item.getAttemptNumber(), hedgedFuture); + } else { + coordinator.isInQueue.set(false); + coordinator.checkCompletionOnQueueExit(); + } + } + + // Reschedule for next items + scheduleQueueProcessing(); + } + } } 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 9ab1dec73471..9424018b382f 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 @@ -81,7 +81,6 @@ public String toString() { @Override public void publish( PublishRequest request, final StreamObserver responseObserver) { - requests.add(request); Response response; try { if (autoPublishResponse) { @@ -98,6 +97,7 @@ public void publish( } if (responseDelay == Duration.ZERO) { sendResponse(response, responseObserver); + requests.add(request); } else { final Response responseToSend = response; executor.schedule( @@ -109,6 +109,7 @@ public void run() { }, responseDelay.toMillis(), TimeUnit.MILLISECONDS); + requests.add(request); } } diff --git a/java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/HedgeSettingsTest.java b/java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/HedgeSettingsTest.java new file mode 100644 index 000000000000..f3b26aedb3c3 --- /dev/null +++ b/java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/HedgeSettingsTest.java @@ -0,0 +1,66 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.pubsub.v1; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertThrows; + +import java.time.Duration; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class HedgeSettingsTest { + + @Test + public void testDefaultSettings() { + HedgeSettings settings = HedgeSettings.newBuilder().build(); + assertNotNull(settings); + assertEquals(Duration.ofMillis(100), settings.getHedgeDelay()); + assertEquals(100, settings.getMaxTokens()); + assertEquals(0.1f, settings.getRefill(), 0.0001f); + } + + @Test + public void testCustomDelay() { + Duration customDelay = Duration.ofMillis(50); + HedgeSettings settings = HedgeSettings.newBuilder().setHedgeDelay(customDelay).build(); + assertNotNull(settings); + assertEquals(customDelay, settings.getHedgeDelay()); + } + + @Test + public void testNegativeDelayThrows() { + assertThrows( + IllegalArgumentException.class, + () -> HedgeSettings.newBuilder().setHedgeDelay(Duration.ofMillis(-10))); + } + + @Test + public void testZeroDelayThrows() { + assertThrows( + IllegalArgumentException.class, + () -> HedgeSettings.newBuilder().setHedgeDelay(Duration.ZERO)); + } + + @Test + public void testNullDelayThrows() { + assertThrows(NullPointerException.class, () -> HedgeSettings.newBuilder().setHedgeDelay(null)); + } +} diff --git a/java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/HedgeTokenBucketTest.java b/java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/HedgeTokenBucketTest.java new file mode 100644 index 000000000000..e92ecd9a6ce2 --- /dev/null +++ b/java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/HedgeTokenBucketTest.java @@ -0,0 +1,71 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.pubsub.v1; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class HedgeTokenBucketTest { + + @Test + public void testInitialState() { + HedgeTokenBucket bucket = new HedgeTokenBucket(10, 0.5f); + assertEquals(10.0f, bucket.getTokens(), 0.0001f); + } + + @Test + public void testAcquire() { + HedgeTokenBucket bucket = new HedgeTokenBucket(10, 0.5f); + assertTrue(bucket.tryAcquire()); + assertEquals(9.0f, bucket.getTokens(), 0.0001f); + } + + @Test + public void testAcquireUntilEmpty() { + HedgeTokenBucket bucket = new HedgeTokenBucket(2, 0.5f); + assertTrue(bucket.tryAcquire()); + assertTrue(bucket.tryAcquire()); + assertFalse(bucket.tryAcquire()); + assertEquals(0.0f, bucket.getTokens(), 0.0001f); + } + + @Test + public void testRefill() { + HedgeTokenBucket bucket = new HedgeTokenBucket(2, 0.5f); + assertTrue(bucket.tryAcquire()); // 1.0 left + assertTrue(bucket.tryAcquire()); // 0.0 left + + bucket.refill(); // 0.5 tokens + assertFalse(bucket.tryAcquire()); // needs 1.0 + + bucket.refill(); // 1.0 tokens + assertTrue(bucket.tryAcquire()); // succeeds, 0.0 left + } + + @Test + public void testRefillCap() { + HedgeTokenBucket bucket = new HedgeTokenBucket(2, 0.5f); + bucket.refill(); // already at max (2) + assertEquals(2.0f, bucket.getTokens(), 0.0001f); + } +} 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 8e6efaf372c9..6acfd76a5870 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 @@ -105,6 +105,7 @@ public void setUp() throws Exception { testServer.start(); fakeExecutor = new FakeScheduledExecutorService(); + testPublisherServiceImpl.setExecutor(fakeExecutor); } @After @@ -1339,6 +1340,197 @@ public void testPublishOpenTelemetryTracing() throws Exception { .hasEnded(); } + @Test + public void testPublisherWithHedgeSettings() throws Exception { + HedgeSettings hedgeSettings = + HedgeSettings.newBuilder().setHedgeDelay(Duration.ofMillis(50)).build(); + Publisher publisher = getTestPublisherBuilder().setHedgeSettings(hedgeSettings).build(); + + assertThat(publisher.getHedgeSettings()).isEqualTo(hedgeSettings); + assertThat(publisher.getHedgeTokenBucket()).isNotNull(); + assertThat(publisher.getHedgeTokenBucket().getTokens()).isWithin(0.0001f).of(100.0f); + + shutdownTestPublisher(publisher); + } + + @Test + public void testPublisherWithoutHedgeSettings() throws Exception { + Publisher publisher = getTestPublisherBuilder().build(); + + assertThat(publisher.getHedgeSettings()).isNull(); + assertThat(publisher.getHedgeTokenBucket()).isNull(); + + shutdownTestPublisher(publisher); + } + + private Publisher getPublisherWithHedge(Duration delay) throws Exception { + HedgeSettings hedgeSettings = HedgeSettings.newBuilder().setHedgeDelay(delay).build(); + return getTestPublisherBuilder() + .setHedgeSettings(hedgeSettings) + .setClock(fakeExecutor.getClock()) + .setBatchingSettings( + Publisher.Builder.DEFAULT_BATCHING_SETTINGS.toBuilder() + .setElementCountThreshold(1L) + .build()) + .build(); + } + + private void waitForRequests(FakePublisherServiceImpl service, int expectedCount) + throws InterruptedException { + long timeout = System.currentTimeMillis() + 5000; + while (service.getCapturedRequests().size() < expectedCount + && System.currentTimeMillis() < timeout) { + Thread.sleep(5); + } + if (service.getCapturedRequests().size() < expectedCount) { + throw new AssertionError( + String.format( + "Timed out waiting for requests. Expected: %d, Got: %d", + expectedCount, service.getCapturedRequests().size())); + } + } + + @Test + public void testHedgingNotTriggeredIfFast() throws Exception { + Publisher publisher = getPublisherWithHedge(Duration.ofMillis(50)); + + // Prepare fast response (10ms delay) + testPublisherServiceImpl.setAutoPublishResponse(false); + testPublisherServiceImpl.setPublishResponseDelay(Duration.ofMillis(10)); + testPublisherServiceImpl.addPublishResponse(PublishResponse.newBuilder().addMessageIds("1")); + + ApiFuture future = sendTestMessage(publisher, "msg-fast"); + waitForRequests(testPublisherServiceImpl, 1); + + // Advance time past response but before hedge delay (e.g. 20ms) + fakeExecutor.advanceTime(Duration.ofMillis(20)); + + // Future should be completed + assertEquals("1", future.get()); + + // Only 1 request should be received by server + assertThat(testPublisherServiceImpl.getCapturedRequests()).hasSize(1); + + shutdownTestPublisher(publisher); + } + + @Test + public void testHedgingTriggeredIfSlow() throws Exception { + Publisher publisher = getPublisherWithHedge(Duration.ofMillis(50)); + + // Set response delay to 100ms (greater than 50ms hedge delay) + testPublisherServiceImpl.setAutoPublishResponse(false); + testPublisherServiceImpl.setPublishResponseDelay(Duration.ofMillis(100)); + // Add two responses (one for main, one for hedge) + testPublisherServiceImpl.addPublishResponse(PublishResponse.newBuilder().addMessageIds("1")); + testPublisherServiceImpl.addPublishResponse(PublishResponse.newBuilder().addMessageIds("2")); + + ApiFuture future = sendTestMessage(publisher, "msg-slow"); + waitForRequests(testPublisherServiceImpl, 1); + + // Advance time to 40ms (before hedge delay) + fakeExecutor.advanceTime(Duration.ofMillis(40)); + assertThat(testPublisherServiceImpl.getCapturedRequests()).hasSize(1); // Only original sent + + // Advance time to 60ms (past 50ms hedge delay) + fakeExecutor.advanceTime(Duration.ofMillis(20)); + waitForRequests(testPublisherServiceImpl, 2); + + // Now attempt 2 should have been triggered + assertThat(testPublisherServiceImpl.getCapturedRequests()).hasSize(2); + + // Advance to 110ms to let responses complete + fakeExecutor.advanceTime(Duration.ofMillis(50)); + assertTrue(future.isDone()); + + shutdownTestPublisher(publisher); + } + + @Test + public void testMultipleHedging() throws Exception { + Publisher publisher = getPublisherWithHedge(Duration.ofMillis(50)); + + // Set delay to 200ms + testPublisherServiceImpl.setAutoPublishResponse(false); + testPublisherServiceImpl.setPublishResponseDelay(Duration.ofMillis(200)); + // Add responses for 3 attempts + testPublisherServiceImpl.addPublishResponse(PublishResponse.newBuilder().addMessageIds("1")); + testPublisherServiceImpl.addPublishResponse(PublishResponse.newBuilder().addMessageIds("2")); + testPublisherServiceImpl.addPublishResponse(PublishResponse.newBuilder().addMessageIds("3")); + + ApiFuture future = sendTestMessage(publisher, "msg-very-slow"); + waitForRequests(testPublisherServiceImpl, 1); + + // T=0: Attempt 1 sent. + // T=60 (Hedge 1): Attempt 2 sent. + fakeExecutor.advanceTime(Duration.ofMillis(60)); + waitForRequests(testPublisherServiceImpl, 2); + assertThat(testPublisherServiceImpl.getCapturedRequests()).hasSize(2); + + // T=120 (Hedge 2): Attempt 3 sent. + fakeExecutor.advanceTime(Duration.ofMillis(60)); + waitForRequests(testPublisherServiceImpl, 3); + assertThat(testPublisherServiceImpl.getCapturedRequests()).hasSize(3); + + // Advance to complete + fakeExecutor.advanceTime(Duration.ofMillis(100)); + assertEquals("1", future.get(5, TimeUnit.SECONDS)); + + shutdownTestPublisher(publisher); + } + + @Test + public void testHedgingBypassedIfNoTokens() throws Exception { + Publisher publisher = getPublisherWithHedge(Duration.ofMillis(50)); + + // Drain the token bucket completely (since it starts full) + while (publisher.getHedgeTokenBucket().tryAcquire()) {} + assertThat(publisher.getHedgeTokenBucket().getTokens()).isEqualTo(0.0f); + + testPublisherServiceImpl.setPublishResponseDelay(Duration.ofMillis(100)); + testPublisherServiceImpl.addPublishResponse(PublishResponse.newBuilder().addMessageIds("1")); + + ApiFuture future = sendTestMessage(publisher, "msg-slow-no-tokens"); + waitForRequests(testPublisherServiceImpl, 1); + + // Advance past hedge delay + fakeExecutor.advanceTime(Duration.ofMillis(60)); + + // Should NOT trigger hedge because token bucket is empty + assertThat(testPublisherServiceImpl.getCapturedRequests()).hasSize(1); + + fakeExecutor.advanceTime(Duration.ofMillis(50)); + assertEquals("1", future.get(5, TimeUnit.SECONDS)); + + shutdownTestPublisher(publisher); + } + + @Test + public void testHedgingCancellationPropagates() throws Exception { + Publisher publisher = getPublisherWithHedge(Duration.ofMillis(50)); + + testPublisherServiceImpl.setAutoPublishResponse(false); + testPublisherServiceImpl.setPublishResponseDelay(Duration.ofMillis(100)); + testPublisherServiceImpl.addPublishResponse(PublishResponse.newBuilder().addMessageIds("1")); + testPublisherServiceImpl.addPublishResponse(PublishResponse.newBuilder().addMessageIds("2")); + + ApiFuture future = sendTestMessage(publisher, "msg-cancel"); + waitForRequests(testPublisherServiceImpl, 1); + + // Trigger hedge + fakeExecutor.advanceTime(Duration.ofMillis(60)); + waitForRequests(testPublisherServiceImpl, 2); + assertThat(testPublisherServiceImpl.getCapturedRequests()).hasSize(2); + + // Cancel the future + future.cancel(true); + + // Verify cancellation propagates to overall future + assertTrue(future.isCancelled()); + + shutdownTestPublisher(publisher); + } + private Builder getTestPublisherBuilder() { return Publisher.newBuilder(TEST_TOPIC) .setExecutorProvider(FixedExecutorProvider.create(fakeExecutor))