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,196 @@
+/*
+ * 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.api.gax.rpc.ApiException;
+import com.google.common.util.concurrent.MoreExecutors;
+import com.google.pubsub.v1.PublishResponse;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantLock;
+
+/**
+ * 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;
+
+  // Guarded by lock
+  private final Map<Integer, ApiFuture<PublishResponse>> runningAttempts = new HashMap<>();
+  private boolean done = false;
+  private Throwable lastError;
+
+  private final Lock lock = new ReentrantLock();
+  private final AtomicBoolean isInQueue = new AtomicBoolean(false);
+
+  CancellationSharer(final Publisher.OutstandingBatch batch, final Publisher publisher) {
+    this.batch = batch;
+    this.publisher = publisher;
+  }
+
+  void addAttempt(final int attemptNumber, final ApiFuture<PublishResponse> future) {
+    lock.lock();
+    try {
+      if (done) {
+        future.cancel(true);
+        return;
+      }
+      runningAttempts.put(attemptNumber, future);
+    } finally {
+      lock.unlock();
+    }
+
+    ApiFutures.addCallback(
+        future,
+        new ApiFutureCallback<PublishResponse>() {
+          @Override
+          public void onSuccess(final PublishResponse result) {
+            handleAttemptSuccess(attemptNumber, result);
+          }
+
+          @Override
+          public void onFailure(final Throwable t) {
+            handleAttemptFailure(attemptNumber, t);
+          }
+        },
+        MoreExecutors.directExecutor());
+  }
+
+  private void handleAttemptSuccess(final int attemptNumber, final PublishResponse response) {
+    lock.lock();
+    try {
+      if (done) {
+        return;
+      }
+      done = true;
+      batch.successfulAttempt = attemptNumber;
+      set(response);
+      cancelAllExceptLocked(attemptNumber);
+    } finally {
+      lock.unlock();
+    }
+    publisher.refillTokenBucket();
+  }
+
+  private void handleAttemptFailure(final int attemptNumber, final Throwable t) {
+    boolean shouldRemoveFromQueue = false;
+    lock.lock();
+    try {
+      if (done) {
+        return; // <-- Exit early before modifying runningAttempts to avoid
+                // ConcurrentModificationException
+      }
+      runningAttempts.remove(attemptNumber);
+      lastError = t;
+
+      boolean isRetryable = true;
+      if (t instanceof ApiException) {
+        isRetryable =
+            publisher.getRetryableCodes().contains(((ApiException) t).getStatusCode().getCode());
+      }
+
+      if (runningAttempts.isEmpty() || !isRetryable) {
+        done = true;
+        setException(lastError);
+        cancelAllLocked();
+        if (isInQueue.get()) {
+          shouldRemoveFromQueue = true;
+        }
+      }
+    } finally {
+      lock.unlock();
+    }
+
+    if (shouldRemoveFromQueue) {
+      publisher.removeFromHedgingQueue(this);
+    }
+  }
+
+  void checkCompletionOnQueueExit() {
+    lock.lock();
+    try {
+      if (!done && runningAttempts.isEmpty() && !isInQueue.get()) {
+        done = true;
+        setException(
+            lastError != null
+                ? lastError
+                : new RuntimeException("Hedging failed with no active attempts"));
+      }
+    } finally {
+      lock.unlock();
+    }
+  }
+
+  @Override
+  public boolean cancel(final boolean mayInterruptIfRunning) {
+    boolean cancelled = false;
+    boolean shouldRemoveFromQueue = false;
+    lock.lock();
+    try {
+      if (super.cancel(mayInterruptIfRunning)) {
+        cancelled = true;
+        done = true;
+        if (isInQueue.get()) {
+          shouldRemoveFromQueue = true;
+        }
+        cancelAllLocked();
+      }
+    } finally {
+      lock.unlock();
+    }
+
+    if (shouldRemoveFromQueue) {
+      publisher.removeFromHedgingQueue(this);
+    }
+    return cancelled;
+  }
+
+  private void cancelAllLocked() {
+    for (ApiFuture<PublishResponse> future : runningAttempts.values()) {
+      future.cancel(true);
+    }
+    runningAttempts.clear();
+  }
+
+  private void cancelAllExceptLocked(final int successfulAttempt) {
+    runningAttempts.forEach(
+        (attempt, future) -> {
+          if (attempt != successfulAttempt) {
+            future.cancel(true);
+          }
+        });
+    runningAttempts.clear();
+  }
+
+  AtomicBoolean isInQueue() {
+    return isInQueue;
+  }
+
+  Publisher.OutstandingBatch getBatch() {
+    return batch;
+  }
+}

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

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

@@ -0,0 +1,152 @@
+/*
+ * 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 HedgingSettings {
+  /** Default hedging delay. */
+  private static final Duration DEFAULT_DELAY = Duration.ofMillis(1000);
+
+  /** Default maximum number of tokens in the bucket. */
+  private static final int DEFAULT_MAX_TOKENS = 50;
+
+  /** Default refill rate (tokens per successful request). */
+  private static final float DEFAULT_REFILL_RATIO = 0.1f;
+
+  /** Minimum refill rate. */
+  private static final float MIN_REFILL_RATIO = 0.001f;
+
+  /** Maximum refill rate. */
+  private static final float MAX_REFILL_RATIO = 0.2f;
+
+  /** Hedging delay. */
+  private final Duration hedgeDelay;
+
+  /** Maximum tokens. */
+  private final int maxTokens;
+
+  /** Refill rate. */
+  private final float refillRatio;
+
+  private HedgingSettings(final Builder builder) {
+    this.hedgeDelay = builder.hedgeDelay;
+    this.maxTokens = builder.maxTokens;
+    this.refillRatio = builder.refillRatio;
+  }
+
+  /**
+   * Returns the configured hedging delay.
+   *
+   * @return the hedging delay.
+   */
+  Duration getHedgeDelay() {
+    return hedgeDelay;
+  }
+
+  int getMaxTokens() {
+    return maxTokens;
+  }
+
+  float getRefillRatio() {
+    return refillRatio;
+  }
+
+  /**
+   * Returns a new builder for {@code HedgingSettings}.
+   *
+   * @return a new builder.
+   */
+  public static Builder newBuilder() {
+    return new Builder();
+  }
+
+  /** Builder for {@code HedgingSettings}. */
+  public static final class Builder {
+    /** Hedging delay. */
+    private Duration hedgeDelay = DEFAULT_DELAY;
+
+    /** Maximum tokens. */
+    private int maxTokens = DEFAULT_MAX_TOKENS;
+
+    /** Refill rate. */
+    private float refillRatio = DEFAULT_REFILL_RATIO;
+
+    private Builder() {}
+
+    /**
+     * Allows hedging delay to be configurable.
+     *
+     * @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.toMillis() < 100 || delay.toMillis() > 10000) {
+        throw new IllegalArgumentException(
+            "hedgeDelay 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 refillRatio the refill rate (tokens per successful request), must be 0.001 <=
+     *     RefillRatio <= 0.2.
+     * @return this builder.
+     */
+    public Builder setRefillRatio(final float refillRatio) {
+      if (refillRatio < MIN_REFILL_RATIO || refillRatio > MAX_REFILL_RATIO) {
+        throw new IllegalArgumentException(
+            "refillRatio must be greater than or equal to "
+                + MIN_REFILL_RATIO
+                + " and less than or equal to "
+                + MAX_REFILL_RATIO);
+      }
+      this.refillRatio = refillRatio;
+      return this;
+    }
+
+    /**
+     * Builds an instance of {@code HedgingSettings}.
+     *
+     * @return the built {@code HedgingSettings} instance.
+     */
+    public HedgingSettings build() {
+      return new HedgingSettings(this);
+    }
+  }
+}