feat(pubsub): implement publish hedging to reduce tail latency

java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/CancellationSharer.java

@@ -0,0 +1,143 @@
+/*
+ * Copyright 2026 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.pubsub.v1;
+
+import com.google.api.core.AbstractApiFuture;
+import com.google.api.core.ApiFuture;
+import com.google.api.core.ApiFutureCallback;
+import com.google.api.core.ApiFutures;
+import com.google.common.util.concurrent.MoreExecutors;
+import com.google.pubsub.v1.PublishResponse;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
+
+/**
+ * Coordinates multiple publish attempts for a single batch of messages.
+ *
+ * <p>Implements {@link ApiFuture} to act as the single future returned to the publisher's client.
+ * It manages the lifecycle of the original attempt and any subsequent hedged attempts.
+ */
+class CancellationSharer extends AbstractApiFuture<PublishResponse> {
+  private final Publisher.OutstandingBatch batch;
+  private final Publisher publisher;
+  private final Map<Integer, ApiFuture<PublishResponse>> 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<Throwable> 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<PublishResponse> future) {
+    runningAttempts.put(attemptNumber, future);
+    totalAttemptsCount.incrementAndGet();
+
+    if (done.get()) {
+      future.cancel(true);
+      runningAttempts.remove(attemptNumber);
+      return;
+    }
+
+    ApiFutures.addCallback(
+        future,
+        new ApiFutureCallback<PublishResponse>() {
+          @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<Integer, ApiFuture<PublishResponse>> 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<PublishResponse> future : runningAttempts.values()) {
+          future.cancel(mayInterruptIfRunning);
+        }
+        return true;
+      }
+    }
+    return false;
+  }
+
+  int getAttemptCount() {
+    return totalAttemptsCount.get();
+  }
+
+  Publisher.OutstandingBatch getBatch() {
+    return batch;
+  }
+}

java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/HedgeSettings.java

@@ -0,0 +1,111 @@
+/*
+ * Copyright 2026 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.pubsub.v1;
+
+import com.google.common.base.Preconditions;
+import java.time.Duration;
+
+/** Settings for configuring publish hedging. */
+public final class HedgeSettings {
+  /** Default hedging delay. */
+  private static final Duration DEFAULT_DELAY = Duration.ofMillis(100);
+
+  /** Default maximum number of tokens in the bucket. */
+  private static final int DEFAULT_MAX_TOKENS = 100;
+
+  /** Default refill rate (tokens per successful request). */
+  private static final float DEFAULT_REFILL = 0.1f;
+
+  /** Hedging delay. */
+  private final Duration hedgeDelay;
+
+  /** Maximum tokens. */
+  private final int maxTokens;
+
+  /** Refill rate. */
+  private final float refill;
+
+  private HedgeSettings(final Builder builder) {
+    this.hedgeDelay = builder.hedgeDelay;
+    this.maxTokens = builder.maxTokens;
+    this.refill = builder.refill;
+  }
+
+  /**
+   * Returns the configured hedging delay.
+   *
+   * @return the hedging delay.
+   */
+  Duration getHedgeDelay() {
+    return hedgeDelay;
+  }
+
+  int getMaxTokens() {
+    return maxTokens;
+  }
+
+  float getRefill() {
+    return refill;
+  }
+
+  /**
+   * Returns a new builder for {@code HedgeSettings}.
+   *
+   * @return a new builder.
+   */
+  public static Builder newBuilder() {
+    return new Builder();
+  }
+
+  /** Builder for {@code HedgeSettings}. */
+  public static final class Builder {
+    /** Hedging delay. */
+    private Duration hedgeDelay = DEFAULT_DELAY;
+
+    /** Maximum tokens. */
+    private int maxTokens = DEFAULT_MAX_TOKENS;
+
+    /** Refill rate. */
+    private float refill = DEFAULT_REFILL;
+
+    private Builder() {}
+
+    /**
+     * Allows hedging delay to be configurable.
+     *
+     * @param delay the hedging delay, must be positive.
+     * @return this builder.
+     */
+    public Builder setHedgeDelay(final Duration delay) {
+      Preconditions.checkNotNull(delay);
+      if (delay.isNegative() || delay.isZero()) {
+        throw new IllegalArgumentException("delay must be positive");
+      }
+      this.hedgeDelay = delay;
+      return this;
+    }
+
+    /**
+     * Builds an instance of {@code HedgeSettings}.
+     *
+     * @return the built {@code HedgeSettings} instance.
+     */
+    public HedgeSettings build() {
+      return new HedgeSettings(this);
+    }
+  }
+}

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;
+  }
+}

java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/HedgedRequest.java

@@ -0,0 +1,42 @@
+/*
+ * Copyright 2026 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.pubsub.v1;
+
+/** Represents a pending hedging check in the publisher's queue. */
+class HedgedRequest {
+  private final CancellationSharer coordinator;
+  private final int attemptNumber;
+  private final long sendAfterMs;
+
+  HedgedRequest(CancellationSharer coordinator, int attemptNumber, long sendAfterMs) {
+    this.coordinator = coordinator;
+    this.attemptNumber = attemptNumber;
+    this.sendAfterMs = sendAfterMs;
+  }
+
+  CancellationSharer getCoordinator() {
+    return coordinator;
+  }
+
+  int getAttemptNumber() {
+    return attemptNumber;
+  }
+
+  long getSendAfterMs() {
+    return sendAfterMs;
+  }
+}