From 61d0b608427a53832b47318b6d076bee32e11719 Mon Sep 17 00:00:00 2001 From: Tony Cui Date: Mon, 13 Jul 2026 15:14:59 +0000 Subject: [PATCH 01/22] Add hedgeSettings with configurable hedgeDelay --- .../google/cloud/pubsub/v1/HedgeSettings.java | 77 +++++++++++++++++++ .../cloud/pubsub/v1/HedgeSettingsTest.java | 68 ++++++++++++++++ 2 files changed, 145 insertions(+) create mode 100644 java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/HedgeSettings.java create mode 100644 java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/HedgeSettingsTest.java 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..9a2dbf66c9dd --- /dev/null +++ b/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/HedgeSettings.java @@ -0,0 +1,77 @@ +/* + * 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 { + private static final Duration DEFAULT_DELAY = Duration.ofMillis(100); + private static final int DEFAULT_MAX_TOKENS = 100; + private static final float DEFAULT_REFILL = 0.1f; + + private final Duration hedgeDelay; + private final int maxTokens; + private final float refill; + + private HedgeSettings(Builder builder) { + this.hedgeDelay = builder.hedgeDelay; + this.maxTokens = builder.maxTokens; + this.refill = builder.refill; + } + + /** Returns the configured hedging delay. */ + public Duration getHedgeDelay() { + return hedgeDelay; + } + + int getMaxTokens() { + return maxTokens; + } + + float getRefill() { + return refill; + } + + /** Returns a new builder for {@code HedgeSettings}. */ + public static Builder newBuilder() { + return new Builder(); + } + + /** Builder for {@code HedgeSettings}. */ + public static final class Builder { + private Duration hedgeDelay = DEFAULT_DELAY; + private int maxTokens = DEFAULT_MAX_TOKENS; + private float refill = DEFAULT_REFILL; + + private Builder() {} + + /** Allows hedging delay to be configurable. */ + public Builder setHedgeDelay(Duration hedgeDelay) { + Preconditions.checkNotNull(hedgeDelay); + Preconditions.checkArgument(!hedgeDelay.isNegative() && !hedgeDelay.isZero(), "hedgeDelay must be positive"); + this.hedgeDelay = hedgeDelay; + return this; + } + + /** Builds an instance of {@code HedgeSettings}. */ + public HedgeSettings build() { + return new HedgeSettings(this); + } + } +} 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..2a01408b7bdf --- /dev/null +++ b/java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/HedgeSettingsTest.java @@ -0,0 +1,68 @@ +/* + * 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)); + } +} From 19afd68374fbd8899e1f5ecd4ad3e15ee2a001d1 Mon Sep 17 00:00:00 2001 From: Tony Cui Date: Mon, 13 Jul 2026 16:08:47 +0000 Subject: [PATCH 02/22] Add HedgeTokenBucket class to allow for token operations --- .../cloud/pubsub/v1/HedgeTokenBucket.java | 65 +++++++++++++++++ .../cloud/pubsub/v1/HedgeTokenBucketTest.java | 71 +++++++++++++++++++ 2 files changed, 136 insertions(+) create mode 100644 java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/HedgeTokenBucket.java create mode 100644 java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/HedgeTokenBucketTest.java 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/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); + } +} From 49426a9b490abb6a6ecb7d0f78cf26e66f1a4816 Mon Sep 17 00:00:00 2001 From: Tony Cui Date: Mon, 13 Jul 2026 16:43:24 +0000 Subject: [PATCH 03/22] Add publisher integration with HedgeSettings --- .../com/google/cloud/pubsub/v1/Publisher.java | 22 +++++++++ .../cloud/pubsub/v1/PublisherImplTest.java | 48 ++++++++++++++----- 2 files changed, 59 insertions(+), 11 deletions(-) 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..1654cb364418 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 @@ -138,6 +138,9 @@ public class Publisher implements PublisherInterface { private final OpenTelemetry openTelemetry; private OpenTelemetryPubsubTracer tracer = new OpenTelemetryPubsubTracer(null, false); + private final HedgeSettings hedgeSettings; + private final HedgeTokenBucket hedgeTokenBucket; + /** The maximum number of messages in one request. Defined by the API. */ public static long getApiMaxRequestElementCount() { return 1000L; @@ -230,6 +233,9 @@ 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.publishContext = GrpcCallContext.createDefault(); this.publishContextWithCompression = GrpcCallContext.createDefault() @@ -246,6 +252,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. @@ -814,6 +829,7 @@ public PubsubMessage apply(PubsubMessage input) { private boolean enableOpenTelemetryTracing = false; private OpenTelemetry openTelemetry = null; + private HedgeSettings hedgeSettings = null; private Builder(String topic) { this.topicName = Preconditions.checkNotNull(topic); @@ -966,6 +982,12 @@ public Builder setOpenTelemetry(OpenTelemetry openTelemetry) { return this; } + /** Configures the Publisher's hedging parameters. */ + public Builder setHedgeSettings(HedgeSettings hedgeSettings) { + this.hedgeSettings = hedgeSettings; + return this; + } + /** Returns the default BatchingSettings used by the client if settings are not provided. */ public static BatchingSettings getDefaultBatchingSettings() { return DEFAULT_BATCHING_SETTINGS; 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..a5bd92c2aa0d 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 @@ -64,7 +64,6 @@ import org.easymock.EasyMock; import org.junit.After; import org.junit.Before; -import org.junit.Ignore; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -513,20 +512,21 @@ public void testEnableMessageOrdering_overwritesMaxAttempts() throws Exception { *
  • publish with key orderA, which should now succeed * */ - @Ignore("https://github.com/googleapis/google-cloud-java/issues/13394") + /* + Temporarily disabled due to https://github.com/googleapis/java-pubsub/issues/1861. + TODO(maitrimangal): Enable once resolved. @Test public void testResumePublish() throws Exception { Publisher publisher = getTestPublisherBuilder() .setBatchingSettings( - Publisher.Builder.DEFAULT_BATCHING_SETTINGS.toBuilder() + Publisher.Builder.DEFAULT_BATCHING_SETTINGS + .toBuilder() .setElementCountThreshold(2L) .build()) .setEnableMessageOrdering(true) .build(); - testPublisherServiceImpl.setExecutor(fakeExecutor); - ApiFuture future1 = sendTestMessageWithOrderingKey(publisher, "m1", "orderA"); ApiFuture future2 = sendTestMessageWithOrderingKey(publisher, "m2", "orderA"); @@ -536,6 +536,7 @@ public void testResumePublish() throws Exception { // This exception should stop future publishing to the same key testPublisherServiceImpl.addPublishError(new StatusException(Status.INVALID_ARGUMENT)); + fakeExecutor.advanceTime(Duration.ZERO); try { @@ -575,8 +576,8 @@ public void testResumePublish() throws Exception { testPublisherServiceImpl.addPublishResponse( PublishResponse.newBuilder().addMessageIds("5").addMessageIds("6")); - assertEquals("5", future5.get()); - assertEquals("6", future6.get()); + Assert.assertEquals("5", future5.get()); + Assert.assertEquals("6", future6.get()); // Resume publishing of "orderA", which should now succeed publisher.resumePublish("orderA"); @@ -587,19 +588,20 @@ public void testResumePublish() throws Exception { testPublisherServiceImpl.addPublishResponse( PublishResponse.newBuilder().addMessageIds("7").addMessageIds("8")); - assertEquals("7", future7.get()); - assertEquals("8", future8.get()); + Assert.assertEquals("7", future7.get()); + Assert.assertEquals("8", future8.get()); shutdownTestPublisher(publisher); } - @Ignore("https://github.com/googleapis/google-cloud-java/issues/13394") @Test public void testPublishThrowExceptionForUnsubmittedOrderingKeyMessage() throws Exception { Publisher publisher = getTestPublisherBuilder() + .setExecutorProvider(SINGLE_THREAD_EXECUTOR) .setBatchingSettings( - Publisher.Builder.DEFAULT_BATCHING_SETTINGS.toBuilder() + Publisher.Builder.DEFAULT_BATCHING_SETTINGS + .toBuilder() .setElementCountThreshold(2L) .setDelayThresholdDuration(Duration.ofSeconds(500)) .build()) @@ -642,6 +644,7 @@ public void testPublishThrowExceptionForUnsubmittedOrderingKeyMessage() throws E assertEquals(SequentialExecutorService.CallbackExecutor.CANCELLATION_EXCEPTION, e.getCause()); } } + */ private ApiFuture sendTestMessageWithOrderingKey( Publisher publisher, String data, String orderingKey) { @@ -1339,6 +1342,29 @@ 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 Builder getTestPublisherBuilder() { return Publisher.newBuilder(TEST_TOPIC) .setExecutorProvider(FixedExecutorProvider.create(fakeExecutor)) From 22417d156d31e10398f225c6f98ccebae9f93595 Mon Sep 17 00:00:00 2001 From: Tony Cui Date: Mon, 13 Jul 2026 16:51:32 +0000 Subject: [PATCH 04/22] Restore original @Ignore annotations and code in PublisherImplTest --- .../cloud/pubsub/v1/PublisherImplTest.java | 48 +++++-------------- 1 file changed, 11 insertions(+), 37 deletions(-) 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 a5bd92c2aa0d..8e6efaf372c9 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 @@ -64,6 +64,7 @@ import org.easymock.EasyMock; import org.junit.After; import org.junit.Before; +import org.junit.Ignore; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -512,21 +513,20 @@ public void testEnableMessageOrdering_overwritesMaxAttempts() throws Exception { *
  • publish with key orderA, which should now succeed * */ - /* - Temporarily disabled due to https://github.com/googleapis/java-pubsub/issues/1861. - TODO(maitrimangal): Enable once resolved. + @Ignore("https://github.com/googleapis/google-cloud-java/issues/13394") @Test public void testResumePublish() throws Exception { Publisher publisher = getTestPublisherBuilder() .setBatchingSettings( - Publisher.Builder.DEFAULT_BATCHING_SETTINGS - .toBuilder() + Publisher.Builder.DEFAULT_BATCHING_SETTINGS.toBuilder() .setElementCountThreshold(2L) .build()) .setEnableMessageOrdering(true) .build(); + testPublisherServiceImpl.setExecutor(fakeExecutor); + ApiFuture future1 = sendTestMessageWithOrderingKey(publisher, "m1", "orderA"); ApiFuture future2 = sendTestMessageWithOrderingKey(publisher, "m2", "orderA"); @@ -536,7 +536,6 @@ public void testResumePublish() throws Exception { // This exception should stop future publishing to the same key testPublisherServiceImpl.addPublishError(new StatusException(Status.INVALID_ARGUMENT)); - fakeExecutor.advanceTime(Duration.ZERO); try { @@ -576,8 +575,8 @@ public void testResumePublish() throws Exception { testPublisherServiceImpl.addPublishResponse( PublishResponse.newBuilder().addMessageIds("5").addMessageIds("6")); - Assert.assertEquals("5", future5.get()); - Assert.assertEquals("6", future6.get()); + assertEquals("5", future5.get()); + assertEquals("6", future6.get()); // Resume publishing of "orderA", which should now succeed publisher.resumePublish("orderA"); @@ -588,20 +587,19 @@ public void testResumePublish() throws Exception { testPublisherServiceImpl.addPublishResponse( PublishResponse.newBuilder().addMessageIds("7").addMessageIds("8")); - Assert.assertEquals("7", future7.get()); - Assert.assertEquals("8", future8.get()); + assertEquals("7", future7.get()); + assertEquals("8", future8.get()); shutdownTestPublisher(publisher); } + @Ignore("https://github.com/googleapis/google-cloud-java/issues/13394") @Test public void testPublishThrowExceptionForUnsubmittedOrderingKeyMessage() throws Exception { Publisher publisher = getTestPublisherBuilder() - .setExecutorProvider(SINGLE_THREAD_EXECUTOR) .setBatchingSettings( - Publisher.Builder.DEFAULT_BATCHING_SETTINGS - .toBuilder() + Publisher.Builder.DEFAULT_BATCHING_SETTINGS.toBuilder() .setElementCountThreshold(2L) .setDelayThresholdDuration(Duration.ofSeconds(500)) .build()) @@ -644,7 +642,6 @@ public void testPublishThrowExceptionForUnsubmittedOrderingKeyMessage() throws E assertEquals(SequentialExecutorService.CallbackExecutor.CANCELLATION_EXCEPTION, e.getCause()); } } - */ private ApiFuture sendTestMessageWithOrderingKey( Publisher publisher, String data, String orderingKey) { @@ -1342,29 +1339,6 @@ 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 Builder getTestPublisherBuilder() { return Publisher.newBuilder(TEST_TOPIC) .setExecutorProvider(FixedExecutorProvider.create(fakeExecutor)) From 91706a2a212828ba92fb22dc9aae51b67bc9f6a6 Mon Sep 17 00:00:00 2001 From: Tony Cui Date: Mon, 13 Jul 2026 17:39:56 +0000 Subject: [PATCH 05/22] style(pubsub): fix checkstyle violations in HedgeSettings --- .../google/cloud/pubsub/v1/HedgeSettings.java | 48 +++++++++++++++---- 1 file changed, 38 insertions(+), 10 deletions(-) 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 index 9a2dbf66c9dd..4641346fca31 100644 --- 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 @@ -21,21 +21,31 @@ /** 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(Builder builder) { + private HedgeSettings(final Builder builder) { this.hedgeDelay = builder.hedgeDelay; this.maxTokens = builder.maxTokens; this.refill = builder.refill; } - /** Returns the configured hedging delay. */ + /** + * Returns the configured hedging delay. + * + * @return the hedging delay. + */ public Duration getHedgeDelay() { return hedgeDelay; } @@ -48,28 +58,46 @@ float getRefill() { return refill; } - /** Returns a new builder for {@code HedgeSettings}. */ + /** + * 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() {} + private Builder() { } - /** Allows hedging delay to be configurable. */ - public Builder setHedgeDelay(Duration hedgeDelay) { - Preconditions.checkNotNull(hedgeDelay); - Preconditions.checkArgument(!hedgeDelay.isNegative() && !hedgeDelay.isZero(), "hedgeDelay must be positive"); - this.hedgeDelay = hedgeDelay; + /** + * 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); + Preconditions.checkArgument( + !delay.isNegative() && !delay.isZero(), + "hedgeDelay must be positive"); + this.hedgeDelay = delay; return this; } - /** Builds an instance of {@code HedgeSettings}. */ + /** + * Builds an instance of {@code HedgeSettings}. + * + * @return the built {@code HedgeSettings} instance. + */ public HedgeSettings build() { return new HedgeSettings(this); } From bd76ede662b5a951773c34cab4f4bd79b7c9373e Mon Sep 17 00:00:00 2001 From: Tony Cui Date: Mon, 13 Jul 2026 17:53:42 +0000 Subject: [PATCH 06/22] style(pubsub): fix formatting and checkstyle violations in HedgeSettings and tests --- .../google/cloud/pubsub/v1/HedgeSettings.java | 16 +++++++++---- .../cloud/pubsub/v1/HedgeSettingsTest.java | 4 +--- .../cloud/pubsub/v1/PublisherImplTest.java | 23 +++++++++++++++++++ 3 files changed, 36 insertions(+), 7 deletions(-) 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 index 4641346fca31..ffa6daed1d91 100644 --- 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 @@ -23,15 +23,19 @@ 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; @@ -71,12 +75,16 @@ public static Builder newBuilder() { 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() { } + private Builder() { + super(); + } /** * Allows hedging delay to be configurable. @@ -86,9 +94,9 @@ private Builder() { } */ public Builder setHedgeDelay(final Duration delay) { Preconditions.checkNotNull(delay); - Preconditions.checkArgument( - !delay.isNegative() && !delay.isZero(), - "hedgeDelay must be positive"); + if (delay.isNegative() || delay.isZero()) { + throw new IllegalArgumentException("delay must be positive"); + } this.hedgeDelay = delay; return this; } 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 index 2a01408b7bdf..f3b26aedb3c3 100644 --- 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 @@ -61,8 +61,6 @@ public void testZeroDelayThrows() { @Test public void testNullDelayThrows() { - assertThrows( - NullPointerException.class, - () -> HedgeSettings.newBuilder().setHedgeDelay(null)); + assertThrows(NullPointerException.class, () -> HedgeSettings.newBuilder().setHedgeDelay(null)); } } 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..95b19d6b0108 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 @@ -1339,6 +1339,29 @@ 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 Builder getTestPublisherBuilder() { return Publisher.newBuilder(TEST_TOPIC) .setExecutorProvider(FixedExecutorProvider.create(fakeExecutor)) From 227f808372e013995803d024d003efd0acd71a48 Mon Sep 17 00:00:00 2001 From: Tony Cui Date: Mon, 20 Jul 2026 15:17:04 +0000 Subject: [PATCH 07/22] Resolve git comments related to access modifiers and constraints --- .../main/java/com/google/cloud/pubsub/v1/HedgeSettings.java | 3 +-- .../src/main/java/com/google/cloud/pubsub/v1/Publisher.java | 3 +++ 2 files changed, 4 insertions(+), 2 deletions(-) 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 index ffa6daed1d91..62e16aa71639 100644 --- 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 @@ -50,7 +50,7 @@ private HedgeSettings(final Builder builder) { * * @return the hedging delay. */ - public Duration getHedgeDelay() { + Duration getHedgeDelay() { return hedgeDelay; } @@ -83,7 +83,6 @@ public static final class Builder { private float refill = DEFAULT_REFILL; private Builder() { - super(); } /** 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 1654cb364418..d982bc03a210 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 @@ -994,6 +994,9 @@ public static BatchingSettings getDefaultBatchingSettings() { } 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); } } From e29eab859ce0904461f5df96898ba23d3bd23910 Mon Sep 17 00:00:00 2001 From: Tony Cui Date: Mon, 20 Jul 2026 15:51:53 +0000 Subject: [PATCH 08/22] Add CancellationSharer and HedgedRequest helper classes --- .../cloud/pubsub/v1/CancellationSharer.java | 141 ++++++++++++++++++ .../google/cloud/pubsub/v1/HedgedRequest.java | 44 ++++++ 2 files changed, 185 insertions(+) create mode 100644 java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/CancellationSharer.java create mode 100644 java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/HedgedRequest.java 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..50995e88f489 --- /dev/null +++ b/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/CancellationSharer.java @@ -0,0 +1,141 @@ +/* + * 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 (CancellationSharer in the diagram). + */ +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/HedgedRequest.java b/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/HedgedRequest.java new file mode 100644 index 000000000000..ac405032c8e1 --- /dev/null +++ b/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/HedgedRequest.java @@ -0,0 +1,44 @@ +/* + * 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; + } +} From ed05c60650c1e8496124abff2a7c49986c8d8e48 Mon Sep 17 00:00:00 2001 From: Tony Cui Date: Mon, 20 Jul 2026 19:23:16 +0000 Subject: [PATCH 09/22] Integrate CancellationSharer and HedgedRequest into Publisher.java. Also add Unit tests --- .../cloud/pubsub/v1/CancellationSharer.java | 3 + .../com/google/cloud/pubsub/v1/Publisher.java | 126 ++++++++++++- .../pubsub/v1/FakePublisherServiceImpl.java | 3 +- .../cloud/pubsub/v1/PublisherImplTest.java | 172 ++++++++++++++++++ 4 files changed, 300 insertions(+), 4 deletions(-) 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 index 50995e88f489..37de9324eedf 100644 --- 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 @@ -28,6 +28,8 @@ import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; + + /** * Coordinates multiple publish attempts for a single batch of messages. * @@ -43,6 +45,7 @@ class CancellationSharer extends AbstractApiFuture { final AtomicBoolean isInQueue = new AtomicBoolean(false); private final AtomicReference lastError = new AtomicReference<>(); + CancellationSharer(Publisher.OutstandingBatch batch, Publisher publisher) { this.batch = batch; this.publisher = publisher; 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 d982bc03a210..b57442d5375f 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; @@ -70,7 +72,10 @@ import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.ScheduledFuture; import java.util.concurrent.TimeUnit; +import java.util.concurrent.ConcurrentLinkedQueue; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; import java.util.concurrent.locks.Lock; import java.util.concurrent.locks.ReentrantLock; import java.util.logging.Level; @@ -140,6 +145,12 @@ public class Publisher implements PublisherInterface { 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() { @@ -236,10 +247,15 @@ private Publisher(Builder builder) throws IOException { 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. */ @@ -587,7 +603,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 = @@ -603,7 +623,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; @@ -615,7 +641,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; } @@ -692,6 +718,9 @@ public void shutdown() { if (currentAlarmFuture != null && activeAlarm.getAndSet(false)) { currentAlarmFuture.cancel(false); } + if (queueProcessingFuture != null) { + queueProcessingFuture.cancel(false); + } publishAllOutstanding(); messagesWaiter.waitComplete(); backgroundResources.shutdown(); @@ -830,6 +859,7 @@ 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); @@ -988,6 +1018,11 @@ public Builder setHedgeSettings(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; @@ -1207,4 +1242,89 @@ && 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/PublisherImplTest.java b/java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/PublisherImplTest.java index 95b19d6b0108..2da52e27c7a2 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 @@ -1362,6 +1363,177 @@ public void testPublisherWithoutHedgeSettings() throws Exception { 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)) From 23b84322f8e958cf5882e3a4bef5e3ddac59d641 Mon Sep 17 00:00:00 2001 From: Tony Cui Date: Mon, 20 Jul 2026 19:30:45 +0000 Subject: [PATCH 10/22] fix lint --- .../cloud/pubsub/v1/CancellationSharer.java | 12 +++--- .../google/cloud/pubsub/v1/HedgeSettings.java | 3 +- .../google/cloud/pubsub/v1/HedgedRequest.java | 4 +- .../com/google/cloud/pubsub/v1/Publisher.java | 42 +++++++++++-------- .../cloud/pubsub/v1/PublisherImplTest.java | 9 ++-- 5 files changed, 35 insertions(+), 35 deletions(-) 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 index 37de9324eedf..702da3667b9a 100644 --- 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 @@ -28,24 +28,23 @@ 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 (CancellationSharer in the diagram). + * It manages the lifecycle of the original attempt and any subsequent hedged attempts + * (CancellationSharer in the diagram). */ class CancellationSharer extends AbstractApiFuture { private final Publisher.OutstandingBatch batch; private final Publisher publisher; - private final Map> runningAttempts = new ConcurrentHashMap<>(); + 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; @@ -108,7 +107,8 @@ 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")); + setException( + error != null ? error : new RuntimeException("Hedging failed with no active attempts")); } } } 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 index 62e16aa71639..fdbfebbe9106 100644 --- 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 @@ -82,8 +82,7 @@ public static final class Builder { /** Refill rate. */ private float refill = DEFAULT_REFILL; - private Builder() { - } + private Builder() {} /** * Allows hedging delay to be configurable. 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 index ac405032c8e1..ad07de9a8a90 100644 --- 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 @@ -16,9 +16,7 @@ package com.google.cloud.pubsub.v1; -/** - * Represents a pending hedging check in the publisher's queue. - */ +/** Represents a pending hedging check in the publisher's queue. */ class HedgedRequest { private final CancellationSharer coordinator; private final int attemptNumber; 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 b57442d5375f..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 @@ -67,15 +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.ConcurrentLinkedQueue; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; -import java.util.concurrent.atomic.AtomicLong; import java.util.concurrent.locks.Lock; import java.util.concurrent.locks.ReentrantLock; import java.util.logging.Level; @@ -1250,16 +1249,18 @@ private ApiFuture startHedgedCall(final OutstandingBatch outsta 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); + outstanding.publishResult.addListener( + new Runnable() { + @Override + public void run() { + if (outstanding.publishResult.isCancelled()) { + if (cancelledCount.incrementAndGet() == batchSize) { + coordinator.cancel(true); + } + } } - } - } - }, directExecutor()); + }, + directExecutor()); } ApiFuture firstAttemptFuture = publishCall(outstandingBatch); @@ -1284,12 +1285,16 @@ private void scheduleQueueProcessing() { long delay = nextItem.getSendAfterMs() - clock.millisTime(); delay = Math.max(0, delay); - queueProcessingFuture = executor.schedule(new Runnable() { - @Override - public void run() { - processQueue(); - } - }, delay, TimeUnit.MILLISECONDS); + queueProcessingFuture = + executor.schedule( + new Runnable() { + @Override + public void run() { + processQueue(); + } + }, + delay, + TimeUnit.MILLISECONDS); } } @@ -1311,7 +1316,8 @@ private void processQueue() { 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); + HedgedRequest nextItem = + new HedgedRequest(coordinator, item.getAttemptNumber() + 1, now + delayMs); hedgingQueue.add(nextItem); // Start Hedged Attempt 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 2da52e27c7a2..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 @@ -1364,9 +1364,7 @@ public void testPublisherWithoutHedgeSettings() throws Exception { } private Publisher getPublisherWithHedge(Duration delay) throws Exception { - HedgeSettings hedgeSettings = HedgeSettings.newBuilder() - .setHedgeDelay(delay) - .build(); + HedgeSettings hedgeSettings = HedgeSettings.newBuilder().setHedgeDelay(delay).build(); return getTestPublisherBuilder() .setHedgeSettings(hedgeSettings) .setClock(fakeExecutor.getClock()) @@ -1380,7 +1378,8 @@ private Publisher getPublisherWithHedge(Duration delay) throws Exception { private void waitForRequests(FakePublisherServiceImpl service, int expectedCount) throws InterruptedException { long timeout = System.currentTimeMillis() + 5000; - while (service.getCapturedRequests().size() < expectedCount && System.currentTimeMillis() < timeout) { + while (service.getCapturedRequests().size() < expectedCount + && System.currentTimeMillis() < timeout) { Thread.sleep(5); } if (service.getCapturedRequests().size() < expectedCount) { @@ -1506,8 +1505,6 @@ public void testHedgingBypassedIfNoTokens() throws Exception { shutdownTestPublisher(publisher); } - - @Test public void testHedgingCancellationPropagates() throws Exception { Publisher publisher = getPublisherWithHedge(Duration.ofMillis(50)); From 3a23d607392c87ea301b87ad188a9683fc227ef0 Mon Sep 17 00:00:00 2001 From: Tony Cui Date: Mon, 20 Jul 2026 21:14:16 +0000 Subject: [PATCH 11/22] Modifying comment line --- .../java/com/google/cloud/pubsub/v1/CancellationSharer.java | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) 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 index 702da3667b9a..571c8a99d8ac 100644 --- 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 @@ -32,8 +32,7 @@ * 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 - * (CancellationSharer in the diagram). + * It manages the lifecycle of the original attempt and any subsequent hedged attempts. */ class CancellationSharer extends AbstractApiFuture { private final Publisher.OutstandingBatch batch; From da4307e697c3873f5695c1158fe00c485ecd6a2c Mon Sep 17 00:00:00 2001 From: Tony Cui Date: Mon, 27 Jul 2026 14:40:55 +0000 Subject: [PATCH 12/22] Modified default values, added configurable settings to refill ratio and maxtokens, updated tests. --- .../google/cloud/pubsub/v1/HedgeSettings.java | 54 ++++++++++++---- .../cloud/pubsub/v1/HedgeTokenBucket.java | 2 +- .../cloud/pubsub/v1/HedgeSettingsTest.java | 62 ++++++++++++++++--- 3 files changed, 98 insertions(+), 20 deletions(-) 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 index fdbfebbe9106..1d509c65f33b 100644 --- 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 @@ -22,13 +22,13 @@ /** Settings for configuring publish hedging. */ public final class HedgeSettings { /** Default hedging delay. */ - private static final Duration DEFAULT_DELAY = Duration.ofMillis(100); + private static final Duration DEFAULT_DELAY = Duration.ofMillis(1000); /** Default maximum number of tokens in the bucket. */ - private static final int DEFAULT_MAX_TOKENS = 100; + private static final int DEFAULT_MAX_TOKENS = 50; /** Default refill rate (tokens per successful request). */ - private static final float DEFAULT_REFILL = 0.1f; + private static final float DEFAULT_REFILL_RATIO = 0.1f; /** Hedging delay. */ private final Duration hedgeDelay; @@ -37,12 +37,12 @@ public final class HedgeSettings { private final int maxTokens; /** Refill rate. */ - private final float refill; + private final float refillRatio; private HedgeSettings(final Builder builder) { this.hedgeDelay = builder.hedgeDelay; this.maxTokens = builder.maxTokens; - this.refill = builder.refill; + this.refillRatio = builder.refillRatio; } /** @@ -58,8 +58,8 @@ int getMaxTokens() { return maxTokens; } - float getRefill() { - return refill; + float getRefillRatio() { + return refillRatio; } /** @@ -80,25 +80,57 @@ public static final class Builder { private int maxTokens = DEFAULT_MAX_TOKENS; /** Refill rate. */ - private float refill = DEFAULT_REFILL; + private float refillRatio = DEFAULT_REFILL_RATIO; private Builder() {} /** * Allows hedging delay to be configurable. * - * @param delay the hedging delay, must be positive. + * @param delay the hedging delay, must be 0.1s <= HedgeDelay <= 10s. * @return this builder. */ public Builder setHedgeDelay(final Duration delay) { Preconditions.checkNotNull(delay); - if (delay.isNegative() || delay.isZero()) { - throw new IllegalArgumentException("delay must be positive"); + if (delay.toMillis() < 100 || delay.toMillis() > 10000) { + throw new IllegalArgumentException( + "delay must be greater than or equal to 100ms and less than or equal to 10s"); } this.hedgeDelay = delay; return this; } + /** + * Allows the maximum number of tokens in the bucket to be configurable. + * + * @param maxTokens the maximum number of tokens, must be 0 < MaxTokens <= 250. + * @return this builder. + */ + public Builder setMaxTokens(final int maxTokens) { + if (maxTokens <= 0 || maxTokens > 250) { + throw new IllegalArgumentException( + "maxTokens must be greater than 0 and less than or equal to 250"); + } + this.maxTokens = maxTokens; + return this; + } + + /** + * Allows the token bucket refill rate to be configurable. + * + * @param refill the refill rate (tokens per successful request), must be 0 < RefillRatio <= + * 0.2. + * @return this builder. + */ + public Builder setRefillRatio(final float refillRatio) { + if (refillRatio <= 0.0f || refillRatio > 0.2f) { + throw new IllegalArgumentException( + "refill must be greater than 0.0 and less than or equal to 0.2"); + } + this.refillRatio = refillRatio; + return this; + } + /** * Builds an instance of {@code HedgeSettings}. * 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 index 5378478cb429..3adc47f931f6 100644 --- 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 @@ -29,7 +29,7 @@ final class HedgeTokenBucket { HedgeTokenBucket(HedgeSettings settings) { this.maxTokens = settings.getMaxTokens(); - this.refillAmount = settings.getRefill(); + this.refillAmount = settings.getRefillRatio(); this.tokens = maxTokens; } 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 index f3b26aedb3c3..8ca202ed30ad 100644 --- 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 @@ -32,35 +32,81 @@ public class HedgeSettingsTest { 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); + assertEquals(Duration.ofMillis(1000), settings.getHedgeDelay()); + assertEquals(50, settings.getMaxTokens()); + assertEquals(0.1f, settings.getRefillRatio(), 0.0001f); } @Test public void testCustomDelay() { - Duration customDelay = Duration.ofMillis(50); + Duration customDelay = Duration.ofMillis(200); HedgeSettings settings = HedgeSettings.newBuilder().setHedgeDelay(customDelay).build(); assertNotNull(settings); assertEquals(customDelay, settings.getHedgeDelay()); } @Test - public void testNegativeDelayThrows() { + public void testDelayTooSmallThrows() { assertThrows( IllegalArgumentException.class, - () -> HedgeSettings.newBuilder().setHedgeDelay(Duration.ofMillis(-10))); + () -> HedgeSettings.newBuilder().setHedgeDelay(Duration.ofMillis(99))); } @Test - public void testZeroDelayThrows() { + public void testDelayTooLargeThrows() { assertThrows( IllegalArgumentException.class, - () -> HedgeSettings.newBuilder().setHedgeDelay(Duration.ZERO)); + () -> HedgeSettings.newBuilder().setHedgeDelay(Duration.ofMillis(10001))); } @Test public void testNullDelayThrows() { assertThrows(NullPointerException.class, () -> HedgeSettings.newBuilder().setHedgeDelay(null)); } + + @Test + public void testCustomMaxTokens() { + HedgeSettings settings = HedgeSettings.newBuilder().setMaxTokens(10).build(); + assertEquals(10, settings.getMaxTokens()); + } + + @Test + public void testNegativeMaxTokensThrows() { + assertThrows(IllegalArgumentException.class, () -> HedgeSettings.newBuilder().setMaxTokens(-5)); + } + + @Test + public void testZeroMaxTokensThrows() { + assertThrows(IllegalArgumentException.class, () -> HedgeSettings.newBuilder().setMaxTokens(0)); + } + + @Test + public void testMaxTokensTooLargeThrows() { + assertThrows( + IllegalArgumentException.class, () -> HedgeSettings.newBuilder().setMaxTokens(251)); + } + + @Test + public void testCustomRefill() { + HedgeSettings settings = HedgeSettings.newBuilder().setRefillRatio(0.15f).build(); + assertEquals(0.15f, settings.getRefillRatio(), 0.0001f); + } + + @Test + public void testNegativeRefillThrows() { + assertThrows( + IllegalArgumentException.class, () -> HedgeSettings.newBuilder().setRefillRatio(-0.1f)); + } + + @Test + public void testZeroRefillThrows() { + assertThrows( + IllegalArgumentException.class, () -> HedgeSettings.newBuilder().setRefillRatio(0.0f)); + } + + @Test + public void testRefillTooLargeThrows() { + assertThrows( + IllegalArgumentException.class, () -> HedgeSettings.newBuilder().setRefillRatio(0.21f)); + } } From 336c8091e4da63d890627f50d49367cb698512b7 Mon Sep 17 00:00:00 2001 From: Tony Cui Date: Mon, 27 Jul 2026 17:30:44 +0000 Subject: [PATCH 13/22] Removed HedgeTokenBucket.java class in favor of using an atomicinteger for hedgeTokenBucket. Also added HEDGE_TOKEN_SCALE. --- .../google/cloud/pubsub/v1/HedgeSettings.java | 4 +- .../cloud/pubsub/v1/HedgeTokenBucket.java | 65 ----------------- .../com/google/cloud/pubsub/v1/Publisher.java | 51 ++++++++++--- .../cloud/pubsub/v1/HedgeTokenBucketTest.java | 71 ------------------- .../cloud/pubsub/v1/PublisherImplTest.java | 66 ++++++++--------- 5 files changed, 78 insertions(+), 179 deletions(-) delete mode 100644 java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/HedgeTokenBucket.java delete mode 100644 java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/HedgeTokenBucketTest.java 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 index 1d509c65f33b..47efec2e42ae 100644 --- 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 @@ -94,7 +94,7 @@ public Builder setHedgeDelay(final Duration delay) { Preconditions.checkNotNull(delay); if (delay.toMillis() < 100 || delay.toMillis() > 10000) { throw new IllegalArgumentException( - "delay must be greater than or equal to 100ms and less than or equal to 10s"); + "hedgeDelay must be greater than or equal to 100ms and less than or equal to 10s"); } this.hedgeDelay = delay; return this; @@ -125,7 +125,7 @@ public Builder setMaxTokens(final int maxTokens) { public Builder setRefillRatio(final float refillRatio) { if (refillRatio <= 0.0f || refillRatio > 0.2f) { throw new IllegalArgumentException( - "refill must be greater than 0.0 and less than or equal to 0.2"); + "refillRatio must be greater than 0.0 and less than or equal to 0.2"); } this.refillRatio = refillRatio; return 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 deleted file mode 100644 index 3adc47f931f6..000000000000 --- a/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/HedgeTokenBucket.java +++ /dev/null @@ -1,65 +0,0 @@ -/* - * 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.getRefillRatio(); - 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/Publisher.java b/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/Publisher.java index 35ab918e5645..f2986402a3a7 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 @@ -143,7 +143,10 @@ public class Publisher implements PublisherInterface { private OpenTelemetryPubsubTracer tracer = new OpenTelemetryPubsubTracer(null, false); private final HedgeSettings hedgeSettings; - private final HedgeTokenBucket hedgeTokenBucket; + private static final int HEDGE_TOKEN_SCALE = 100; + private final AtomicInteger hedgingTokenBucket = new AtomicInteger(); + private int scaledMaxHedgeTokens; + private int scaledHedgeRefillAmount; private final ApiClock clock; private final ConcurrentLinkedQueue hedgingQueue; @@ -244,8 +247,12 @@ private Publisher(Builder builder) throws IOException { shutdown = new AtomicBoolean(false); messagesWaiter = new Waiter(); this.hedgeSettings = builder.hedgeSettings; - this.hedgeTokenBucket = - this.hedgeSettings != null ? new HedgeTokenBucket(this.hedgeSettings) : null; + if (this.hedgeSettings != null) { + this.scaledMaxHedgeTokens = this.hedgeSettings.getMaxTokens() * HEDGE_TOKEN_SCALE; + this.scaledHedgeRefillAmount = + (int) (this.hedgeSettings.getRefillRatio() * HEDGE_TOKEN_SCALE); + this.hedgingTokenBucket.set(scaledMaxHedgeTokens); + } this.clock = builder.clock != null ? builder.clock : CurrentMillisClock.getDefaultClock(); this.publishContext = GrpcCallContext.createDefault(); this.publishContextWithCompression = @@ -272,8 +279,11 @@ public HedgeSettings getHedgeSettings() { return hedgeSettings; } - HedgeTokenBucket getHedgeTokenBucket() { - return hedgeTokenBucket; + Float getHedgeTokenBalance() { + if (hedgeSettings == null) { + return null; + } + return (float) hedgingTokenBucket.get() / HEDGE_TOKEN_SCALE; } /** @@ -623,8 +633,33 @@ public ApiFuture call() { } void refillTokenBucket() { - if (hedgeTokenBucket != null) { - hedgeTokenBucket.refill(); + if (hedgeSettings != null) { + while (true) { + int current = hedgingTokenBucket.get(); + if (current >= scaledMaxHedgeTokens) { + return; + } + int next = Math.min(scaledMaxHedgeTokens, current + scaledHedgeRefillAmount); + if (hedgingTokenBucket.compareAndSet(current, next)) { + return; + } + } + } + } + + boolean tryAcquireHedgeToken() { + if (hedgeSettings == null) { + return false; + } + while (true) { + int current = hedgingTokenBucket.get(); + if (current < HEDGE_TOKEN_SCALE) { + return false; + } + int next = current - HEDGE_TOKEN_SCALE; + if (hedgingTokenBucket.compareAndSet(current, next)) { + return true; + } } } @@ -1313,7 +1348,7 @@ private void processQueue() { continue; } - if (hedgeTokenBucket.tryAcquire()) { + if (tryAcquireHedgeToken()) { // Clone and schedule next attempt check (Attempt + 1) long delayMs = hedgeSettings.getHedgeDelay().toMillis(); HedgedRequest nextItem = 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 deleted file mode 100644 index e92ecd9a6ce2..000000000000 --- a/java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/HedgeTokenBucketTest.java +++ /dev/null @@ -1,71 +0,0 @@ -/* - * 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 6acfd76a5870..f0e17ec83afd 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 @@ -1343,12 +1343,12 @@ public void testPublishOpenTelemetryTracing() throws Exception { @Test public void testPublisherWithHedgeSettings() throws Exception { HedgeSettings hedgeSettings = - HedgeSettings.newBuilder().setHedgeDelay(Duration.ofMillis(50)).build(); + HedgeSettings.newBuilder().setHedgeDelay(Duration.ofMillis(100)).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); + assertThat(publisher.getHedgeTokenBalance()).isNotNull(); + assertThat(publisher.getHedgeTokenBalance()).isWithin(0.0001f).of(50.0f); shutdownTestPublisher(publisher); } @@ -1358,7 +1358,7 @@ public void testPublisherWithoutHedgeSettings() throws Exception { Publisher publisher = getTestPublisherBuilder().build(); assertThat(publisher.getHedgeSettings()).isNull(); - assertThat(publisher.getHedgeTokenBucket()).isNull(); + assertThat(publisher.getHedgeTokenBalance()).isNull(); shutdownTestPublisher(publisher); } @@ -1392,7 +1392,7 @@ private void waitForRequests(FakePublisherServiceImpl service, int expectedCount @Test public void testHedgingNotTriggeredIfFast() throws Exception { - Publisher publisher = getPublisherWithHedge(Duration.ofMillis(50)); + Publisher publisher = getPublisherWithHedge(Duration.ofMillis(100)); // Prepare fast response (10ms delay) testPublisherServiceImpl.setAutoPublishResponse(false); @@ -1402,8 +1402,8 @@ public void testHedgingNotTriggeredIfFast() throws Exception { 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)); + // Advance time past response but before hedge delay (e.g. 50ms) + fakeExecutor.advanceTime(Duration.ofMillis(50)); // Future should be completed assertEquals("1", future.get()); @@ -1416,11 +1416,11 @@ public void testHedgingNotTriggeredIfFast() throws Exception { @Test public void testHedgingTriggeredIfSlow() throws Exception { - Publisher publisher = getPublisherWithHedge(Duration.ofMillis(50)); + Publisher publisher = getPublisherWithHedge(Duration.ofMillis(100)); - // Set response delay to 100ms (greater than 50ms hedge delay) + // Set response delay to 200ms (greater than 100ms hedge delay) testPublisherServiceImpl.setAutoPublishResponse(false); - testPublisherServiceImpl.setPublishResponseDelay(Duration.ofMillis(100)); + testPublisherServiceImpl.setPublishResponseDelay(Duration.ofMillis(200)); // Add two responses (one for main, one for hedge) testPublisherServiceImpl.addPublishResponse(PublishResponse.newBuilder().addMessageIds("1")); testPublisherServiceImpl.addPublishResponse(PublishResponse.newBuilder().addMessageIds("2")); @@ -1428,19 +1428,19 @@ public void testHedgingTriggeredIfSlow() throws Exception { ApiFuture future = sendTestMessage(publisher, "msg-slow"); waitForRequests(testPublisherServiceImpl, 1); - // Advance time to 40ms (before hedge delay) - fakeExecutor.advanceTime(Duration.ofMillis(40)); + // Advance time to 80ms (before hedge delay) + fakeExecutor.advanceTime(Duration.ofMillis(80)); assertThat(testPublisherServiceImpl.getCapturedRequests()).hasSize(1); // Only original sent - // Advance time to 60ms (past 50ms hedge delay) - fakeExecutor.advanceTime(Duration.ofMillis(20)); + // Advance time to 120ms (past 100ms hedge delay) + fakeExecutor.advanceTime(Duration.ofMillis(40)); 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)); + // Advance to 220ms to let responses complete + fakeExecutor.advanceTime(Duration.ofMillis(100)); assertTrue(future.isDone()); shutdownTestPublisher(publisher); @@ -1448,11 +1448,11 @@ public void testHedgingTriggeredIfSlow() throws Exception { @Test public void testMultipleHedging() throws Exception { - Publisher publisher = getPublisherWithHedge(Duration.ofMillis(50)); + Publisher publisher = getPublisherWithHedge(Duration.ofMillis(100)); - // Set delay to 200ms + // Set delay to 400ms testPublisherServiceImpl.setAutoPublishResponse(false); - testPublisherServiceImpl.setPublishResponseDelay(Duration.ofMillis(200)); + testPublisherServiceImpl.setPublishResponseDelay(Duration.ofMillis(400)); // Add responses for 3 attempts testPublisherServiceImpl.addPublishResponse(PublishResponse.newBuilder().addMessageIds("1")); testPublisherServiceImpl.addPublishResponse(PublishResponse.newBuilder().addMessageIds("2")); @@ -1462,18 +1462,18 @@ public void testMultipleHedging() throws Exception { waitForRequests(testPublisherServiceImpl, 1); // T=0: Attempt 1 sent. - // T=60 (Hedge 1): Attempt 2 sent. - fakeExecutor.advanceTime(Duration.ofMillis(60)); + // T=120 (Hedge 1): Attempt 2 sent. + fakeExecutor.advanceTime(Duration.ofMillis(120)); waitForRequests(testPublisherServiceImpl, 2); assertThat(testPublisherServiceImpl.getCapturedRequests()).hasSize(2); - // T=120 (Hedge 2): Attempt 3 sent. - fakeExecutor.advanceTime(Duration.ofMillis(60)); + // T=240 (Hedge 2): Attempt 3 sent. + fakeExecutor.advanceTime(Duration.ofMillis(120)); waitForRequests(testPublisherServiceImpl, 3); assertThat(testPublisherServiceImpl.getCapturedRequests()).hasSize(3); // Advance to complete - fakeExecutor.advanceTime(Duration.ofMillis(100)); + fakeExecutor.advanceTime(Duration.ofMillis(200)); assertEquals("1", future.get(5, TimeUnit.SECONDS)); shutdownTestPublisher(publisher); @@ -1481,25 +1481,25 @@ public void testMultipleHedging() throws Exception { @Test public void testHedgingBypassedIfNoTokens() throws Exception { - Publisher publisher = getPublisherWithHedge(Duration.ofMillis(50)); + Publisher publisher = getPublisherWithHedge(Duration.ofMillis(100)); // Drain the token bucket completely (since it starts full) - while (publisher.getHedgeTokenBucket().tryAcquire()) {} - assertThat(publisher.getHedgeTokenBucket().getTokens()).isEqualTo(0.0f); + while (publisher.tryAcquireHedgeToken()) {} + assertThat(publisher.getHedgeTokenBalance()).isEqualTo(0.0f); - testPublisherServiceImpl.setPublishResponseDelay(Duration.ofMillis(100)); + testPublisherServiceImpl.setPublishResponseDelay(Duration.ofMillis(200)); 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)); + fakeExecutor.advanceTime(Duration.ofMillis(120)); // Should NOT trigger hedge because token bucket is empty assertThat(testPublisherServiceImpl.getCapturedRequests()).hasSize(1); - fakeExecutor.advanceTime(Duration.ofMillis(50)); + fakeExecutor.advanceTime(Duration.ofMillis(100)); assertEquals("1", future.get(5, TimeUnit.SECONDS)); shutdownTestPublisher(publisher); @@ -1507,10 +1507,10 @@ public void testHedgingBypassedIfNoTokens() throws Exception { @Test public void testHedgingCancellationPropagates() throws Exception { - Publisher publisher = getPublisherWithHedge(Duration.ofMillis(50)); + Publisher publisher = getPublisherWithHedge(Duration.ofMillis(100)); testPublisherServiceImpl.setAutoPublishResponse(false); - testPublisherServiceImpl.setPublishResponseDelay(Duration.ofMillis(100)); + testPublisherServiceImpl.setPublishResponseDelay(Duration.ofMillis(200)); testPublisherServiceImpl.addPublishResponse(PublishResponse.newBuilder().addMessageIds("1")); testPublisherServiceImpl.addPublishResponse(PublishResponse.newBuilder().addMessageIds("2")); @@ -1518,7 +1518,7 @@ public void testHedgingCancellationPropagates() throws Exception { waitForRequests(testPublisherServiceImpl, 1); // Trigger hedge - fakeExecutor.advanceTime(Duration.ofMillis(60)); + fakeExecutor.advanceTime(Duration.ofMillis(120)); waitForRequests(testPublisherServiceImpl, 2); assertThat(testPublisherServiceImpl.getCapturedRequests()).hasSize(2); From 59a7a09c412f4cbc0d96500c32b65c8fe501f897 Mon Sep 17 00:00:00 2001 From: Tony Cui Date: Mon, 27 Jul 2026 17:41:06 +0000 Subject: [PATCH 14/22] Add documentation for scaling and remove unused getAttemptCount() in CancellationSharer.java --- .../com/google/cloud/pubsub/v1/CancellationSharer.java | 7 ------- .../main/java/com/google/cloud/pubsub/v1/Publisher.java | 9 +++++++++ 2 files changed, 9 insertions(+), 7 deletions(-) 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 index 571c8a99d8ac..b52e3f4f660f 100644 --- 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 @@ -25,7 +25,6 @@ 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; /** @@ -39,7 +38,6 @@ class CancellationSharer extends AbstractApiFuture { 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<>(); @@ -57,7 +55,6 @@ class CancellationSharer extends AbstractApiFuture { */ void addAttempt(final int attemptNumber, ApiFuture future) { runningAttempts.put(attemptNumber, future); - totalAttemptsCount.incrementAndGet(); if (done.get()) { future.cancel(true); @@ -133,10 +130,6 @@ public boolean cancel(boolean mayInterruptIfRunning) { 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/Publisher.java b/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/Publisher.java index f2986402a3a7..2b251d71a158 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 @@ -143,7 +143,14 @@ public class Publisher implements PublisherInterface { private OpenTelemetryPubsubTracer tracer = new OpenTelemetryPubsubTracer(null, false); private final HedgeSettings hedgeSettings; + + /** + * 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. + */ private static final int HEDGE_TOKEN_SCALE = 100; + private final AtomicInteger hedgingTokenBucket = new AtomicInteger(); private int scaledMaxHedgeTokens; private int scaledHedgeRefillAmount; @@ -248,6 +255,8 @@ private Publisher(Builder builder) throws IOException { messagesWaiter = new Waiter(); this.hedgeSettings = builder.hedgeSettings; if (this.hedgeSettings != null) { + // Scale all configuration values to integers to allow lock-free atomic operations + // without the garbage collection overhead of autoboxing Float objects. this.scaledMaxHedgeTokens = this.hedgeSettings.getMaxTokens() * HEDGE_TOKEN_SCALE; this.scaledHedgeRefillAmount = (int) (this.hedgeSettings.getRefillRatio() * HEDGE_TOKEN_SCALE); From 79411a8476051f9cf92c7168b3349ce120d11332 Mon Sep 17 00:00:00 2001 From: Tony Cui Date: Mon, 27 Jul 2026 19:23:43 +0000 Subject: [PATCH 15/22] Added method overloading to handle OTel span for hedged publish events --- .../pubsub/v1/OpenTelemetryPubsubTracer.java | 7 ++- .../com/google/cloud/pubsub/v1/Publisher.java | 12 ++++- .../cloud/pubsub/v1/PubsubMessageWrapper.java | 15 +++++- .../cloud/pubsub/v1/OpenTelemetryTest.java | 50 +++++++++++++++++++ 4 files changed, 80 insertions(+), 4 deletions(-) diff --git a/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/OpenTelemetryPubsubTracer.java b/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/OpenTelemetryPubsubTracer.java index 3de4484586d0..a811e30e1452 100644 --- a/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/OpenTelemetryPubsubTracer.java +++ b/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/OpenTelemetryPubsubTracer.java @@ -179,6 +179,11 @@ void endPublishBatchingSpan(PubsubMessageWrapper message) { * links with the publisher parent span are created for sampled messages in the batch. */ Span startPublishRpcSpan(TopicName topicName, List messages) { + return startPublishRpcSpan(topicName, messages, 1); + } + + Span startPublishRpcSpan( + TopicName topicName, List messages, int attemptNumber) { if (!enabled) { return null; } @@ -203,7 +208,7 @@ Span startPublishRpcSpan(TopicName topicName, List message for (PubsubMessageWrapper message : messages) { if (publishRpcSpan.getSpanContext().isSampled()) { message.getPublisherSpan().addLink(publishRpcSpan.getSpanContext(), linkAttributes); - message.addPublishStartEvent(); + message.addPublishStartEvent(attemptNumber); } } return publishRpcSpan; 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 2b251d71a158..92fe52038a48 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 @@ -530,6 +530,11 @@ private void publishAllWithoutInflightForKey(final String orderingKey) { } private ApiFuture publishCall(OutstandingBatch outstandingBatch) { + return publishCall(outstandingBatch, 1); + } + + private ApiFuture publishCall( + OutstandingBatch outstandingBatch, int attemptNumber) { GrpcCallContext context = publishContext; if (enableCompression && outstandingBatch.batchSizeBytes >= compressionBytesThreshold) { context = publishContextWithCompression; @@ -543,7 +548,9 @@ private ApiFuture publishCall(OutstandingBatch outstandingBatch pubsubMessagesList.add(messageWrapper.getPubsubMessage()); } - outstandingBatch.publishRpcSpan = tracer.startPublishRpcSpan(topicNameObject, messageWrappers); + outstandingBatch.publishRpcSpan = + tracer.startPublishRpcSpan(top + icNameObject, messageWrappers, attemptNumber); return publisherStub .publishCallable() @@ -1365,7 +1372,8 @@ private void processQueue() { hedgingQueue.add(nextItem); // Start Hedged Attempt - ApiFuture hedgedFuture = publishCall(coordinator.getBatch()); + ApiFuture hedgedFuture = + publishCall(coordinator.getBatch(), item.getAttemptNumber()); coordinator.addAttempt(item.getAttemptNumber(), hedgedFuture); } else { coordinator.isInQueue.set(false); diff --git a/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/PubsubMessageWrapper.java b/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/PubsubMessageWrapper.java index 19864a26f5a1..91009aa135d5 100644 --- a/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/PubsubMessageWrapper.java +++ b/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/PubsubMessageWrapper.java @@ -42,6 +42,7 @@ public class PubsubMessageWrapper { private final int deliveryAttempt; private static final String PUBLISH_START_EVENT = "publish start"; + private static final String HEDGED_PUBLISH_START_EVENT = "publish start (hedged)"; private static final String PUBLISH_END_EVENT = "publish end"; private static final String MODACK_START_EVENT = "modack start"; @@ -183,8 +184,20 @@ void setSubscribeProcessSpan(Span span) { /** Creates a publish start event that is tied to the publish RPC span time. */ void addPublishStartEvent() { + addPublishStartEvent(1); + } + + /** + * Creates a publish start event that is tied to the publish RPC span time, marking hedged + * attempts explicitly. + */ + void addPublishStartEvent(int attemptNumber) { if (publisherSpan != null) { - publisherSpan.addEvent(PUBLISH_START_EVENT); + if (attemptNumber > 1) { + publisherSpan.addEvent(HEDGED_PUBLISH_START_EVENT); + } else { + publisherSpan.addEvent(PUBLISH_START_EVENT); + } } } diff --git a/java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/OpenTelemetryTest.java b/java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/OpenTelemetryTest.java index 52351ddef466..af3d70d79990 100644 --- a/java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/OpenTelemetryTest.java +++ b/java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/OpenTelemetryTest.java @@ -57,6 +57,7 @@ public class OpenTelemetryTest { private static final String PUBLISH_BATCHING_SPAN_NAME = "publisher batching"; private static final String PUBLISH_RPC_SPAN_NAME = FULL_TOPIC_NAME.getTopic() + " publish"; private static final String PUBLISH_START_EVENT = "publish start"; + private static final String HEDGED_PUBLISH_START_EVENT = "publish start (hedged)"; private static final String PUBLISH_END_EVENT = "publish end"; private static final String SUBSCRIBER_SPAN_NAME = @@ -656,6 +657,55 @@ public void testSubscribeRpcSpanFailures() { .hasEnded(); } + @Test + public void testHedgedPublishSpanEvents() { + PubsubMessage message = getPubsubMessage(); + PubsubMessageWrapper messageWrapper = + PubsubMessageWrapper.newBuilder(message, FULL_TOPIC_NAME).build(); + List messageWrappers = + java.util.Collections.singletonList(messageWrapper); + + Tracer openTelemetryTracer = openTelemetryTesting.getOpenTelemetry().getTracer("test"); + OpenTelemetryPubsubTracer tracer = new OpenTelemetryPubsubTracer(openTelemetryTracer, true); + + // Start Publisher span + tracer.startPublisherSpan(messageWrapper); + + // Original Attempt 1 + Span publishRpcSpan1 = tracer.startPublishRpcSpan(FULL_TOPIC_NAME, messageWrappers, 1); + tracer.endPublishRpcSpan(publishRpcSpan1); + + // Hedged Attempt 2 + Span publishRpcSpan2 = tracer.startPublishRpcSpan(FULL_TOPIC_NAME, messageWrappers, 2); + tracer.endPublishRpcSpan(publishRpcSpan2); + + // End Publisher span + tracer.endPublisherSpan(messageWrapper); + + List allSpans = openTelemetryTesting.getSpans(); + // 3 Spans: publishRpcSpan1, publishRpcSpan2, publisherSpan + assertEquals(3, allSpans.size()); + SpanData publisherSpanData = allSpans.get(2); + + // The publisher parent span should have 3 events: + // 1. "publish start" (from attempt 1) + // 2. "publish start (hedged)" (from attempt 2) + // 3. "publish end" (when publisher span ends) + assertEquals(3, publisherSpanData.getEvents().size()); + + EventDataAssert startEvent1Assert = + OpenTelemetryAssertions.assertThat(publisherSpanData.getEvents().get(0)); + startEvent1Assert.hasName(PUBLISH_START_EVENT); + + EventDataAssert startEvent2Assert = + OpenTelemetryAssertions.assertThat(publisherSpanData.getEvents().get(1)); + startEvent2Assert.hasName(HEDGED_PUBLISH_START_EVENT); + + EventDataAssert endEventAssert = + OpenTelemetryAssertions.assertThat(publisherSpanData.getEvents().get(2)); + endEventAssert.hasName(PUBLISH_END_EVENT); + } + private PubsubMessage getPubsubMessage() { return PubsubMessage.newBuilder() .setData(ByteString.copyFromUtf8("test-data")) From e545cd6e894a07addb89ba8ffe5db4ea6445ef18 Mon Sep 17 00:00:00 2001 From: Tony Cui Date: Mon, 27 Jul 2026 19:38:48 +0000 Subject: [PATCH 16/22] Verify that the per-RPC timeout is strictly greater than the hedging delay --- .../com/google/cloud/pubsub/v1/Publisher.java | 16 ++++-- .../cloud/pubsub/v1/PublisherImplTest.java | 49 +++++++++++++++++++ 2 files changed, 61 insertions(+), 4 deletions(-) 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 92fe52038a48..17c1b56a7137 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 @@ -255,8 +255,17 @@ private Publisher(Builder builder) throws IOException { messagesWaiter = new Waiter(); this.hedgeSettings = builder.hedgeSettings; if (this.hedgeSettings != null) { - // Scale all configuration values to integers to allow lock-free atomic operations - // without the garbage collection overhead of autoboxing Float objects. + // Verify that the hedge delay is strictly less than the initial RPC timeout. + Duration hedgeDelay = this.hedgeSettings.getHedgeDelay(); + Duration initialRpcTimeout = builder.retrySettings.getInitialRpcTimeoutDuration(); + if (hedgeDelay.compareTo(initialRpcTimeout) >= 0) { + throw new IllegalArgumentException( + "hedgeDelay (" + + hedgeDelay.toMillis() + + "ms) must be strictly less than the initial RPC timeout duration (" + + initialRpcTimeout.toMillis() + + "ms)"); + } this.scaledMaxHedgeTokens = this.hedgeSettings.getMaxTokens() * HEDGE_TOKEN_SCALE; this.scaledHedgeRefillAmount = (int) (this.hedgeSettings.getRefillRatio() * HEDGE_TOKEN_SCALE); @@ -549,8 +558,7 @@ private ApiFuture publishCall( } outstandingBatch.publishRpcSpan = - tracer.startPublishRpcSpan(top - icNameObject, messageWrappers, attemptNumber); + tracer.startPublishRpcSpan(topicNameObject, messageWrappers, attemptNumber); return publisherStub .publishCallable() 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 f0e17ec83afd..693197e475b8 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 @@ -1353,6 +1353,54 @@ public void testPublisherWithHedgeSettings() throws Exception { shutdownTestPublisher(publisher); } + @Test + public void testPublisherThrowsIfHedgeDelayGtRpcTimeout() throws Exception { + HedgeSettings hedgeSettings = + HedgeSettings.newBuilder().setHedgeDelay(Duration.ofMillis(500)).build(); + com.google.api.gax.retrying.RetrySettings retrySettings = + com.google.api.gax.retrying.RetrySettings.newBuilder() + .setInitialRpcTimeoutDuration(Duration.ofMillis(400)) + .setMaxRpcTimeoutDuration(Duration.ofMillis(400)) + .setTotalTimeoutDuration(Duration.ofSeconds(10)) + .build(); + + try { + getTestPublisherBuilder() + .setHedgeSettings(hedgeSettings) + .setRetrySettings(retrySettings) + .build(); + fail( + "Should have thrown IllegalArgumentException because hedgeDelay (500ms) > RPC timeout (400ms)"); + } catch (IllegalArgumentException e) { + assertThat(e.getMessage()) + .contains("must be strictly less than the initial RPC timeout duration"); + } + } + + @Test + public void testPublisherThrowsIfHedgeDelayEqRpcTimeout() throws Exception { + HedgeSettings hedgeSettings = + HedgeSettings.newBuilder().setHedgeDelay(Duration.ofMillis(500)).build(); + com.google.api.gax.retrying.RetrySettings retrySettings = + com.google.api.gax.retrying.RetrySettings.newBuilder() + .setInitialRpcTimeoutDuration(Duration.ofMillis(500)) + .setMaxRpcTimeoutDuration(Duration.ofMillis(500)) + .setTotalTimeoutDuration(Duration.ofSeconds(10)) + .build(); + + try { + getTestPublisherBuilder() + .setHedgeSettings(hedgeSettings) + .setRetrySettings(retrySettings) + .build(); + fail( + "Should have thrown IllegalArgumentException because hedgeDelay (500ms) == RPC timeout (500ms)"); + } catch (IllegalArgumentException e) { + assertThat(e.getMessage()) + .contains("must be strictly less than the initial RPC timeout duration"); + } + } + @Test public void testPublisherWithoutHedgeSettings() throws Exception { Publisher publisher = getTestPublisherBuilder().build(); @@ -1441,6 +1489,7 @@ public void testHedgingTriggeredIfSlow() throws Exception { // Advance to 220ms to let responses complete fakeExecutor.advanceTime(Duration.ofMillis(100)); + fakeExecutor.advanceTime(Duration.ZERO); // Drain pending tasks assertTrue(future.isDone()); shutdownTestPublisher(publisher); From 8543d2345fdfc0641dc83605da88277e54955026 Mon Sep 17 00:00:00 2001 From: Tony Cui Date: Mon, 27 Jul 2026 20:04:55 +0000 Subject: [PATCH 17/22] Moved static methods in Publisher.java above the builder --- .../com/google/cloud/pubsub/v1/Publisher.java | 188 +++++++++--------- 1 file changed, 94 insertions(+), 94 deletions(-) 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 17c1b56a7137..7a0f9d58896a 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 @@ -656,7 +656,7 @@ public ApiFuture call() { ApiFutures.addCallback(future, futureCallback, callbackExecutor); } - void refillTokenBucket() { + private void refillTokenBucket() { if (hedgeSettings != null) { while (true) { int current = hedgingTokenBucket.get(); @@ -687,6 +687,99 @@ boolean tryAcquireHedgeToken() { } } + 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 (tryAcquireHedgeToken()) { + // 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(), item.getAttemptNumber()); + coordinator.addAttempt(item.getAttemptNumber(), hedgedFuture); + } else { + coordinator.isInQueue.set(false); + coordinator.checkCompletionOnQueueExit(); + } + } + + // Reschedule for next items + scheduleQueueProcessing(); + } + } + final class OutstandingBatch { final List outstandingPublishes; final long creationTime; @@ -1300,97 +1393,4 @@ && 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 (tryAcquireHedgeToken()) { - // 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(), item.getAttemptNumber()); - coordinator.addAttempt(item.getAttemptNumber(), hedgedFuture); - } else { - coordinator.isInQueue.set(false); - coordinator.checkCompletionOnQueueExit(); - } - } - - // Reschedule for next items - scheduleQueueProcessing(); - } - } } From 3c48a1a0fb0377b0a1b05a9818e114048338717d Mon Sep 17 00:00:00 2001 From: Tony Cui Date: Mon, 27 Jul 2026 21:47:25 +0000 Subject: [PATCH 18/22] Configure empty token bucket and fix naming --- .../com/google/cloud/pubsub/v1/Publisher.java | 19 ++++++----- .../pubsub/v1/FakePublisherServiceImpl.java | 4 +++ .../cloud/pubsub/v1/PublisherImplTest.java | 32 +++++++++++++++---- 3 files changed, 39 insertions(+), 16 deletions(-) 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 7a0f9d58896a..de044f5e1824 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 @@ -151,7 +151,7 @@ public class Publisher implements PublisherInterface { */ private static final int HEDGE_TOKEN_SCALE = 100; - private final AtomicInteger hedgingTokenBucket = new AtomicInteger(); + private final AtomicInteger hedgeTokenBucket = new AtomicInteger(); private int scaledMaxHedgeTokens; private int scaledHedgeRefillAmount; private final ApiClock clock; @@ -269,7 +269,7 @@ private Publisher(Builder builder) throws IOException { this.scaledMaxHedgeTokens = this.hedgeSettings.getMaxTokens() * HEDGE_TOKEN_SCALE; this.scaledHedgeRefillAmount = (int) (this.hedgeSettings.getRefillRatio() * HEDGE_TOKEN_SCALE); - this.hedgingTokenBucket.set(scaledMaxHedgeTokens); + this.hedgeTokenBucket.set(0); } this.clock = builder.clock != null ? builder.clock : CurrentMillisClock.getDefaultClock(); this.publishContext = GrpcCallContext.createDefault(); @@ -301,7 +301,7 @@ Float getHedgeTokenBalance() { if (hedgeSettings == null) { return null; } - return (float) hedgingTokenBucket.get() / HEDGE_TOKEN_SCALE; + return (float) hedgeTokenBucket.get() / HEDGE_TOKEN_SCALE; } /** @@ -656,15 +656,15 @@ public ApiFuture call() { ApiFutures.addCallback(future, futureCallback, callbackExecutor); } - private void refillTokenBucket() { + void refillTokenBucket() { if (hedgeSettings != null) { while (true) { - int current = hedgingTokenBucket.get(); + int current = hedgeTokenBucket.get(); if (current >= scaledMaxHedgeTokens) { return; } int next = Math.min(scaledMaxHedgeTokens, current + scaledHedgeRefillAmount); - if (hedgingTokenBucket.compareAndSet(current, next)) { + if (hedgeTokenBucket.compareAndSet(current, next)) { return; } } @@ -676,12 +676,12 @@ boolean tryAcquireHedgeToken() { return false; } while (true) { - int current = hedgingTokenBucket.get(); + int current = hedgeTokenBucket.get(); if (current < HEDGE_TOKEN_SCALE) { return false; } int next = current - HEDGE_TOKEN_SCALE; - if (hedgingTokenBucket.compareAndSet(current, next)) { + if (hedgeTokenBucket.compareAndSet(current, next)) { return true; } } @@ -727,8 +727,7 @@ private void scheduleQueueProcessing() { return; } - long delay = nextItem.getSendAfterMs() - clock.millisTime(); - delay = Math.max(0, delay); + long delay = Math.max(0, nextItem.getSendAfterMs() - clock.millisTime()); queueProcessingFuture = executor.schedule( 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 9424018b382f..6a20300e44e1 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 @@ -161,4 +161,8 @@ public FakePublisherServiceImpl addPublishError(Throwable error) { public List getCapturedRequests() { return new ArrayList(requests); } + + public void clearRequests() { + requests.clear(); + } } 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 693197e475b8..b22138671d18 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 @@ -1348,7 +1348,7 @@ public void testPublisherWithHedgeSettings() throws Exception { assertThat(publisher.getHedgeSettings()).isEqualTo(hedgeSettings); assertThat(publisher.getHedgeTokenBalance()).isNotNull(); - assertThat(publisher.getHedgeTokenBalance()).isWithin(0.0001f).of(50.0f); + assertThat(publisher.getHedgeTokenBalance()).isWithin(0.0001f).of(0.0f); shutdownTestPublisher(publisher); } @@ -1412,7 +1412,15 @@ public void testPublisherWithoutHedgeSettings() throws Exception { } private Publisher getPublisherWithHedge(Duration delay) throws Exception { - HedgeSettings hedgeSettings = HedgeSettings.newBuilder().setHedgeDelay(delay).build(); + return getPublisherWithHedge(delay, 0.1f, 20); + } + + private Publisher getPublisherWithHedge(Duration delay, float refillRatio, int maxTokens) throws Exception { + HedgeSettings hedgeSettings = HedgeSettings.newBuilder() + .setHedgeDelay(delay) + .setRefillRatio(refillRatio) + .setMaxTokens(maxTokens) + .build(); return getTestPublisherBuilder() .setHedgeSettings(hedgeSettings) .setClock(fakeExecutor.getClock()) @@ -1423,6 +1431,15 @@ private Publisher getPublisherWithHedge(Duration delay) throws Exception { .build(); } + private void fillTokenBucket(Publisher publisher, int tokensToFill) throws Exception { + testPublisherServiceImpl.setAutoPublishResponse(true); + for (int i = 0; i < tokensToFill; i++) { + ApiFuture future = sendTestMessage(publisher, "warmup-msg-" + i); + future.get(); + } + testPublisherServiceImpl.clearRequests(); + } + private void waitForRequests(FakePublisherServiceImpl service, int expectedCount) throws InterruptedException { long timeout = System.currentTimeMillis() + 5000; @@ -1464,7 +1481,8 @@ public void testHedgingNotTriggeredIfFast() throws Exception { @Test public void testHedgingTriggeredIfSlow() throws Exception { - Publisher publisher = getPublisherWithHedge(Duration.ofMillis(100)); + Publisher publisher = getPublisherWithHedge(Duration.ofMillis(100), 0.2f, 20); + fillTokenBucket(publisher, 5); // Set response delay to 200ms (greater than 100ms hedge delay) testPublisherServiceImpl.setAutoPublishResponse(false); @@ -1490,14 +1508,15 @@ public void testHedgingTriggeredIfSlow() throws Exception { // Advance to 220ms to let responses complete fakeExecutor.advanceTime(Duration.ofMillis(100)); fakeExecutor.advanceTime(Duration.ZERO); // Drain pending tasks - assertTrue(future.isDone()); + assertEquals("1", future.get(5, TimeUnit.SECONDS)); shutdownTestPublisher(publisher); } @Test public void testMultipleHedging() throws Exception { - Publisher publisher = getPublisherWithHedge(Duration.ofMillis(100)); + Publisher publisher = getPublisherWithHedge(Duration.ofMillis(100), 0.2f, 20); + fillTokenBucket(publisher, 10); // Set delay to 400ms testPublisherServiceImpl.setAutoPublishResponse(false); @@ -1556,7 +1575,8 @@ public void testHedgingBypassedIfNoTokens() throws Exception { @Test public void testHedgingCancellationPropagates() throws Exception { - Publisher publisher = getPublisherWithHedge(Duration.ofMillis(100)); + Publisher publisher = getPublisherWithHedge(Duration.ofMillis(100), 0.2f, 20); + fillTokenBucket(publisher, 5); testPublisherServiceImpl.setAutoPublishResponse(false); testPublisherServiceImpl.setPublishResponseDelay(Duration.ofMillis(200)); From 0929a1bff7cc98bc9c3e9d77ab7c0eb596f80f17 Mon Sep 17 00:00:00 2001 From: Tony Cui Date: Mon, 27 Jul 2026 22:30:53 +0000 Subject: [PATCH 19/22] Change from synchronization to actual lock object for processQueue method --- .../main/java/com/google/cloud/pubsub/v1/Publisher.java | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) 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 de044f5e1824..6b14bf4c5375 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 @@ -159,7 +159,7 @@ public class Publisher implements PublisherInterface { private final ConcurrentLinkedQueue hedgingQueue; private final AtomicBoolean isQueueProcessingScheduled; private ScheduledFuture queueProcessingFuture; - private final Object queueLock; + private final Lock queueLock; /** The maximum number of messages in one request. Defined by the API. */ public static long getApiMaxRequestElementCount() { @@ -278,7 +278,7 @@ private Publisher(Builder builder) throws IOException { .withCallOptions(CallOptions.DEFAULT.withCompression(GZIP_COMPRESSION)); this.hedgingQueue = new ConcurrentLinkedQueue<>(); this.isQueueProcessingScheduled = new AtomicBoolean(false); - this.queueLock = new Object(); + this.queueLock = new ReentrantLock(); this.queueProcessingFuture = null; } @@ -743,7 +743,8 @@ public void run() { } private void processQueue() { - synchronized (queueLock) { + queueLock.lock(); + try { isQueueProcessingScheduled.set(false); long now = clock.millisTime(); @@ -776,6 +777,8 @@ private void processQueue() { // Reschedule for next items scheduleQueueProcessing(); + } finally { + queueLock.unlock(); } } From 02951441186f8741824061e90ce78764997f8930 Mon Sep 17 00:00:00 2001 From: Tony Cui Date: Tue, 28 Jul 2026 19:10:05 +0000 Subject: [PATCH 20/22] Add header injection for hedged publish requests --- .../com/google/cloud/pubsub/v1/Publisher.java | 10 +++++++ .../pubsub/v1/FakePublisherServiceImpl.java | 11 ++++++++ .../cloud/pubsub/v1/PublisherImplTest.java | 27 ++++++++++++++++++- 3 files changed, 47 insertions(+), 1 deletion(-) 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 6b14bf4c5375..582ca1ee499b 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 @@ -48,6 +48,8 @@ import com.google.cloud.pubsub.v1.stub.PublisherStub; import com.google.cloud.pubsub.v1.stub.PublisherStubSettings; import com.google.common.base.Preconditions; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; import com.google.protobuf.CodedOutputStream; import com.google.pubsub.v1.PublishRequest; import com.google.pubsub.v1.PublishResponse; @@ -61,6 +63,7 @@ import java.io.IOException; import java.time.Duration; import java.util.ArrayList; + import java.util.HashMap; import java.util.Iterator; import java.util.LinkedList; @@ -143,6 +146,7 @@ public class Publisher implements PublisherInterface { private OpenTelemetryPubsubTracer tracer = new OpenTelemetryPubsubTracer(null, false); private final HedgeSettings hedgeSettings; + private final Map> hedgingMetadata; /** * Scale factor to represent decimal token values (e.g. 0.1 refill ratio) as integers inside the @@ -151,6 +155,7 @@ public class Publisher implements PublisherInterface { */ private static final int HEDGE_TOKEN_SCALE = 100; + private final AtomicInteger hedgeTokenBucket = new AtomicInteger(); private int scaledMaxHedgeTokens; private int scaledHedgeRefillAmount; @@ -273,6 +278,8 @@ private Publisher(Builder builder) throws IOException { } this.clock = builder.clock != null ? builder.clock : CurrentMillisClock.getDefaultClock(); this.publishContext = GrpcCallContext.createDefault(); + this.hedgingMetadata = + ImmutableMap.of("x-goog-pubsub-hedged", ImmutableList.of("true")); this.publishContextWithCompression = GrpcCallContext.createDefault() .withCallOptions(CallOptions.DEFAULT.withCompression(GZIP_COMPRESSION)); @@ -548,6 +555,9 @@ private ApiFuture publishCall( if (enableCompression && outstandingBatch.batchSizeBytes >= compressionBytesThreshold) { context = publishContextWithCompression; } + if (attemptNumber > 1) { + context = context.withExtraHeaders(hedgingMetadata); + } int numMessagesInBatch = outstandingBatch.size(); List pubsubMessagesList = new ArrayList(numMessagesInBatch); 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 6a20300e44e1..8e0cf6c6083c 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 @@ -22,6 +22,7 @@ import com.google.pubsub.v1.PublisherGrpc.PublisherImplBase; import io.grpc.stub.StreamObserver; import java.time.Duration; +import io.grpc.Metadata; import java.util.ArrayList; import java.util.List; import java.util.concurrent.LinkedBlockingQueue; @@ -36,6 +37,7 @@ class FakePublisherServiceImpl extends PublisherImplBase { private final LinkedBlockingQueue requests = new LinkedBlockingQueue<>(); + private final LinkedBlockingQueue capturedHeaders = new LinkedBlockingQueue<>(); private final LinkedBlockingQueue publishResponses = new LinkedBlockingQueue<>(); private final AtomicInteger nextMessageId = new AtomicInteger(1); private boolean autoPublishResponse; @@ -162,7 +164,16 @@ public List getCapturedRequests() { return new ArrayList(requests); } + public void recordHeaders(Metadata headers) { + capturedHeaders.add(headers); + } + + public List getCapturedHeaders() { + return new ArrayList<>(capturedHeaders); + } + public void clearRequests() { requests.clear(); + capturedHeaders.clear(); } } 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 b22138671d18..012cd3c7d8a0 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 @@ -43,6 +43,11 @@ import com.google.pubsub.v1.PublishResponse; import com.google.pubsub.v1.PubsubMessage; import io.grpc.ManagedChannel; +import io.grpc.Metadata; +import io.grpc.ServerCall; +import io.grpc.ServerCallHandler; +import io.grpc.ServerInterceptor; +import io.grpc.ServerInterceptors; import io.grpc.Server; import io.grpc.Status; import io.grpc.StatusException; @@ -98,8 +103,19 @@ public class PublisherImplTest { public void setUp() throws Exception { testPublisherServiceImpl = new FakePublisherServiceImpl(); + ServerInterceptor headerInterceptor = + new ServerInterceptor() { + @Override + public ServerCall.Listener interceptCall( + ServerCall call, Metadata headers, ServerCallHandler next) { + testPublisherServiceImpl.recordHeaders(headers); + return next.startCall(call, headers); + } + }; + InProcessServerBuilder serverBuilder = InProcessServerBuilder.forName("test-server"); - serverBuilder.addService(testPublisherServiceImpl); + serverBuilder.addService( + ServerInterceptors.intercept(testPublisherServiceImpl, headerInterceptor)); testServer = serverBuilder.build(); testChannel = InProcessChannelBuilder.forName("test-server").build(); testServer.start(); @@ -1510,6 +1526,15 @@ public void testHedgingTriggeredIfSlow() throws Exception { fakeExecutor.advanceTime(Duration.ZERO); // Drain pending tasks assertEquals("1", future.get(5, TimeUnit.SECONDS)); + List capturedHeaders = testPublisherServiceImpl.getCapturedHeaders(); + assertThat(capturedHeaders).hasSize(2); + Metadata.Key hedgedHeaderKey = + Metadata.Key.of("x-goog-pubsub-hedged", Metadata.ASCII_STRING_MARSHALLER); + // Original request should NOT have the header + assertThat(capturedHeaders.get(0).get(hedgedHeaderKey)).isNull(); + // Hedged request SHOULD have the header set to "true" + assertThat(capturedHeaders.get(1).get(hedgedHeaderKey)).isEqualTo("true"); + shutdownTestPublisher(publisher); } From c8296f5313d262d67c6fe7a0ad39e48a5d69f5a3 Mon Sep 17 00:00:00 2001 From: Tony Cui Date: Tue, 28 Jul 2026 20:19:20 +0000 Subject: [PATCH 21/22] Cleanup code and fix bug with handling failure attempts in the first rpc call --- .../cloud/pubsub/v1/CancellationSharer.java | 26 ++++++++------- .../com/google/cloud/pubsub/v1/Publisher.java | 15 +++++++++ .../cloud/pubsub/v1/PublisherImplTest.java | 33 +++++++++++++++++++ 3 files changed, 63 insertions(+), 11 deletions(-) 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 index b52e3f4f660f..59ecf13f5891 100644 --- 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 @@ -89,12 +89,14 @@ private void handleAttemptSuccess(int attemptNumber, PublishResponse response) { 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()); - } + if (done.get()) { + return; + } + lastError.set(t); + if (runningAttempts.isEmpty() && done.compareAndSet(false, true)) { + setException(lastError.get()); + if (isInQueue.get()) { + publisher.removeFromHedgingQueue(this); } } } @@ -120,12 +122,14 @@ private void cancelAllExcept(int successfulAttempt) { @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; + done.set(true); + if (isInQueue.get()) { + publisher.removeFromHedgingQueue(this); + } + for (ApiFuture future : runningAttempts.values()) { + future.cancel(mayInterruptIfRunning); } + return true; } return false; } 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 582ca1ee499b..1db26e1f571d 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 @@ -752,6 +752,21 @@ public void run() { } } + void removeFromHedgingQueue(CancellationSharer coordinator) { + queueLock.lock(); + try { + Iterator iterator = hedgingQueue.iterator(); + while (iterator.hasNext()) { + if (iterator.next().getCoordinator() == coordinator) { + iterator.remove(); + coordinator.isInQueue.set(false); + } + } + } finally { + queueLock.unlock(); + } + } + private void processQueue() { queueLock.lock(); try { 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 012cd3c7d8a0..3c6093f8c998 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 @@ -34,6 +34,7 @@ import com.google.api.gax.grpc.GrpcTransportChannel; import com.google.api.gax.grpc.testing.LocalChannelProvider; import com.google.api.gax.rpc.DataLossException; +import com.google.api.gax.rpc.InvalidArgumentException; import com.google.api.gax.rpc.FixedTransportChannelProvider; import com.google.api.gax.rpc.TransportChannelProvider; import com.google.cloud.pubsub.v1.Publisher.Builder; @@ -1625,6 +1626,38 @@ public void testHedgingCancellationPropagates() throws Exception { shutdownTestPublisher(publisher); } + @Test + public void testNoHedgingIfOriginalFailsImmediately() throws Exception { + Publisher publisher = getPublisherWithHedge(Duration.ofMillis(100), 0.2f, 20); + fillTokenBucket(publisher, 5); + + // Configure the fake to immediately return an INVALID_ARGUMENT error + testPublisherServiceImpl.setAutoPublishResponse(false); + testPublisherServiceImpl.addPublishError(new StatusException(Status.INVALID_ARGUMENT)); + + ApiFuture future = sendTestMessage(publisher, "msg-fail-fast"); + + // The request should fail immediately without waiting or advancing time + try { + future.get(1, TimeUnit.SECONDS); + fail("Should have failed with ExecutionException"); + } catch (ExecutionException e) { + // expected + assertThat(e.getCause()).isInstanceOf(InvalidArgumentException.class); + } + + // Server should receive exactly 1 request (the original attempt) + assertThat(testPublisherServiceImpl.getCapturedRequests()).hasSize(1); + + // Advance time past the 100ms hedge delay and check that no hedge was sent + fakeExecutor.advanceTime(Duration.ofMillis(200)); + + // Captured requests should still be 1 (no hedge triggered) + assertThat(testPublisherServiceImpl.getCapturedRequests()).hasSize(1); + + shutdownTestPublisher(publisher); + } + private Builder getTestPublisherBuilder() { return Publisher.newBuilder(TEST_TOPIC) .setExecutorProvider(FixedExecutorProvider.create(fakeExecutor)) From c8fd6e99f0fc77c0831ed36e9280898b248a5945 Mon Sep 17 00:00:00 2001 From: Tony Cui Date: Tue, 28 Jul 2026 20:47:36 +0000 Subject: [PATCH 22/22] Disable retries for failing rpcs when publish hedging --- .../com/google/cloud/pubsub/v1/Publisher.java | 12 +++++----- .../pubsub/v1/FakePublisherServiceImpl.java | 2 +- .../cloud/pubsub/v1/PublisherImplTest.java | 22 ++++++++++--------- 3 files changed, 20 insertions(+), 16 deletions(-) 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 1db26e1f571d..85ff8c95c78e 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 @@ -63,7 +63,7 @@ import java.io.IOException; import java.time.Duration; import java.util.ArrayList; - +import java.util.Collections; import java.util.HashMap; import java.util.Iterator; import java.util.LinkedList; @@ -155,7 +155,6 @@ public class Publisher implements PublisherInterface { */ private static final int HEDGE_TOKEN_SCALE = 100; - private final AtomicInteger hedgeTokenBucket = new AtomicInteger(); private int scaledMaxHedgeTokens; private int scaledHedgeRefillAmount; @@ -278,8 +277,7 @@ private Publisher(Builder builder) throws IOException { } this.clock = builder.clock != null ? builder.clock : CurrentMillisClock.getDefaultClock(); this.publishContext = GrpcCallContext.createDefault(); - this.hedgingMetadata = - ImmutableMap.of("x-goog-pubsub-hedged", ImmutableList.of("true")); + this.hedgingMetadata = ImmutableMap.of("x-goog-pubsub-hedged", ImmutableList.of("true")); this.publishContextWithCompression = GrpcCallContext.createDefault() .withCallOptions(CallOptions.DEFAULT.withCompression(GZIP_COMPRESSION)); @@ -556,7 +554,11 @@ private ApiFuture publishCall( context = publishContextWithCompression; } if (attemptNumber > 1) { - context = context.withExtraHeaders(hedgingMetadata); + logger.log(Level.FINER, "Publishing hedged attempt {0}", attemptNumber); + context = + context + .withExtraHeaders(hedgingMetadata) + .withRetryableCodes(Collections.emptySet()); } int numMessagesInBatch = outstandingBatch.size(); 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 8e0cf6c6083c..315c01a0ff55 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 @@ -20,9 +20,9 @@ import com.google.pubsub.v1.PublishRequest; import com.google.pubsub.v1.PublishResponse; import com.google.pubsub.v1.PublisherGrpc.PublisherImplBase; +import io.grpc.Metadata; import io.grpc.stub.StreamObserver; import java.time.Duration; -import io.grpc.Metadata; import java.util.ArrayList; import java.util.List; import java.util.concurrent.LinkedBlockingQueue; 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 3c6093f8c998..aa2f2e660a09 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 @@ -34,8 +34,8 @@ import com.google.api.gax.grpc.GrpcTransportChannel; import com.google.api.gax.grpc.testing.LocalChannelProvider; import com.google.api.gax.rpc.DataLossException; -import com.google.api.gax.rpc.InvalidArgumentException; import com.google.api.gax.rpc.FixedTransportChannelProvider; +import com.google.api.gax.rpc.InvalidArgumentException; import com.google.api.gax.rpc.TransportChannelProvider; import com.google.cloud.pubsub.v1.Publisher.Builder; import com.google.protobuf.ByteString; @@ -45,11 +45,11 @@ import com.google.pubsub.v1.PubsubMessage; import io.grpc.ManagedChannel; import io.grpc.Metadata; +import io.grpc.Server; import io.grpc.ServerCall; import io.grpc.ServerCallHandler; import io.grpc.ServerInterceptor; import io.grpc.ServerInterceptors; -import io.grpc.Server; import io.grpc.Status; import io.grpc.StatusException; import io.grpc.inprocess.InProcessChannelBuilder; @@ -1432,12 +1432,14 @@ private Publisher getPublisherWithHedge(Duration delay) throws Exception { return getPublisherWithHedge(delay, 0.1f, 20); } - private Publisher getPublisherWithHedge(Duration delay, float refillRatio, int maxTokens) throws Exception { - HedgeSettings hedgeSettings = HedgeSettings.newBuilder() - .setHedgeDelay(delay) - .setRefillRatio(refillRatio) - .setMaxTokens(maxTokens) - .build(); + private Publisher getPublisherWithHedge(Duration delay, float refillRatio, int maxTokens) + throws Exception { + HedgeSettings hedgeSettings = + HedgeSettings.newBuilder() + .setHedgeDelay(delay) + .setRefillRatio(refillRatio) + .setMaxTokens(maxTokens) + .build(); return getTestPublisherBuilder() .setHedgeSettings(hedgeSettings) .setClock(fakeExecutor.getClock()) @@ -1636,7 +1638,7 @@ public void testNoHedgingIfOriginalFailsImmediately() throws Exception { testPublisherServiceImpl.addPublishError(new StatusException(Status.INVALID_ARGUMENT)); ApiFuture future = sendTestMessage(publisher, "msg-fail-fast"); - + // The request should fail immediately without waiting or advancing time try { future.get(1, TimeUnit.SECONDS); @@ -1651,7 +1653,7 @@ public void testNoHedgingIfOriginalFailsImmediately() throws Exception { // Advance time past the 100ms hedge delay and check that no hedge was sent fakeExecutor.advanceTime(Duration.ofMillis(200)); - + // Captured requests should still be 1 (no hedge triggered) assertThat(testPublisherServiceImpl.getCapturedRequests()).hasSize(1);