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..ffa6daed1d91 --- /dev/null +++ b/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/HedgeSettings.java @@ -0,0 +1,113 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.pubsub.v1; + +import com.google.common.base.Preconditions; +import java.time.Duration; + +/** Settings for configuring publish hedging. */ +public final class HedgeSettings { + /** Default hedging delay. */ + private static final Duration DEFAULT_DELAY = Duration.ofMillis(100); + + /** Default maximum number of tokens in the bucket. */ + private static final int DEFAULT_MAX_TOKENS = 100; + + /** Default refill rate (tokens per successful request). */ + private static final float DEFAULT_REFILL = 0.1f; + + /** Hedging delay. */ + private final Duration hedgeDelay; + + /** Maximum tokens. */ + private final int maxTokens; + + /** Refill rate. */ + private final float refill; + + private HedgeSettings(final Builder builder) { + this.hedgeDelay = builder.hedgeDelay; + this.maxTokens = builder.maxTokens; + this.refill = builder.refill; + } + + /** + * Returns the configured hedging delay. + * + * @return the hedging delay. + */ + public Duration getHedgeDelay() { + return hedgeDelay; + } + + int getMaxTokens() { + return maxTokens; + } + + float getRefill() { + return refill; + } + + /** + * Returns a new builder for {@code HedgeSettings}. + * + * @return a new builder. + */ + public static Builder newBuilder() { + return new Builder(); + } + + /** Builder for {@code HedgeSettings}. */ + public static final class Builder { + /** Hedging delay. */ + private Duration hedgeDelay = DEFAULT_DELAY; + + /** Maximum tokens. */ + private int maxTokens = DEFAULT_MAX_TOKENS; + + /** Refill rate. */ + private float refill = DEFAULT_REFILL; + + private Builder() { + super(); + } + + /** + * Allows hedging delay to be configurable. + * + * @param delay the hedging delay, must be positive. + * @return this builder. + */ + public Builder setHedgeDelay(final Duration delay) { + Preconditions.checkNotNull(delay); + if (delay.isNegative() || delay.isZero()) { + throw new IllegalArgumentException("delay must be positive"); + } + this.hedgeDelay = delay; + return this; + } + + /** + * Builds an instance of {@code HedgeSettings}. + * + * @return the built {@code HedgeSettings} instance. + */ + public HedgeSettings build() { + return new HedgeSettings(this); + } + } +} diff --git a/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/HedgeTokenBucket.java b/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/HedgeTokenBucket.java new file mode 100644 index 000000000000..5378478cb429 --- /dev/null +++ b/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/HedgeTokenBucket.java @@ -0,0 +1,65 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.pubsub.v1; + +import com.google.common.annotations.VisibleForTesting; +import javax.annotation.concurrent.GuardedBy; + +/** Token bucket for limiting hedged requests. Thread-safe. */ +final class HedgeTokenBucket { + private final int maxTokens; + private final float refillAmount; + + @GuardedBy("this") + private float tokens; + + HedgeTokenBucket(HedgeSettings settings) { + this.maxTokens = settings.getMaxTokens(); + this.refillAmount = settings.getRefill(); + this.tokens = maxTokens; + } + + @VisibleForTesting + HedgeTokenBucket(int maxTokens, float refillAmount) { + this.maxTokens = maxTokens; + this.refillAmount = refillAmount; + this.tokens = maxTokens; + } + + /** + * Attempts to acquire a token for a hedged request. + * + * @return {@code true} if a token was acquired, {@code false} otherwise. + */ + synchronized boolean tryAcquire() { + if (tokens >= 1.0f) { + tokens -= 1.0f; + return true; + } + return false; + } + + /** Refills the bucket by the configured refill amount, capped at max tokens. */ + synchronized void refill() { + tokens = Math.min(maxTokens, tokens + refillAmount); + } + + @VisibleForTesting + synchronized float getTokens() { + return tokens; + } +} diff --git a/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/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/HedgeSettingsTest.java b/java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/HedgeSettingsTest.java new file mode 100644 index 000000000000..f3b26aedb3c3 --- /dev/null +++ b/java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/HedgeSettingsTest.java @@ -0,0 +1,66 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.pubsub.v1; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertThrows; + +import java.time.Duration; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class HedgeSettingsTest { + + @Test + public void testDefaultSettings() { + HedgeSettings settings = HedgeSettings.newBuilder().build(); + assertNotNull(settings); + assertEquals(Duration.ofMillis(100), settings.getHedgeDelay()); + assertEquals(100, settings.getMaxTokens()); + assertEquals(0.1f, settings.getRefill(), 0.0001f); + } + + @Test + public void testCustomDelay() { + Duration customDelay = Duration.ofMillis(50); + HedgeSettings settings = HedgeSettings.newBuilder().setHedgeDelay(customDelay).build(); + assertNotNull(settings); + assertEquals(customDelay, settings.getHedgeDelay()); + } + + @Test + public void testNegativeDelayThrows() { + assertThrows( + IllegalArgumentException.class, + () -> HedgeSettings.newBuilder().setHedgeDelay(Duration.ofMillis(-10))); + } + + @Test + public void testZeroDelayThrows() { + assertThrows( + IllegalArgumentException.class, + () -> HedgeSettings.newBuilder().setHedgeDelay(Duration.ZERO)); + } + + @Test + public void testNullDelayThrows() { + assertThrows(NullPointerException.class, () -> HedgeSettings.newBuilder().setHedgeDelay(null)); + } +} diff --git a/java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/HedgeTokenBucketTest.java b/java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/HedgeTokenBucketTest.java new file mode 100644 index 000000000000..e92ecd9a6ce2 --- /dev/null +++ b/java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/HedgeTokenBucketTest.java @@ -0,0 +1,71 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.pubsub.v1; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class HedgeTokenBucketTest { + + @Test + public void testInitialState() { + HedgeTokenBucket bucket = new HedgeTokenBucket(10, 0.5f); + assertEquals(10.0f, bucket.getTokens(), 0.0001f); + } + + @Test + public void testAcquire() { + HedgeTokenBucket bucket = new HedgeTokenBucket(10, 0.5f); + assertTrue(bucket.tryAcquire()); + assertEquals(9.0f, bucket.getTokens(), 0.0001f); + } + + @Test + public void testAcquireUntilEmpty() { + HedgeTokenBucket bucket = new HedgeTokenBucket(2, 0.5f); + assertTrue(bucket.tryAcquire()); + assertTrue(bucket.tryAcquire()); + assertFalse(bucket.tryAcquire()); + assertEquals(0.0f, bucket.getTokens(), 0.0001f); + } + + @Test + public void testRefill() { + HedgeTokenBucket bucket = new HedgeTokenBucket(2, 0.5f); + assertTrue(bucket.tryAcquire()); // 1.0 left + assertTrue(bucket.tryAcquire()); // 0.0 left + + bucket.refill(); // 0.5 tokens + assertFalse(bucket.tryAcquire()); // needs 1.0 + + bucket.refill(); // 1.0 tokens + assertTrue(bucket.tryAcquire()); // succeeds, 0.0 left + } + + @Test + public void testRefillCap() { + HedgeTokenBucket bucket = new HedgeTokenBucket(2, 0.5f); + bucket.refill(); // already at max (2) + assertEquals(2.0f, bucket.getTokens(), 0.0001f); + } +} diff --git a/java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/PublisherImplTest.java b/java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/PublisherImplTest.java index 8e6efaf372c9..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))