Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
@@ -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);
}
}
}
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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

In performance-sensitive code like the Pub/Sub Publisher, using synchronized methods on the critical path can introduce significant lock contention and performance overhead. To protect shared state while ensuring thread safety and visibility, prefer using explicit locks (such as ReentrantLock) over the synchronized keyword.

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
  1. In performance-sensitive code, prefer using explicit locks over the 'synchronized' keyword to protect shared state while ensuring thread safety and visibility.

}
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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()
Expand All @@ -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.
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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;
Expand Down
Original file line number Diff line number Diff line change
@@ -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));
}
}
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);
}
}
Loading
Loading