-
Notifications
You must be signed in to change notification settings - Fork 1.2k
feat(pubsub): implement HedgeSettings and integrate with Publisher #13735
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from 4 commits
61d0b60
19afd68
49426a9
22417d1
91706a2
bd76ede
61aef08
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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() { | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. nit: Does this method need to be public? I think leaving it without an access modifier (i.e., making it package-private) should be sufficient. |
||
| 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); | ||
| } | ||
| } | ||
| } | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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; | ||
| } | ||
|
Comment on lines
+19
to
+64
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. In performance-sensitive code like the Pub/Sub import com.google.common.annotations.VisibleForTesting;
import java.util.concurrent.locks.ReentrantLock;
/** Token bucket for limiting hedged requests. Thread-safe. */
final class HedgeTokenBucket {
private final ReentrantLock lock = new ReentrantLock();
private final int maxTokens;
private final float refillAmount;
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.
*/
boolean tryAcquire() {
lock.lock();
try {
if (tokens < 1.0f) {
return false;
}
tokens -= 1.0f;
return true;
} finally {
lock.unlock();
}
}
/** Refills the bucket by the configured refill amount, capped at max tokens. */
void refill() {
lock.lock();
try {
tokens = Math.min((float) maxTokens, tokens + refillAmount);
} finally {
lock.unlock();
}
}
@VisibleForTesting
float getTokens() {
lock.lock();
try {
return tokens;
} finally {
lock.unlock();
}
}References
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I think this is a valid concern regarding high contention from the use of synchronized for lock synchronization in this class. I'll take a deeper look on Monday at what better options would be. |
||
| } | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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) { | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This method should enforce the constraint that ordering is not enabled. |
||
| 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; | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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)); | ||
| } | ||
| } |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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); | ||
| } | ||
| } |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
As a note, this value will be changing (potentially other default values as well). Leaving this as a placeholder comment to ensure these meet the agreed upon values before merging.