diff --git a/handwritten/pubsub/src/default-options.ts b/handwritten/pubsub/src/default-options.ts index 7f09392d51d1..d499f087cfc0 100644 --- a/handwritten/pubsub/src/default-options.ts +++ b/handwritten/pubsub/src/default-options.ts @@ -55,5 +55,16 @@ export const defaultOptions = { // The maximum time we'll wait to send batched messages, in milliseconds. maxDelayMillis: 10, + + hedging: { + // The default delay before sending a hedged publish RPC. + hedgeDelay: Duration.from({milliseconds: 1000}), + + // The default maximum number of tokens in the hedging token bucket. + maxTokens: 50, + + // The default token refill ratio per successful publish batch. + refillRatio: 0.1, + }, }, }; diff --git a/handwritten/pubsub/src/index.ts b/handwritten/pubsub/src/index.ts index 2e8ae3fff203..6b9064748433 100644 --- a/handwritten/pubsub/src/index.ts +++ b/handwritten/pubsub/src/index.ts @@ -98,7 +98,12 @@ export { TestIamPermissionsCallback, IAM, } from './iam'; -export {Attributes, PublishCallback, PublishOptions} from './publisher'; +export { + Attributes, + HedgingOptions, + PublishCallback, + PublishOptions, +} from './publisher'; export {BatchPublishOptions} from './publisher/message-batch'; export {PublishError} from './publisher/publish-error'; export {FlowControlOptions} from './publisher/flow-control'; diff --git a/handwritten/pubsub/src/publisher/hedging.ts b/handwritten/pubsub/src/publisher/hedging.ts new file mode 100644 index 000000000000..0fd20bc86375 --- /dev/null +++ b/handwritten/pubsub/src/publisher/hedging.ts @@ -0,0 +1,250 @@ +/*! + * 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. + */ + +import {CallOptions} from 'google-gax'; +import {defaultOptions} from '../default-options'; +import {Duration} from '../temporal'; + +/** + * Minimum allowed hedging delay (100ms). + */ +const MIN_HEDGE_DELAY = Duration.from({milliseconds: 100}); + +/** + * Maximum allowed hedging delay (10s). + */ +const MAX_HEDGE_DELAY = Duration.from({seconds: 10}); + +/** + * Minimum allowed maxTokens value. + */ +const MIN_MAX_TOKENS = 1; + +/** + * Maximum allowed maxTokens value. + */ +const MAX_MAX_TOKENS = 250; + +/** + * Minimum allowed refill ratio. + */ +const MIN_REFILL_RATIO = 0.001; + +/** + * Maximum allowed refill ratio. + */ +const MAX_REFILL_RATIO = 0.2; + +/** + * Default initial RPC timeout for Publish calls (from publisher_client_config.json). + */ +const DEFAULT_INITIAL_RPC_TIMEOUT = Duration.from({seconds: 60}); + +/** + * Default total timeout for Publish calls (from publisher_client_config.json). + */ +const DEFAULT_TOTAL_TIMEOUT = Duration.from({seconds: 600}); + +/** + * Scale factor to represent decimal token values (e.g. 0.1 refill ratio) as + * integers inside the token bucket. A scale of 1000 allows representing decimal + * ratios down to 0.001 (1.0 logical token = 1000). + */ +const HEDGE_TOKEN_SCALE = 1000; + +/** + * Settings for configuring publish hedging. + */ +export interface HedgingOptions { + /** + * Delay before sending a hedged (backup) publish RPC. + * Default: 1000ms. Must satisfy 100ms <= hedgeDelay <= 10s. + */ + hedgeDelay?: Duration; + + /** + * Maximum number of tokens in the token bucket rate limiter. + * Default: 50. Must satisfy 0 < maxTokens <= 250 (integer). + */ + maxTokens?: number; + + /** + * Token bucket refill ratio per successful publish batch. + * Default: 0.1. Must satisfy 0.001 <= refillRatio <= 0.2. + */ + refillRatio?: number; +} + +/** + * Validated and defaulted hedging configuration. + */ +export interface ResolvedHedgingOptions { + readonly hedgeDelay: Duration; + readonly maxTokens: number; + readonly refillRatio: number; +} + +/** + * Fixed-point token bucket rate limiter for publish hedging. + * + * Starts empty (0 tokens), refills by `refillRatio` on each successful publish + * batch (capped at `maxTokens`), and requires 1.0 token per hedged attempt. + */ +export class HedgingTokenBucket { + private hedgeTokenBucket = 0; + private readonly scaledMaxHedgeTokens: number; + private readonly scaledHedgeRefillAmount: number; + + constructor(options: ResolvedHedgingOptions) { + this.scaledMaxHedgeTokens = options.maxTokens * HEDGE_TOKEN_SCALE; + this.scaledHedgeRefillAmount = Math.round( + options.refillRatio * HEDGE_TOKEN_SCALE, + ); + } + + /** + * Resolves the effective initial RPC timeout Duration from CallOptions. + */ + static resolveInitialRpcTimeout(gaxOpts?: CallOptions): Duration { + const initialRpcTimeoutMillis = + gaxOpts?.retry?.backoffSettings?.initialRpcTimeoutMillis ?? + gaxOpts?.timeout; + if (typeof initialRpcTimeoutMillis === 'number') { + return Duration.from({milliseconds: initialRpcTimeoutMillis}); + } + return DEFAULT_INITIAL_RPC_TIMEOUT; + } + + /** + * Resolves the effective total timeout Duration from CallOptions. + */ + static resolveTotalTimeout(gaxOpts?: CallOptions): Duration { + const totalTimeoutMillis = + gaxOpts?.retry?.backoffSettings?.totalTimeoutMillis ?? gaxOpts?.timeout; + if (typeof totalTimeoutMillis === 'number') { + return Duration.from({milliseconds: totalTimeoutMillis}); + } + return DEFAULT_TOTAL_TIMEOUT; + } + + /** + * Validates user-supplied HedgingOptions, applies defaults, and verifies + * compatibility with messageOrdering and RPC/total timeouts. + */ + static validateAndResolveHedgingOptions( + hedging?: HedgingOptions, + messageOrdering?: boolean, + gaxOpts?: CallOptions, + ): ResolvedHedgingOptions | undefined { + if (!hedging) { + return undefined; + } + + if (messageOrdering) { + throw new Error( + 'Publish hedging and message ordering cannot be enabled at the same time.', + ); + } + + const hedgeDelay = hedging.hedgeDelay + ? Duration.from(hedging.hedgeDelay) + : defaultOptions.publish.hedging.hedgeDelay; + + if ( + Duration.compare(hedgeDelay, MIN_HEDGE_DELAY) < 0 || + Duration.compare(hedgeDelay, MAX_HEDGE_DELAY) > 0 + ) { + throw new RangeError( + 'hedgeDelay must be greater than or equal to 100ms and less than or equal to 10s', + ); + } + + const maxTokens = + hedging.maxTokens ?? defaultOptions.publish.hedging.maxTokens; + if ( + !Number.isInteger(maxTokens) || + maxTokens < MIN_MAX_TOKENS || + maxTokens > MAX_MAX_TOKENS + ) { + throw new RangeError( + 'maxTokens must be greater than 0 and less than or equal to 250', + ); + } + + const refillRatio = + hedging.refillRatio ?? defaultOptions.publish.hedging.refillRatio; + if ( + typeof refillRatio !== 'number' || + Number.isNaN(refillRatio) || + refillRatio < MIN_REFILL_RATIO || + refillRatio > MAX_REFILL_RATIO + ) { + throw new RangeError( + `refillRatio must be greater than or equal to ${MIN_REFILL_RATIO} and less than or equal to ${MAX_REFILL_RATIO}`, + ); + } + + const initialRpcTimeout = + HedgingTokenBucket.resolveInitialRpcTimeout(gaxOpts); + if (Duration.compare(hedgeDelay, initialRpcTimeout) >= 0) { + throw new RangeError( + `hedgeDelay (${hedgeDelay.milliseconds}ms) must be strictly less than the initial RPC timeout duration (${initialRpcTimeout.milliseconds}ms)`, + ); + } + + const totalTimeout = HedgingTokenBucket.resolveTotalTimeout(gaxOpts); + if (Duration.compare(hedgeDelay, totalTimeout) >= 0) { + throw new RangeError( + `hedgeDelay (${hedgeDelay.milliseconds}ms) must be strictly less than the total timeout duration (${totalTimeout.milliseconds}ms)`, + ); + } + + return { + hedgeDelay, + maxTokens, + refillRatio, + }; + } + + /** + * Refills the token bucket after a successful publish batch. + */ + refillTokenBucket(): void { + this.hedgeTokenBucket = Math.min( + this.scaledMaxHedgeTokens, + this.hedgeTokenBucket + this.scaledHedgeRefillAmount, + ); + } + + /** + * Attempts to consume 1.0 token for a hedged request. + * Returns true if a token was available and consumed, false otherwise. + */ + tryAcquireHedgeToken(): boolean { + if (this.hedgeTokenBucket < HEDGE_TOKEN_SCALE) { + return false; + } + this.hedgeTokenBucket -= HEDGE_TOKEN_SCALE; + return true; + } + + /** + * Returns the current logical token balance (for testing/inspection). + */ + getTokenBalance(): number { + return this.hedgeTokenBucket / HEDGE_TOKEN_SCALE; + } +} diff --git a/handwritten/pubsub/src/publisher/index.ts b/handwritten/pubsub/src/publisher/index.ts index 2046c61699f5..86bc80713ef7 100644 --- a/handwritten/pubsub/src/publisher/index.ts +++ b/handwritten/pubsub/src/publisher/index.ts @@ -26,10 +26,20 @@ import {defaultOptions} from '../default-options'; import * as tracing from '../telemetry-tracing'; import {FlowControl, FlowControlOptions} from './flow-control'; +import { + HedgingOptions, + ResolvedHedgingOptions, + HedgingTokenBucket, +} from './hedging'; import {promisifySome} from '../util'; import {PubsubMessage, Attributes} from './pubsub-message'; export {PubsubMessage, Attributes} from './pubsub-message'; +export { + HedgingOptions, + ResolvedHedgingOptions, + HedgingTokenBucket, +} from './hedging'; export type PublishCallback = RequestCallback; @@ -38,6 +48,7 @@ export interface PublishOptions { flowControlOptions?: FlowControlOptions; gaxOpts?: CallOptions; messageOrdering?: boolean; + hedging?: HedgingOptions; } /** @@ -52,6 +63,7 @@ export interface PublishOptions { * same order key in Message will be delivered to the subscribers in the order in which they * are received by the Pub/Sub system. Otherwise, they may be delivered in * any order. + * @property {HedgingOptions} [hedging] Settings for configuring publish hedging. */ export const BATCH_LIMITS: BatchPublishOptions = { @@ -81,6 +93,8 @@ export class Publisher { queue: Queue; orderedQueues: Map; flowControl: FlowControl; + hedgingOptions?: ResolvedHedgingOptions; + private tokenBucket_?: HedgingTokenBucket; constructor(topic: Topic, options?: PublishOptions) { this.flowControl = new FlowControl( @@ -131,6 +145,7 @@ export class Publisher { .then(() => allDrains) .then(() => { definedCallback(null); + return undefined; }) .catch(definedCallback); } @@ -292,18 +307,32 @@ export class Publisher { options, ); + const resolvedHedging = HedgingTokenBucket.validateAndResolveHedgingOptions( + options?.hedging ?? this.settings?.hedging, + messageOrdering, + gaxOpts, + ); + this.hedgingOptions = resolvedHedging; + this.tokenBucket_ = resolvedHedging + ? new HedgingTokenBucket(resolvedHedging) + : undefined; + this.settings = { batching: { - maxBytes: Math.min(batching!.maxBytes!, BATCH_LIMITS.maxBytes!), + maxBytes: Math.min( + batching?.maxBytes ?? BATCH_LIMITS.maxBytes ?? 0, + BATCH_LIMITS.maxBytes ?? 0, + ), maxMessages: Math.min( - batching!.maxMessages!, - BATCH_LIMITS.maxMessages!, + batching?.maxMessages ?? BATCH_LIMITS.maxMessages ?? 0, + BATCH_LIMITS.maxMessages ?? 0, ), - maxMilliseconds: batching!.maxMilliseconds, + maxMilliseconds: batching?.maxMilliseconds, }, gaxOpts, messageOrdering, flowControlOptions, + ...(resolvedHedging ? {hedging: resolvedHedging} : {}), }; // We also need to let all of our queues know that they need to update their options. @@ -318,7 +347,18 @@ export class Publisher { } // This will always be filled in by our defaults if nothing else. - this.flowControl.setOptions(this.settings.flowControlOptions!); + if (this.settings.flowControlOptions) { + this.flowControl.setOptions(this.settings.flowControlOptions); + } + } + + /** + * Returns the hedging token bucket instance, or undefined if hedging is disabled. + * + * @private + */ + get tokenBucket(): HedgingTokenBucket | undefined { + return this.tokenBucket_; } /** diff --git a/handwritten/pubsub/test/publisher/hedging.ts b/handwritten/pubsub/test/publisher/hedging.ts new file mode 100644 index 000000000000..265bcdabb074 --- /dev/null +++ b/handwritten/pubsub/test/publisher/hedging.ts @@ -0,0 +1,250 @@ +/*! + * 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. + */ + +import * as assert from 'assert'; +import {describe, it} from 'mocha'; +import {defaultOptions} from '../../src/default-options'; +import {Duration} from '../../src/temporal'; +import {HedgingTokenBucket} from '../../src/publisher/hedging'; + +describe('Publisher Hedging (Options & Token Bucket)', () => { + describe('validateAndResolveHedgingOptions', () => { + it('returns undefined when hedging options are not provided', () => { + assert.strictEqual( + HedgingTokenBucket.validateAndResolveHedgingOptions(undefined), + undefined, + ); + }); + + it('applies default settings when an empty hedging object is provided', () => { + const resolved = HedgingTokenBucket.validateAndResolveHedgingOptions({}); + assert.ok(resolved); + assert.strictEqual( + resolved.hedgeDelay.milliseconds, + defaultOptions.publish.hedging.hedgeDelay.milliseconds, + ); + assert.strictEqual( + resolved.maxTokens, + defaultOptions.publish.hedging.maxTokens, + ); + assert.strictEqual( + resolved.refillRatio, + defaultOptions.publish.hedging.refillRatio, + ); + }); + + it('accepts valid custom hedging settings', () => { + const customDelay = Duration.from({milliseconds: 250}); + const resolved = HedgingTokenBucket.validateAndResolveHedgingOptions({ + hedgeDelay: customDelay, + maxTokens: 100, + refillRatio: 0.05, + }); + assert.ok(resolved); + assert.strictEqual(resolved.hedgeDelay.milliseconds, 250); + assert.strictEqual(resolved.maxTokens, 100); + assert.strictEqual(resolved.refillRatio, 0.05); + }); + + it('throws RangeError when hedgeDelay is less than 100ms', () => { + assert.throws( + () => + HedgingTokenBucket.validateAndResolveHedgingOptions({ + hedgeDelay: Duration.from({milliseconds: 99}), + }), + RangeError, + ); + }); + + it('throws RangeError when hedgeDelay is greater than 10s', () => { + assert.throws( + () => + HedgingTokenBucket.validateAndResolveHedgingOptions({ + hedgeDelay: Duration.from({milliseconds: 10001}), + }), + RangeError, + ); + }); + + it('throws RangeError when maxTokens is out of bounds or non-integer', () => { + assert.throws( + () => + HedgingTokenBucket.validateAndResolveHedgingOptions({maxTokens: 0}), + RangeError, + ); + assert.throws( + () => + HedgingTokenBucket.validateAndResolveHedgingOptions({maxTokens: -5}), + RangeError, + ); + assert.throws( + () => + HedgingTokenBucket.validateAndResolveHedgingOptions({maxTokens: 251}), + RangeError, + ); + assert.throws( + () => + HedgingTokenBucket.validateAndResolveHedgingOptions({ + maxTokens: 10.5, + }), + RangeError, + ); + }); + + it('throws RangeError when refillRatio is out of bounds or NaN', () => { + assert.throws( + () => + HedgingTokenBucket.validateAndResolveHedgingOptions({ + refillRatio: 0.0009, + }), + RangeError, + ); + assert.throws( + () => + HedgingTokenBucket.validateAndResolveHedgingOptions({ + refillRatio: 0.21, + }), + RangeError, + ); + assert.throws( + () => + HedgingTokenBucket.validateAndResolveHedgingOptions({ + refillRatio: Number.NaN, + }), + RangeError, + ); + }); + + it('throws Error when both hedging and messageOrdering are enabled', () => { + assert.throws( + () => HedgingTokenBucket.validateAndResolveHedgingOptions({}, true), + /Publish hedging and message ordering cannot be enabled at the same time\./, + ); + }); + + it('throws RangeError when hedgeDelay is greater than or equal to initialRpcTimeoutMillis', () => { + assert.throws( + () => + HedgingTokenBucket.validateAndResolveHedgingOptions( + {hedgeDelay: Duration.from({milliseconds: 500})}, + false, + { + retry: { + backoffSettings: { + initialRetryDelayMillis: 100, + retryDelayMultiplier: 1.3, + maxRetryDelayMillis: 60000, + initialRpcTimeoutMillis: 500, + rpcTimeoutMultiplier: 1, + maxRpcTimeoutMillis: 60000, + totalTimeoutMillis: 600000, + }, + }, + }, + ), + /strictly less than the initial RPC timeout duration/, + ); + }); + + it('throws RangeError when hedgeDelay is greater than or equal to totalTimeoutMillis', () => { + assert.throws( + () => + HedgingTokenBucket.validateAndResolveHedgingOptions( + {hedgeDelay: Duration.from({milliseconds: 1000})}, + false, + { + retry: { + backoffSettings: { + initialRetryDelayMillis: 100, + retryDelayMultiplier: 1.3, + maxRetryDelayMillis: 60000, + initialRpcTimeoutMillis: 5000, + rpcTimeoutMultiplier: 1, + maxRpcTimeoutMillis: 60000, + totalTimeoutMillis: 1000, + }, + }, + }, + ), + /strictly less than the total timeout duration/, + ); + }); + }); + + describe('HedgingTokenBucket', () => { + it('starts empty and cannot acquire a token initially', () => { + const resolved = HedgingTokenBucket.validateAndResolveHedgingOptions({}); + assert.ok(resolved); + const bucket = new HedgingTokenBucket(resolved); + assert.strictEqual(bucket.getTokenBalance(), 0); + assert.strictEqual(bucket.tryAcquireHedgeToken(), false); + }); + + it('refills tokens on success and consumes 1.0 token per acquisition', () => { + const resolved = HedgingTokenBucket.validateAndResolveHedgingOptions({ + maxTokens: 2, + refillRatio: 0.2, + }); + assert.ok(resolved); + const bucket = new HedgingTokenBucket(resolved); + + for (let i = 0; i < 4; i++) { + bucket.refillTokenBucket(); + } + assert.strictEqual(bucket.getTokenBalance(), 0.8); + assert.strictEqual(bucket.tryAcquireHedgeToken(), false); + + bucket.refillTokenBucket(); + assert.strictEqual(bucket.getTokenBalance(), 1.0); + assert.strictEqual(bucket.tryAcquireHedgeToken(), true); + assert.strictEqual(bucket.getTokenBalance(), 0); + assert.strictEqual(bucket.tryAcquireHedgeToken(), false); + }); + + it('caps token balance at maxTokens', () => { + const resolved = HedgingTokenBucket.validateAndResolveHedgingOptions({ + maxTokens: 1, + refillRatio: 0.2, + }); + assert.ok(resolved); + const bucket = new HedgingTokenBucket(resolved); + + for (let i = 0; i < 10; i++) { + bucket.refillTokenBucket(); + } + assert.strictEqual(bucket.getTokenBalance(), 1.0); + }); + + it('maintains exact fixed-point precision for 0.001 refillRatio', () => { + const resolved = HedgingTokenBucket.validateAndResolveHedgingOptions({ + maxTokens: 10, + refillRatio: 0.001, + }); + assert.ok(resolved); + const bucket = new HedgingTokenBucket(resolved); + + for (let i = 0; i < 999; i++) { + bucket.refillTokenBucket(); + } + assert.strictEqual(bucket.tryAcquireHedgeToken(), false); + + bucket.refillTokenBucket(); + assert.strictEqual(bucket.getTokenBalance(), 1.0); + assert.strictEqual(bucket.tryAcquireHedgeToken(), true); + assert.strictEqual(bucket.getTokenBalance(), 0); + }); + }); +}); diff --git a/handwritten/pubsub/test/publisher/index.ts b/handwritten/pubsub/test/publisher/index.ts index 6930b3baa67e..32907341937d 100644 --- a/handwritten/pubsub/test/publisher/index.ts +++ b/handwritten/pubsub/test/publisher/index.ts @@ -28,6 +28,7 @@ import {PublishError} from '../../src/publisher/publish-error'; import * as util from '../../src/util'; import {defaultOptions} from '../../src/default-options'; +import {Duration} from '../../src/temporal'; import * as tracing from '../../src/telemetry-tracing'; import {exporter} from '../tracing'; import {SpanKind} from '@opentelemetry/api'; @@ -485,6 +486,58 @@ describe('Publisher', () => { stubs.forEach(s => assert.ok(s.calledOnce)); }); + + it('should configure hedging options and token bucket when hedging is enabled', () => { + publisher.setOptions({ + hedging: { + hedgeDelay: Duration.from({milliseconds: 300}), + maxTokens: 20, + refillRatio: 0.2, + }, + }); + + assert.ok(publisher.hedgingOptions); + assert.strictEqual(publisher.hedgingOptions.hedgeDelay.milliseconds, 300); + assert.strictEqual(publisher.hedgingOptions.maxTokens, 20); + assert.strictEqual(publisher.hedgingOptions.refillRatio, 0.2); + assert.ok(publisher.tokenBucket); + assert.strictEqual(publisher.tokenBucket.getTokenBalance(), 0); + + for (let i = 0; i < 5; i++) { + publisher.tokenBucket.refillTokenBucket(); + } + assert.strictEqual(publisher.tokenBucket.getTokenBalance(), 1.0); + assert.strictEqual(publisher.tokenBucket.tryAcquireHedgeToken(), true); + assert.strictEqual(publisher.tokenBucket.getTokenBalance(), 0); + }); + + it('should preserve hedging options when setOptions is called with other options', () => { + publisher.setOptions({ + hedging: { + hedgeDelay: Duration.from({milliseconds: 300}), + maxTokens: 20, + refillRatio: 0.2, + }, + }); + publisher.setOptions({ + batching: {maxMessages: 50}, + }); + assert.ok(publisher.hedgingOptions); + assert.strictEqual(publisher.hedgingOptions.hedgeDelay.milliseconds, 300); + assert.strictEqual(publisher.hedgingOptions.maxTokens, 20); + assert.strictEqual(publisher.hedgingOptions.refillRatio, 0.2); + }); + + it('should throw when hedging and messageOrdering are both enabled', () => { + assert.throws( + () => + publisher.setOptions({ + messageOrdering: true, + hedging: {}, + }), + /Publish hedging and message ordering cannot be enabled at the same time\./, + ); + }); }); describe('flush', () => {