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,140 @@
+/*
+ * 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.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 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);
+
+    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()) {
+      return;
+    }
+    lastError.set(t);
+    if (runningAttempts.isEmpty() && done.compareAndSet(false, true)) {
+      setException(lastError.get());
+      if (isInQueue.get()) {
+        publisher.removeFromHedgingQueue(this);
+      }
+    }
+  }
+
+  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)) {
+      done.set(true);
+      if (isInQueue.get()) {
+        publisher.removeFromHedgingQueue(this);
+      }
+      for (ApiFuture<PublishResponse> future : runningAttempts.values()) {
+        future.cancel(mayInterruptIfRunning);
+      }
+      return true;
+    }
+    return false;
+  }
+
+  Publisher.OutstandingBatch getBatch() {
+    return batch;
+  }
+}

java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/HedgeSettings.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
+ *
+ *      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(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;
+
+  /** Hedging delay. */
+  private final Duration hedgeDelay;
+
+  /** Maximum tokens. */
+  private final int maxTokens;
+
+  /** Refill rate. */
+  private final float refillRatio;
+
+  private HedgeSettings(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 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 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 refill the refill rate (tokens per successful request), must be 0 < RefillRatio <=
+     *     0.2.
+     * @return this builder.
+     */
+    public Builder setRefillRatio(final float refillRatio) {
+      if (refillRatio <= 0.0f || refillRatio > 0.2f) {
+        throw new IllegalArgumentException(
+            "refillRatio must be greater than 0.0 and less than or equal to 0.2");
+      }
+      this.refillRatio = refillRatio;
+      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/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/OpenTelemetryPubsubTracer.java

@@ -179,6 +179,11 @@ void endPublishBatchingSpan(PubsubMessageWrapper message) {
    * links with the publisher parent span are created for sampled messages in the batch.
    */
   Span startPublishRpcSpan(TopicName topicName, List<PubsubMessageWrapper> messages) {
+    return startPublishRpcSpan(topicName, messages, 1);
+  }
+
+  Span startPublishRpcSpan(
+      TopicName topicName, List<PubsubMessageWrapper> messages, int attemptNumber) {
     if (!enabled) {
       return null;
     }
@@ -203,7 +208,7 @@ Span startPublishRpcSpan(TopicName topicName, List<PubsubMessageWrapper> message
     for (PubsubMessageWrapper message : messages) {
       if (publishRpcSpan.getSpanContext().isSampled()) {
         message.getPublisherSpan().addLink(publishRpcSpan.getSpanContext(), linkAttributes);
-        message.addPublishStartEvent();
+        message.addPublishStartEvent(attemptNumber);
       }
     }
     return publishRpcSpan;