diff --git a/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/CancellationSharer.java b/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/CancellationSharer.java new file mode 100644 index 000000000000..fd6759bddc59 --- /dev/null +++ b/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/CancellationSharer.java @@ -0,0 +1,200 @@ +/* + * 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. + * + *

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 { + private Publisher.OutstandingBatch batch; + private final Publisher publisher; + private final long deadlineMs; + + // Guarded by lock + private final Map> runningAttempts = new HashMap<>(); + private boolean done = false; + private Throwable lastError; + + private final Lock lock = new ReentrantLock(); + private final AtomicBoolean isInQueue = new AtomicBoolean(false); + + private void cleanupLocked() { + runningAttempts.clear(); + this.batch = null; + } + + CancellationSharer( + final Publisher.OutstandingBatch batch, final Publisher publisher, final long deadlineMs) { + this.batch = batch; + this.publisher = publisher; + this.deadlineMs = deadlineMs; + } + + void addAttempt(final int attemptNumber, final ApiFuture future) { + lock.lock(); + try { + if (done) { + future.cancel(true); + return; + } + runningAttempts.put(attemptNumber, future); + } finally { + lock.unlock(); + } + + ApiFutures.addCallback( + future, + new ApiFutureCallback() { + @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; + publisher.refillTokenBucket(); + set(response); + cancelAllExceptLocked(attemptNumber); + cleanupLocked(); + } finally { + lock.unlock(); + } + } + + private void handleAttemptFailure(final int attemptNumber, final Throwable t) { + lock.lock(); + try { + if (done) { + return; + } + 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(); + cleanupLocked(); + } + } finally { + lock.unlock(); + } + } + + 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")); + cleanupLocked(); + } + } finally { + lock.unlock(); + } + } + + @Override + public boolean cancel(final boolean mayInterruptIfRunning) { + boolean cancelled = false; + lock.lock(); + try { + if (super.cancel(mayInterruptIfRunning)) { + cancelled = true; + done = true; + cancelAllLocked(); + cleanupLocked(); + } + } finally { + lock.unlock(); + } + return cancelled; + } + + private void cancelAllLocked() { + for (ApiFuture 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 getBatchIfActive() { + lock.lock(); + try { + return done ? null : batch; + } finally { + lock.unlock(); + } + } + + long getDeadlineMs() { + return deadlineMs; + } +} diff --git a/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/HedgedRequest.java b/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/HedgedRequest.java new file mode 100644 index 000000000000..ad07de9a8a90 --- /dev/null +++ b/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; + } +} diff --git a/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/HedgingSettings.java b/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/HedgingSettings.java new file mode 100644 index 000000000000..d72857766ddf --- /dev/null +++ b/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); + } + } +} diff --git a/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/LoggingUtil.java b/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/LoggingUtil.java index dbc3a5d86e88..12a87e89dd05 100644 --- a/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/LoggingUtil.java +++ b/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/LoggingUtil.java @@ -33,6 +33,7 @@ public final class LoggingUtil { private static final Logger ackNackLogger = Logger.getLogger("ack-nack"); private static final Logger publishBatchLogger = Logger.getLogger("publish-batch"); private static final Logger subscriberStreamsLogger = Logger.getLogger("subscriber-streams"); + private static final Logger publishHedgedLogger = Logger.getLogger("publish-hedged"); public enum SubSystem { SLOW_ACK(slowAckLogger), @@ -43,7 +44,8 @@ public enum SubSystem { SUBSCRIBER_FLOW_CONTROL(subscriberFlowControlLogger), ACK_NACK(ackNackLogger), PUBLISH_BATCH(publishBatchLogger), - SUBSCRIBER_STREAMS(subscriberStreamsLogger); + SUBSCRIBER_STREAMS(subscriberStreamsLogger), + PUBLISH_HEDGED(publishHedgedLogger); private final Logger logger; diff --git a/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/OpenTelemetryPubsubTracer.java b/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/OpenTelemetryPubsubTracer.java index 3de4484586d0..8bbb5f7d426b 100644 --- a/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/OpenTelemetryPubsubTracer.java +++ b/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/OpenTelemetryPubsubTracer.java @@ -118,10 +118,14 @@ void startPublisherSpan(PubsubMessageWrapper message) { } void endPublisherSpan(PubsubMessageWrapper message) { + endPublisherSpan(message, false); + } + + void endPublisherSpan(PubsubMessageWrapper message, boolean wasHedged) { if (!enabled) { return; } - message.endPublisherSpan(); + message.endPublisherSpan(wasHedged); } void setPublisherMessageIdSpanAttribute(PubsubMessageWrapper message, String messageId) { @@ -179,6 +183,11 @@ void endPublishBatchingSpan(PubsubMessageWrapper message) { * links with the publisher parent span are created for sampled messages in the batch. */ Span startPublishRpcSpan(TopicName topicName, List messages) { + return startPublishRpcSpan(topicName, messages, 0); + } + + Span startPublishRpcSpan( + TopicName topicName, List messages, int attemptNumber) { if (!enabled) { return null; } @@ -203,7 +212,7 @@ Span startPublishRpcSpan(TopicName topicName, List message for (PubsubMessageWrapper message : messages) { if (publishRpcSpan.getSpanContext().isSampled()) { message.getPublisherSpan().addLink(publishRpcSpan.getSpanContext(), linkAttributes); - message.addPublishStartEvent(); + message.addPublishStartEvent(attemptNumber); } } return publishRpcSpan; diff --git a/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/Publisher.java b/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/Publisher.java index 56c920bcfdc1..cf4708575027 100644 --- a/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/Publisher.java +++ b/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/Publisher.java @@ -18,11 +18,13 @@ import static com.google.common.util.concurrent.MoreExecutors.directExecutor; +import com.google.api.core.ApiClock; import com.google.api.core.ApiFunction; import com.google.api.core.ApiFuture; import com.google.api.core.ApiFutureCallback; import com.google.api.core.ApiFutures; import com.google.api.core.BetaApi; +import com.google.api.core.CurrentMillisClock; import com.google.api.core.SettableApiFuture; import com.google.api.gax.batching.BatchingSettings; import com.google.api.gax.batching.FlowControlSettings; @@ -45,7 +47,10 @@ import com.google.cloud.pubsub.v1.stub.GrpcPublisherStub; import com.google.cloud.pubsub.v1.stub.PublisherStub; import com.google.cloud.pubsub.v1.stub.PublisherStubSettings; +import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Preconditions; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; import com.google.protobuf.CodedOutputStream; import com.google.pubsub.v1.PublishRequest; import com.google.pubsub.v1.PublishResponse; @@ -59,18 +64,22 @@ import java.io.IOException; import java.time.Duration; import java.util.ArrayList; +import java.util.Collections; import java.util.HashMap; import java.util.Iterator; import java.util.LinkedList; import java.util.List; import java.util.Map; +import java.util.Set; import java.util.concurrent.Callable; +import java.util.concurrent.ConcurrentLinkedQueue; import java.util.concurrent.CountDownLatch; import java.util.concurrent.Executor; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.ScheduledFuture; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.locks.Lock; import java.util.concurrent.locks.ReentrantLock; import java.util.logging.Level; @@ -122,6 +131,7 @@ public class Publisher implements PublisherInterface { private final AtomicBoolean shutdown; private final BackgroundResource backgroundResources; + private final RetrySettings retrySettings; private final Waiter messagesWaiter; private ScheduledFuture currentAlarmFuture; private final ApiFunction messageTransform; @@ -138,6 +148,26 @@ public class Publisher implements PublisherInterface { private final OpenTelemetry openTelemetry; private OpenTelemetryPubsubTracer tracer = new OpenTelemetryPubsubTracer(null, false); + private final HedgingSettings hedgingSettings; + private final Set retryableCodes; + + /** + * Scale factor to represent decimal token values (e.g. 0.1 refill ratio) as integers inside the + * AtomicInteger token bucket. A scale of 1000 allows representing decimal ratios down to 0.001. + * For example, 1.0 logical token is represented as 1000. + */ + private static final int HEDGE_TOKEN_SCALE = 1000; + + private final AtomicInteger hedgeTokenBucket = new AtomicInteger(); + private int scaledMaxHedgeTokens; + private int scaledHedgeRefillAmount; + private final ApiClock clock; + + private final ConcurrentLinkedQueue hedgingQueue; + private final AtomicBoolean isQueueProcessingScheduled; + private ScheduledFuture queueProcessingFuture; + private final Lock queueLock; + /** The maximum number of messages in one request. Defined by the API. */ public static long getApiMaxRequestElementCount() { return 1000L; @@ -205,6 +235,7 @@ private Publisher(Builder builder) throws IOException { .setTotalTimeoutDuration(Duration.ofNanos(Long.MAX_VALUE)); } + this.retrySettings = retrySettingsBuilder.build(); PublisherStubSettings.Builder stubSettings = PublisherStubSettings.newBuilder() .setCredentialsProvider(builder.credentialsProvider) @@ -223,17 +254,30 @@ private Publisher(Builder builder) throws IOException { StatusCode.Code.RESOURCE_EXHAUSTED, StatusCode.Code.UNKNOWN, StatusCode.Code.UNAVAILABLE) - .setRetrySettings(retrySettingsBuilder.build()) + .setRetrySettings(this.retrySettings) .setBatchingSettings(BatchingSettings.newBuilder().setIsEnabled(false).build()); this.publisherStub = GrpcPublisherStub.create(stubSettings.build()); + this.retryableCodes = ImmutableSet.copyOf(stubSettings.publishSettings().getRetryableCodes()); backgroundResourceList.add(publisherStub); backgroundResources = new BackgroundResourceAggregation(backgroundResourceList); shutdown = new AtomicBoolean(false); messagesWaiter = new Waiter(); + this.hedgingSettings = builder.hedgingSettings; + if (this.hedgingSettings != null) { + this.scaledMaxHedgeTokens = this.hedgingSettings.getMaxTokens() * HEDGE_TOKEN_SCALE; + this.scaledHedgeRefillAmount = + (int) (this.hedgingSettings.getRefillRatio() * HEDGE_TOKEN_SCALE); + this.hedgeTokenBucket.set(0); + } + this.clock = builder.clock != null ? builder.clock : CurrentMillisClock.getDefaultClock(); this.publishContext = GrpcCallContext.createDefault(); this.publishContextWithCompression = GrpcCallContext.createDefault() .withCallOptions(CallOptions.DEFAULT.withCompression(GZIP_COMPRESSION)); + this.hedgingQueue = new ConcurrentLinkedQueue<>(); + this.isQueueProcessingScheduled = new AtomicBoolean(false); + this.queueLock = new ReentrantLock(); + this.queueProcessingFuture = null; } /** Topic which the publisher publishes to. */ @@ -246,6 +290,19 @@ public String getTopicNameString() { return topicName; } + /** Returns the configured hedging settings, or null if hedging is disabled. */ + public HedgingSettings getHedgingSettings() { + return hedgingSettings; + } + + @VisibleForTesting + Float getHedgeTokenBalance() { + if (hedgingSettings == null) { + return null; + } + return (float) hedgeTokenBucket.get() / HEDGE_TOKEN_SCALE; + } + /** * Schedules the publishing of a message. The publishing of the message may occur immediately or * be delayed based on the publisher batching options. @@ -403,6 +460,10 @@ public void run() { * wait for the send operations to complete. To wait for messages to send, call {@code get} on the * futures returned from {@code publish}. */ + Set getRetryableCodes() { + return retryableCodes; + } + public void publishAllOutstanding() { OutstandingBatch unorderedOutstandingBatch = null; messagesBatchLock.lock(); @@ -481,10 +542,32 @@ private void publishAllWithoutInflightForKey(final String orderingKey) { } private ApiFuture publishCall(OutstandingBatch outstandingBatch) { + return publishCall(outstandingBatch, 0, null); + } + + private ApiFuture publishCall( + OutstandingBatch outstandingBatch, int attemptNumber, Duration timeout) { GrpcCallContext context = publishContext; if (enableCompression && outstandingBatch.batchSizeBytes >= compressionBytesThreshold) { context = publishContextWithCompression; } + if (timeout != null) { + context = context.withTimeoutDuration(timeout); + } + if (attemptNumber > 0) { + loggingUtil.logPublisher( + LoggingUtil.SubSystem.PUBLISH_HEDGED, + Level.FINER, + String.format("Publishing hedged attempt %d", attemptNumber), + outstandingBatch.getMessageWrappers().get(0)); + context = + context + .withExtraHeaders( + ImmutableMap.of( + "x-goog-pubsub-hedged-count", + Collections.singletonList(Integer.toString(attemptNumber)))) + .withRetryableCodes(Collections.emptySet()); + } int numMessagesInBatch = outstandingBatch.size(); List pubsubMessagesList = new ArrayList(numMessagesInBatch); @@ -494,7 +577,8 @@ private ApiFuture publishCall(OutstandingBatch outstandingBatch pubsubMessagesList.add(messageWrapper.getPubsubMessage()); } - outstandingBatch.publishRpcSpan = tracer.startPublishRpcSpan(topicNameObject, messageWrappers); + outstandingBatch.publishRpcSpan = + tracer.startPublishRpcSpan(topicNameObject, messageWrappers, attemptNumber); return publisherStub .publishCallable() @@ -572,7 +656,11 @@ public void onFailure(Throwable t) { ApiFuture future; Executor callbackExecutor = directExecutor(); if (outstandingBatch.orderingKey == null || outstandingBatch.orderingKey.isEmpty()) { - future = publishCall(outstandingBatch); + if (hedgingSettings != null) { + future = startHedgedCall(outstandingBatch); + } else { + future = publishCall(outstandingBatch); + } } else { // If ordering key is specified, publish the batch using the sequential executor. future = @@ -588,8 +676,141 @@ public ApiFuture call() { ApiFutures.addCallback(future, futureCallback, callbackExecutor); } - private final class OutstandingBatch { + void refillTokenBucket() { + if (hedgingSettings != null) { + hedgeTokenBucket.accumulateAndGet( + scaledHedgeRefillAmount, + (current, refill) -> Math.min(scaledMaxHedgeTokens, current + refill)); + } + } + + boolean tryAcquireHedgeToken() { + if (hedgingSettings == null) { + return false; + } + int previous = + hedgeTokenBucket.getAndUpdate( + current -> { + if (current < HEDGE_TOKEN_SCALE) { + return current; + } + return current - HEDGE_TOKEN_SCALE; + }); + return previous >= HEDGE_TOKEN_SCALE; + } + + private ApiFuture startHedgedCall(final OutstandingBatch outstandingBatch) { + long deadlineMs = clock.millisTime() + retrySettings.getTotalTimeoutDuration().toMillis(); + final CancellationSharer coordinator = + new CancellationSharer(outstandingBatch, this, deadlineMs); + + // Register cancellation listeners on client futures to propagate cancel to coordinator + final AtomicInteger cancelledCount = new AtomicInteger(0); + final int batchSize = outstandingBatch.outstandingPublishes.size(); + for (final OutstandingPublish outstanding : outstandingBatch.outstandingPublishes) { + outstanding.publishResult.addListener( + new Runnable() { + @Override + public void run() { + if (outstanding.publishResult.isCancelled()) { + if (cancelledCount.incrementAndGet() == batchSize) { + coordinator.cancel(true); + } + } + } + }, + directExecutor()); + } + + ApiFuture firstAttemptFuture = publishCall(outstandingBatch); + coordinator.addAttempt(0, firstAttemptFuture); + long delayMs = hedgingSettings.getHedgeDelay().toMillis(); + HedgedRequest item = new HedgedRequest(coordinator, 1, clock.millisTime() + delayMs); + hedgingQueue.add(item); + coordinator.isInQueue().set(true); + scheduleQueueProcessing(); + + return coordinator; + } + + private void scheduleQueueProcessing() { + if (isQueueProcessingScheduled.compareAndSet(false, true)) { + HedgedRequest nextItem = hedgingQueue.peek(); + if (nextItem == null) { + isQueueProcessingScheduled.set(false); + return; + } + + long delay = Math.max(0, nextItem.getSendAfterMs() - clock.millisTime()); + + queueProcessingFuture = + executor.schedule( + new Runnable() { + @Override + public void run() { + processQueue(); + } + }, + delay, + TimeUnit.MILLISECONDS); + } + } + + private void processQueue() { + queueLock.lock(); + try { + long now = clock.millisTime(); + + HedgedRequest item; + while ((item = hedgingQueue.peek()) != null && item.getSendAfterMs() <= now) { + hedgingQueue.poll(); + + CancellationSharer coordinator = item.getCoordinator(); + OutstandingBatch batch = coordinator.getBatchIfActive(); + if (batch == null) { + coordinator.isInQueue().set(false); + continue; + } + + long remainingTimeoutMs = coordinator.getDeadlineMs() - clock.millisTime(); + if (remainingTimeoutMs <= 0) { + coordinator.isInQueue().set(false); + coordinator.checkCompletionOnQueueExit(); + continue; + } + long attemptTimeoutMs = Math.min(10000, remainingTimeoutMs); + + if (tryAcquireHedgeToken()) { + // Clone and schedule next attempt check (Attempt + 1) + long delayMs = hedgingSettings.getHedgeDelay().toMillis(); + HedgedRequest nextItem = + new HedgedRequest(coordinator, item.getAttemptNumber() + 1, now + delayMs); + hedgingQueue.add(nextItem); + + // Start Hedged Attempt + ApiFuture hedgedFuture = + publishCall(batch, item.getAttemptNumber(), Duration.ofMillis(attemptTimeoutMs)); + coordinator.addAttempt(item.getAttemptNumber(), hedgedFuture); + } else { + loggingUtil.logPublisher( + LoggingUtil.SubSystem.PUBLISH_HEDGED, + Level.FINER, + "Hedging rate limited due to lack of tokens.", + batch.getMessageWrappers().get(0)); + coordinator.isInQueue().set(false); + coordinator.checkCompletionOnQueueExit(); + } + } + isQueueProcessingScheduled.set(false); + scheduleQueueProcessing(); + } finally { + queueLock.unlock(); + } + } + + final class OutstandingBatch { final List outstandingPublishes; + int successfulAttempt = 0; final long creationTime; int attempt; int batchSizeBytes; @@ -600,7 +821,7 @@ private final class OutstandingBatch { List outstandingPublishes, int batchSizeBytes, String orderingKey) { this.outstandingPublishes = outstandingPublishes; attempt = 1; - creationTime = System.currentTimeMillis(); + creationTime = clock.millisTime(); this.batchSizeBytes = batchSizeBytes; this.orderingKey = orderingKey; } @@ -631,6 +852,7 @@ private void onFailure(Throwable t) { private void onSuccess(Iterable results) { tracer.endPublishRpcSpan(publishRpcSpan); + boolean wasHedged = successfulAttempt > 0; Iterator messagesResultsIt = outstandingPublishes.iterator(); for (String messageId : results) { @@ -640,7 +862,7 @@ private void onSuccess(Iterable results) { } nextPublish.publishResult.set(messageId); tracer.setPublisherMessageIdSpanAttribute(nextPublish.messageWrapper, messageId); - tracer.endPublisherSpan(nextPublish.messageWrapper); + tracer.endPublisherSpan(nextPublish.messageWrapper, wasHedged); } } } @@ -677,6 +899,9 @@ public void shutdown() { if (currentAlarmFuture != null && activeAlarm.getAndSet(false)) { currentAlarmFuture.cancel(false); } + if (queueProcessingFuture != null) { + queueProcessingFuture.cancel(false); + } publishAllOutstanding(); messagesWaiter.waitComplete(); backgroundResources.shutdown(); @@ -814,6 +1039,8 @@ public PubsubMessage apply(PubsubMessage input) { private boolean enableOpenTelemetryTracing = false; private OpenTelemetry openTelemetry = null; + private HedgingSettings hedgingSettings = null; + ApiClock clock = null; private Builder(String topic) { this.topicName = Preconditions.checkNotNull(topic); @@ -966,12 +1193,38 @@ public Builder setOpenTelemetry(OpenTelemetry openTelemetry) { return this; } + /** Configures the Publisher's hedging parameters. */ + public Builder setHedgingSettings(HedgingSettings hedgingSettings) { + this.hedgingSettings = hedgingSettings; + return this; + } + + Builder setClock(ApiClock clock) { + this.clock = clock; + return this; + } + /** Returns the default BatchingSettings used by the client if settings are not provided. */ public static BatchingSettings getDefaultBatchingSettings() { return DEFAULT_BATCHING_SETTINGS; } public Publisher build() throws IOException { + Preconditions.checkState( + !(enableMessageOrdering && hedgingSettings != null), + "Publish hedging and message ordering cannot be enabled at the same time."); + if (hedgingSettings != null) { + Duration hedgeDelay = hedgingSettings.getHedgeDelay(); + Duration initialRpcTimeout = retrySettings.getInitialRpcTimeoutDuration(); + if (hedgeDelay.compareTo(initialRpcTimeout) >= 0) { + throw new IllegalArgumentException( + "hedgeDelay (" + + hedgeDelay.toMillis() + + "ms) must be strictly less than the initial RPC timeout duration (" + + initialRpcTimeout.toMillis() + + "ms)"); + } + } return new Publisher(this); } } diff --git a/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/PubsubMessageWrapper.java b/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/PubsubMessageWrapper.java index 19864a26f5a1..0db40fa89c73 100644 --- a/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/PubsubMessageWrapper.java +++ b/java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/PubsubMessageWrapper.java @@ -43,6 +43,8 @@ public class PubsubMessageWrapper { private static final String PUBLISH_START_EVENT = "publish start"; private static final String PUBLISH_END_EVENT = "publish end"; + private static final String HEDGED_PUBLISH_START_EVENT = "publish start (hedged)"; + private static final String HEDGED_PUBLISH_END_EVENT = "publish end (hedged)"; private static final String MODACK_START_EVENT = "modack start"; private static final String MODACK_END_EVENT = "modack end"; @@ -183,8 +185,20 @@ void setSubscribeProcessSpan(Span span) { /** Creates a publish start event that is tied to the publish RPC span time. */ void addPublishStartEvent() { + addPublishStartEvent(0); + } + + /** + * Creates a publish start event that is tied to the publish RPC span time, marking hedged + * attempts explicitly. + */ + void addPublishStartEvent(int attemptNumber) { if (publisherSpan != null) { - publisherSpan.addEvent(PUBLISH_START_EVENT); + if (attemptNumber > 0) { + publisherSpan.addEvent(HEDGED_PUBLISH_START_EVENT); + } else { + publisherSpan.addEvent(PUBLISH_START_EVENT); + } } } @@ -200,8 +214,17 @@ void setPublisherMessageIdSpanAttribute(String messageId) { /** Ends the publisher parent span if it exists. */ void endPublisherSpan() { + endPublisherSpan(false); + } + + /** Ends the publisher parent span if it exists, marking if the operation finished via a hedge. */ + void endPublisherSpan(boolean wasHedged) { if (publisherSpan != null) { - publisherSpan.addEvent(PUBLISH_END_EVENT); + if (wasHedged) { + publisherSpan.addEvent(HEDGED_PUBLISH_END_EVENT); + } else { + publisherSpan.addEvent(PUBLISH_END_EVENT); + } publisherSpan.end(); } } diff --git a/java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/FakePublisherServiceImpl.java b/java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/FakePublisherServiceImpl.java index 9ab1dec73471..a247c240d650 100644 --- a/java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/FakePublisherServiceImpl.java +++ b/java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/FakePublisherServiceImpl.java @@ -20,6 +20,7 @@ import com.google.pubsub.v1.PublishRequest; import com.google.pubsub.v1.PublishResponse; import com.google.pubsub.v1.PublisherGrpc.PublisherImplBase; +import io.grpc.Metadata; import io.grpc.stub.StreamObserver; import java.time.Duration; import java.util.ArrayList; @@ -36,6 +37,7 @@ class FakePublisherServiceImpl extends PublisherImplBase { private final LinkedBlockingQueue requests = new LinkedBlockingQueue<>(); + private final LinkedBlockingQueue capturedHeaders = new LinkedBlockingQueue<>(); private final LinkedBlockingQueue publishResponses = new LinkedBlockingQueue<>(); private final AtomicInteger nextMessageId = new AtomicInteger(1); private boolean autoPublishResponse; @@ -81,7 +83,6 @@ public String toString() { @Override public void publish( PublishRequest request, final StreamObserver responseObserver) { - requests.add(request); Response response; try { if (autoPublishResponse) { @@ -97,6 +98,7 @@ public void publish( throw new IllegalArgumentException(e); } if (responseDelay == Duration.ZERO) { + requests.add(request); sendResponse(response, responseObserver); } else { final Response responseToSend = response; @@ -109,6 +111,7 @@ public void run() { }, responseDelay.toMillis(), TimeUnit.MILLISECONDS); + requests.add(request); } } @@ -160,4 +163,17 @@ public FakePublisherServiceImpl addPublishError(Throwable error) { public List getCapturedRequests() { return new ArrayList(requests); } + + public void recordHeaders(Metadata headers) { + capturedHeaders.add(headers); + } + + public List getCapturedHeaders() { + return new ArrayList<>(capturedHeaders); + } + + public void clearRequests() { + requests.clear(); + capturedHeaders.clear(); + } } diff --git a/java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/HedgingSettingsTest.java b/java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/HedgingSettingsTest.java new file mode 100644 index 000000000000..abcc9a821db2 --- /dev/null +++ b/java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/HedgingSettingsTest.java @@ -0,0 +1,127 @@ +/* + * 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 HedgingSettingsTest { + + @Test + public void testDefaultSettings() { + HedgingSettings settings = HedgingSettings.newBuilder().build(); + assertNotNull(settings); + assertEquals(Duration.ofMillis(1000), settings.getHedgeDelay()); + assertEquals(50, settings.getMaxTokens()); + assertEquals(0.1f, settings.getRefillRatio(), 0.0001f); + } + + @Test + public void testCustomDelay() { + Duration customDelay = Duration.ofMillis(200); + HedgingSettings settings = HedgingSettings.newBuilder().setHedgeDelay(customDelay).build(); + assertNotNull(settings); + assertEquals(customDelay, settings.getHedgeDelay()); + } + + @Test + public void testDelayTooSmallThrows() { + assertThrows( + IllegalArgumentException.class, + () -> HedgingSettings.newBuilder().setHedgeDelay(Duration.ofMillis(99))); + } + + @Test + public void testDelayTooLargeThrows() { + assertThrows( + IllegalArgumentException.class, + () -> HedgingSettings.newBuilder().setHedgeDelay(Duration.ofMillis(10001))); + } + + @Test + public void testNullDelayThrows() { + assertThrows( + NullPointerException.class, () -> HedgingSettings.newBuilder().setHedgeDelay(null)); + } + + @Test + public void testCustomMaxTokens() { + HedgingSettings settings = HedgingSettings.newBuilder().setMaxTokens(10).build(); + assertEquals(10, settings.getMaxTokens()); + } + + @Test + public void testNegativeMaxTokensThrows() { + assertThrows( + IllegalArgumentException.class, () -> HedgingSettings.newBuilder().setMaxTokens(-5)); + } + + @Test + public void testZeroMaxTokensThrows() { + assertThrows( + IllegalArgumentException.class, () -> HedgingSettings.newBuilder().setMaxTokens(0)); + } + + @Test + public void testMaxTokensTooLargeThrows() { + assertThrows( + IllegalArgumentException.class, () -> HedgingSettings.newBuilder().setMaxTokens(251)); + } + + @Test + public void testCustomRefill() { + HedgingSettings settings = HedgingSettings.newBuilder().setRefillRatio(0.15f).build(); + assertEquals(0.15f, settings.getRefillRatio(), 0.0001f); + } + + @Test + public void testNegativeRefillThrows() { + assertThrows( + IllegalArgumentException.class, () -> HedgingSettings.newBuilder().setRefillRatio(-0.1f)); + } + + @Test + public void testZeroRefillThrows() { + assertThrows( + IllegalArgumentException.class, () -> HedgingSettings.newBuilder().setRefillRatio(0.0f)); + } + + @Test + public void testRefillTooLargeThrows() { + assertThrows( + IllegalArgumentException.class, () -> HedgingSettings.newBuilder().setRefillRatio(0.21f)); + } + + @Test + public void testRefillTooSmallThrows() { + assertThrows( + IllegalArgumentException.class, () -> HedgingSettings.newBuilder().setRefillRatio(0.0009f)); + } + + @Test + public void testMinimumRefillValid() { + HedgingSettings settings = HedgingSettings.newBuilder().setRefillRatio(0.001f).build(); + assertEquals(0.001f, settings.getRefillRatio(), 0.0001f); + } +} diff --git a/java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/OpenTelemetryTest.java b/java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/OpenTelemetryTest.java index 52351ddef466..da068c794379 100644 --- a/java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/OpenTelemetryTest.java +++ b/java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/OpenTelemetryTest.java @@ -57,6 +57,7 @@ public class OpenTelemetryTest { private static final String PUBLISH_BATCHING_SPAN_NAME = "publisher batching"; private static final String PUBLISH_RPC_SPAN_NAME = FULL_TOPIC_NAME.getTopic() + " publish"; private static final String PUBLISH_START_EVENT = "publish start"; + private static final String HEDGED_PUBLISH_START_EVENT = "publish start (hedged)"; private static final String PUBLISH_END_EVENT = "publish end"; private static final String SUBSCRIBER_SPAN_NAME = @@ -656,6 +657,55 @@ public void testSubscribeRpcSpanFailures() { .hasEnded(); } + @Test + public void testHedgedPublishSpanEvents() { + PubsubMessage message = getPubsubMessage(); + PubsubMessageWrapper messageWrapper = + PubsubMessageWrapper.newBuilder(message, FULL_TOPIC_NAME).build(); + List messageWrappers = + java.util.Collections.singletonList(messageWrapper); + + Tracer openTelemetryTracer = openTelemetryTesting.getOpenTelemetry().getTracer("test"); + OpenTelemetryPubsubTracer tracer = new OpenTelemetryPubsubTracer(openTelemetryTracer, true); + + // Start Publisher span + tracer.startPublisherSpan(messageWrapper); + + // Original Attempt 0 + Span publishRpcSpan1 = tracer.startPublishRpcSpan(FULL_TOPIC_NAME, messageWrappers, 0); + tracer.endPublishRpcSpan(publishRpcSpan1); + + // Hedged Attempt 1 + Span publishRpcSpan2 = tracer.startPublishRpcSpan(FULL_TOPIC_NAME, messageWrappers, 1); + tracer.endPublishRpcSpan(publishRpcSpan2); + + // End Publisher span + tracer.endPublisherSpan(messageWrapper); + + List allSpans = openTelemetryTesting.getSpans(); + // 3 Spans: publishRpcSpan1, publishRpcSpan2, publisherSpan + assertEquals(3, allSpans.size()); + SpanData publisherSpanData = allSpans.get(2); + + // The publisher parent span should have 3 events: + // 1. "publish start" (from attempt 1) + // 2. "publish start (hedged)" (from attempt 2) + // 3. "publish end" (when publisher span ends) + assertEquals(3, publisherSpanData.getEvents().size()); + + EventDataAssert startEvent1Assert = + OpenTelemetryAssertions.assertThat(publisherSpanData.getEvents().get(0)); + startEvent1Assert.hasName(PUBLISH_START_EVENT); + + EventDataAssert startEvent2Assert = + OpenTelemetryAssertions.assertThat(publisherSpanData.getEvents().get(1)); + startEvent2Assert.hasName(HEDGED_PUBLISH_START_EVENT); + + EventDataAssert endEventAssert = + OpenTelemetryAssertions.assertThat(publisherSpanData.getEvents().get(2)); + endEventAssert.hasName(PUBLISH_END_EVENT); + } + private PubsubMessage getPubsubMessage() { return PubsubMessage.newBuilder() .setData(ByteString.copyFromUtf8("test-data")) diff --git a/java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/PublisherImplTest.java b/java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/PublisherImplTest.java index 8e6efaf372c9..b10c9e292563 100644 --- a/java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/PublisherImplTest.java +++ b/java-pubsub/google-cloud-pubsub/src/test/java/com/google/cloud/pubsub/v1/PublisherImplTest.java @@ -35,6 +35,8 @@ import com.google.api.gax.grpc.testing.LocalChannelProvider; import com.google.api.gax.rpc.DataLossException; import com.google.api.gax.rpc.FixedTransportChannelProvider; +import com.google.api.gax.rpc.InvalidArgumentException; +import com.google.api.gax.rpc.PermissionDeniedException; import com.google.api.gax.rpc.TransportChannelProvider; import com.google.cloud.pubsub.v1.Publisher.Builder; import com.google.protobuf.ByteString; @@ -43,7 +45,12 @@ import com.google.pubsub.v1.PublishResponse; import com.google.pubsub.v1.PubsubMessage; import io.grpc.ManagedChannel; +import io.grpc.Metadata; import io.grpc.Server; +import io.grpc.ServerCall; +import io.grpc.ServerCallHandler; +import io.grpc.ServerInterceptor; +import io.grpc.ServerInterceptors; import io.grpc.Status; import io.grpc.StatusException; import io.grpc.inprocess.InProcessChannelBuilder; @@ -98,13 +105,25 @@ public class PublisherImplTest { public void setUp() throws Exception { testPublisherServiceImpl = new FakePublisherServiceImpl(); + ServerInterceptor headerInterceptor = + new ServerInterceptor() { + @Override + public ServerCall.Listener interceptCall( + ServerCall call, Metadata headers, ServerCallHandler next) { + testPublisherServiceImpl.recordHeaders(headers); + return next.startCall(call, headers); + } + }; + InProcessServerBuilder serverBuilder = InProcessServerBuilder.forName("test-server"); - serverBuilder.addService(testPublisherServiceImpl); + serverBuilder.addService( + ServerInterceptors.intercept(testPublisherServiceImpl, headerInterceptor)); testServer = serverBuilder.build(); testChannel = InProcessChannelBuilder.forName("test-server").build(); testServer.start(); fakeExecutor = new FakeScheduledExecutorService(); + testPublisherServiceImpl.setExecutor(fakeExecutor); } @After @@ -1339,6 +1358,369 @@ public void testPublishOpenTelemetryTracing() throws Exception { .hasEnded(); } + @Test + public void testPublisherWithHedgingSettings() throws Exception { + HedgingSettings hedgingSettings = + HedgingSettings.newBuilder().setHedgeDelay(Duration.ofMillis(100)).build(); + Publisher publisher = getTestPublisherBuilder().setHedgingSettings(hedgingSettings).build(); + + assertThat(publisher.getHedgingSettings()).isEqualTo(hedgingSettings); + assertThat(publisher.getHedgeTokenBalance()).isNotNull(); + assertThat(publisher.getHedgeTokenBalance()).isWithin(0.0001f).of(0.0f); + + shutdownTestPublisher(publisher); + } + + @Test + public void testPublisherThrowsIfHedgeDelayGtRpcTimeout() throws Exception { + HedgingSettings hedgingSettings = + HedgingSettings.newBuilder().setHedgeDelay(Duration.ofMillis(500)).build(); + com.google.api.gax.retrying.RetrySettings retrySettings = + com.google.api.gax.retrying.RetrySettings.newBuilder() + .setInitialRpcTimeoutDuration(Duration.ofMillis(400)) + .setMaxRpcTimeoutDuration(Duration.ofMillis(400)) + .setTotalTimeoutDuration(Duration.ofSeconds(10)) + .build(); + + try { + getTestPublisherBuilder() + .setHedgingSettings(hedgingSettings) + .setRetrySettings(retrySettings) + .build(); + fail( + "Should have thrown IllegalArgumentException because hedgeDelay (500ms) > RPC timeout (400ms)"); + } catch (IllegalArgumentException e) { + assertThat(e.getMessage()) + .contains("must be strictly less than the initial RPC timeout duration"); + } + } + + @Test + public void testPublisherThrowsIfHedgeDelayEqRpcTimeout() throws Exception { + HedgingSettings hedgingSettings = + HedgingSettings.newBuilder().setHedgeDelay(Duration.ofMillis(500)).build(); + com.google.api.gax.retrying.RetrySettings retrySettings = + com.google.api.gax.retrying.RetrySettings.newBuilder() + .setInitialRpcTimeoutDuration(Duration.ofMillis(500)) + .setMaxRpcTimeoutDuration(Duration.ofMillis(500)) + .setTotalTimeoutDuration(Duration.ofSeconds(10)) + .build(); + + try { + getTestPublisherBuilder() + .setHedgingSettings(hedgingSettings) + .setRetrySettings(retrySettings) + .build(); + fail( + "Should have thrown IllegalArgumentException because hedgeDelay (500ms) == RPC timeout (500ms)"); + } catch (IllegalArgumentException e) { + assertThat(e.getMessage()) + .contains("must be strictly less than the initial RPC timeout duration"); + } + } + + @Test + public void testPublisherWithoutHedgingSettings() throws Exception { + Publisher publisher = getTestPublisherBuilder().build(); + + assertThat(publisher.getHedgingSettings()).isNull(); + assertThat(publisher.getHedgeTokenBalance()).isNull(); + + shutdownTestPublisher(publisher); + } + + private Publisher getPublisherWithHedge(Duration delay) throws Exception { + return getPublisherWithHedge(delay, 0.1f, 20); + } + + private Publisher getPublisherWithHedge(Duration delay, float refillRatio, int maxTokens) + throws Exception { + HedgingSettings hedgingSettings = + HedgingSettings.newBuilder() + .setHedgeDelay(delay) + .setRefillRatio(refillRatio) + .setMaxTokens(maxTokens) + .build(); + return getTestPublisherBuilder() + .setHedgingSettings(hedgingSettings) + .setClock(fakeExecutor.getClock()) + .setBatchingSettings( + Publisher.Builder.DEFAULT_BATCHING_SETTINGS.toBuilder() + .setElementCountThreshold(1L) + .build()) + .build(); + } + + private void fillTokenBucket(Publisher publisher, int tokensToFill) throws Exception { + testPublisherServiceImpl.setAutoPublishResponse(true); + for (int i = 0; i < tokensToFill; i++) { + ApiFuture future = sendTestMessage(publisher, "warmup-msg-" + i); + future.get(); + } + testPublisherServiceImpl.clearRequests(); + } + + private void waitForRequests(FakePublisherServiceImpl service, int expectedCount) + throws InterruptedException { + long timeout = System.currentTimeMillis() + 5000; + while (service.getCapturedRequests().size() < expectedCount + && System.currentTimeMillis() < timeout) { + Thread.sleep(5); + } + if (service.getCapturedRequests().size() < expectedCount) { + throw new AssertionError( + String.format( + "Timed out waiting for requests. Expected: %d, Got: %d", + expectedCount, service.getCapturedRequests().size())); + } + } + + @Test + public void testTokenBucketRefillRate() throws Exception { + Publisher publisher = getPublisherWithHedge(Duration.ofMillis(100), 0.125f, 10); + // Starts at 0 + assertThat(publisher.getHedgeTokenBalance()).isEqualTo(0.0f); + + // Warm up 1 message (should succeed and refill by 0.125) + testPublisherServiceImpl.setAutoPublishResponse(true); + sendTestMessage(publisher, "refill-warmup").get(); + + // Balance should be exactly 0.125 + assertThat(publisher.getHedgeTokenBalance()).isWithin(0.0001f).of(0.125f); + shutdownTestPublisher(publisher); + } + + @Test + public void testHedgingNotTriggeredIfFast() throws Exception { + Publisher publisher = getPublisherWithHedge(Duration.ofMillis(100)); + + // Prepare fast response (10ms delay) + testPublisherServiceImpl.setAutoPublishResponse(false); + testPublisherServiceImpl.setPublishResponseDelay(Duration.ofMillis(10)); + testPublisherServiceImpl.addPublishResponse(PublishResponse.newBuilder().addMessageIds("1")); + + ApiFuture future = sendTestMessage(publisher, "msg-fast"); + waitForRequests(testPublisherServiceImpl, 1); + + // Advance time past response but before hedge delay (e.g. 50ms) + fakeExecutor.advanceTime(Duration.ofMillis(50)); + + // Future should be completed + assertEquals("1", future.get()); + + // Only 1 request should be received by server + assertThat(testPublisherServiceImpl.getCapturedRequests()).hasSize(1); + + shutdownTestPublisher(publisher); + } + + @Test + public void testHedgingTriggeredIfSlow() throws Exception { + Publisher publisher = getPublisherWithHedge(Duration.ofMillis(100), 0.2f, 20); + fillTokenBucket(publisher, 5); + + // Set response delay to 200ms (greater than 100ms hedge delay) + testPublisherServiceImpl.setAutoPublishResponse(false); + testPublisherServiceImpl.setPublishResponseDelay(Duration.ofMillis(200)); + // Add two responses (one for main, one for hedge) + testPublisherServiceImpl.addPublishResponse(PublishResponse.newBuilder().addMessageIds("1")); + testPublisherServiceImpl.addPublishResponse(PublishResponse.newBuilder().addMessageIds("2")); + + ApiFuture future = sendTestMessage(publisher, "msg-slow"); + waitForRequests(testPublisherServiceImpl, 1); + + // Advance time to 80ms (before hedge delay) + fakeExecutor.advanceTime(Duration.ofMillis(80)); + assertThat(testPublisherServiceImpl.getCapturedRequests()).hasSize(1); // Only original sent + + // Advance time to 120ms (past 100ms hedge delay) + fakeExecutor.advanceTime(Duration.ofMillis(40)); + waitForRequests(testPublisherServiceImpl, 2); + + // Now attempt 2 should have been triggered + assertThat(testPublisherServiceImpl.getCapturedRequests()).hasSize(2); + + // Advance to 220ms to let responses complete + fakeExecutor.advanceTime(Duration.ofMillis(100)); + fakeExecutor.advanceTime(Duration.ZERO); // Drain pending tasks + assertEquals("1", future.get(5, TimeUnit.SECONDS)); + + List capturedHeaders = testPublisherServiceImpl.getCapturedHeaders(); + assertThat(capturedHeaders).hasSize(2); + Metadata.Key hedgedHeaderKey = + Metadata.Key.of("x-goog-pubsub-hedged-count", Metadata.ASCII_STRING_MARSHALLER); + // Original request should NOT have the header + assertThat(capturedHeaders.get(0).get(hedgedHeaderKey)).isNull(); + // First hedged request should have value "1" + assertThat(capturedHeaders.get(1).get(hedgedHeaderKey)).isEqualTo("1"); + + shutdownTestPublisher(publisher); + } + + @Test + public void testMultipleHedging() throws Exception { + Publisher publisher = getPublisherWithHedge(Duration.ofMillis(100), 0.2f, 20); + fillTokenBucket(publisher, 10); + + // Set delay to 400ms + testPublisherServiceImpl.setAutoPublishResponse(false); + testPublisherServiceImpl.setPublishResponseDelay(Duration.ofMillis(400)); + // Add responses for 3 attempts + testPublisherServiceImpl.addPublishResponse(PublishResponse.newBuilder().addMessageIds("1")); + testPublisherServiceImpl.addPublishResponse(PublishResponse.newBuilder().addMessageIds("2")); + testPublisherServiceImpl.addPublishResponse(PublishResponse.newBuilder().addMessageIds("3")); + + ApiFuture future = sendTestMessage(publisher, "msg-very-slow"); + waitForRequests(testPublisherServiceImpl, 1); + + // T=0: Attempt 1 sent. + // T=120 (Hedge 1): Attempt 2 sent. + fakeExecutor.advanceTime(Duration.ofMillis(120)); + waitForRequests(testPublisherServiceImpl, 2); + assertThat(testPublisherServiceImpl.getCapturedRequests()).hasSize(2); + + // T=240 (Hedge 2): Attempt 3 sent. + fakeExecutor.advanceTime(Duration.ofMillis(120)); + waitForRequests(testPublisherServiceImpl, 3); + assertThat(testPublisherServiceImpl.getCapturedRequests()).hasSize(3); + + // Advance to complete + fakeExecutor.advanceTime(Duration.ofMillis(200)); + assertEquals("1", future.get(5, TimeUnit.SECONDS)); + + List capturedHeaders = testPublisherServiceImpl.getCapturedHeaders(); + assertThat(capturedHeaders).hasSize(3); + Metadata.Key hedgedHeaderKey = + Metadata.Key.of("x-goog-pubsub-hedged-count", Metadata.ASCII_STRING_MARSHALLER); + assertThat(capturedHeaders.get(0).get(hedgedHeaderKey)).isNull(); + assertThat(capturedHeaders.get(1).get(hedgedHeaderKey)).isEqualTo("1"); + assertThat(capturedHeaders.get(2).get(hedgedHeaderKey)).isEqualTo("2"); + + shutdownTestPublisher(publisher); + } + + @Test + public void testHedgingBypassedIfNoTokens() throws Exception { + Publisher publisher = getPublisherWithHedge(Duration.ofMillis(100)); + + // Drain the token bucket completely (since it starts full) + while (publisher.tryAcquireHedgeToken()) {} + assertThat(publisher.getHedgeTokenBalance()).isEqualTo(0.0f); + + testPublisherServiceImpl.setPublishResponseDelay(Duration.ofMillis(200)); + testPublisherServiceImpl.addPublishResponse(PublishResponse.newBuilder().addMessageIds("1")); + + ApiFuture future = sendTestMessage(publisher, "msg-slow-no-tokens"); + waitForRequests(testPublisherServiceImpl, 1); + + // Advance past hedge delay + fakeExecutor.advanceTime(Duration.ofMillis(120)); + + // Should NOT trigger hedge because token bucket is empty + assertThat(testPublisherServiceImpl.getCapturedRequests()).hasSize(1); + + fakeExecutor.advanceTime(Duration.ofMillis(100)); + assertEquals("1", future.get(5, TimeUnit.SECONDS)); + + shutdownTestPublisher(publisher); + } + + @Test + public void testHedgingCancellationPropagates() throws Exception { + Publisher publisher = getPublisherWithHedge(Duration.ofMillis(100), 0.2f, 20); + fillTokenBucket(publisher, 5); + + testPublisherServiceImpl.setAutoPublishResponse(false); + testPublisherServiceImpl.setPublishResponseDelay(Duration.ofMillis(200)); + testPublisherServiceImpl.addPublishResponse(PublishResponse.newBuilder().addMessageIds("1")); + testPublisherServiceImpl.addPublishResponse(PublishResponse.newBuilder().addMessageIds("2")); + + ApiFuture future = sendTestMessage(publisher, "msg-cancel"); + waitForRequests(testPublisherServiceImpl, 1); + + // Trigger hedge + fakeExecutor.advanceTime(Duration.ofMillis(120)); + waitForRequests(testPublisherServiceImpl, 2); + assertThat(testPublisherServiceImpl.getCapturedRequests()).hasSize(2); + + // Cancel the future + future.cancel(true); + + // Verify cancellation propagates to overall future + assertTrue(future.isCancelled()); + + shutdownTestPublisher(publisher); + } + + @Test + public void testNoHedgingIfOriginalFailsImmediately() throws Exception { + Publisher publisher = getPublisherWithHedge(Duration.ofMillis(100), 0.2f, 20); + fillTokenBucket(publisher, 5); + + // Configure the fake to immediately return an INVALID_ARGUMENT error + testPublisherServiceImpl.setAutoPublishResponse(false); + testPublisherServiceImpl.addPublishError(new StatusException(Status.INVALID_ARGUMENT)); + + ApiFuture future = sendTestMessage(publisher, "msg-fail-fast"); + + // The request should fail immediately without waiting or advancing time + try { + future.get(1, TimeUnit.SECONDS); + fail("Should have failed with ExecutionException"); + } catch (ExecutionException e) { + // expected + assertThat(e.getCause()).isInstanceOf(InvalidArgumentException.class); + } + + // Server should receive exactly 1 request (the original attempt) + assertThat(testPublisherServiceImpl.getCapturedRequests()).hasSize(1); + + // Advance time past the 100ms hedge delay and check that no hedge was sent + fakeExecutor.advanceTime(Duration.ofMillis(200)); + + // Captured requests should still be 1 (no hedge triggered) + assertThat(testPublisherServiceImpl.getCapturedRequests()).hasSize(1); + + shutdownTestPublisher(publisher); + } + + @Test + public void testPermanentErrorTerminatesMetaRequest() throws Exception { + Publisher publisher = getPublisherWithHedge(Duration.ofMillis(100), 0.2f, 20); + fillTokenBucket(publisher, 5); + + // 1. Configure the first response to be slow (200ms) + testPublisherServiceImpl.setAutoPublishResponse(false); + testPublisherServiceImpl.setPublishResponseDelay(Duration.ofMillis(200)); + testPublisherServiceImpl.addPublishResponse(PublishResponse.newBuilder().addMessageIds("1")); + + // Start original attempt (Attempt 1) + ApiFuture future = sendTestMessage(publisher, "msg-permanent-error-test"); + waitForRequests(testPublisherServiceImpl, 1); + + // 2. Before the hedge is triggered, change response delay to 0ms (fast) + // and enqueue a PERMISSION_DENIED error for Attempt 2 + fakeExecutor.advanceTime(Duration.ofMillis(50)); + testPublisherServiceImpl.setPublishResponseDelay(Duration.ZERO); + testPublisherServiceImpl.addPublishError(new StatusException(Status.PERMISSION_DENIED)); + + // 3. Advance past the hedge delay to trigger the hedge (Attempt 2) at t=120ms + fakeExecutor.advanceTime(Duration.ofMillis(70)); + waitForRequests(testPublisherServiceImpl, 2); + + // Attempt 2 fails immediately with a permanent error. + // It should fail the client future immediately at t=120ms, + // without waiting for Attempt 1 to complete at t=200ms. + try { + future.get(1, TimeUnit.SECONDS); + fail("Should have failed with ExecutionException"); + } catch (ExecutionException e) { + assertThat(e.getCause()).isInstanceOf(PermissionDeniedException.class); + } + + shutdownTestPublisher(publisher); + } + private Builder getTestPublisherBuilder() { return Publisher.newBuilder(TEST_TOPIC) .setExecutorProvider(FixedExecutorProvider.create(fakeExecutor))