diff --git a/handwritten/pubsub/src/publisher/hedging.ts b/handwritten/pubsub/src/publisher/hedging.ts index 0fd20bc8637..383fdf51093 100644 --- a/handwritten/pubsub/src/publisher/hedging.ts +++ b/handwritten/pubsub/src/publisher/hedging.ts @@ -15,8 +15,9 @@ */ import {CallOptions} from 'google-gax'; +import {google} from '../../protos/protos'; import {defaultOptions} from '../default-options'; -import {Duration} from '../temporal'; +import {Duration, atLeast, atMost} from '../temporal'; /** * Minimum allowed hedging delay (100ms). @@ -28,6 +29,16 @@ const MIN_HEDGE_DELAY = Duration.from({milliseconds: 100}); */ const MAX_HEDGE_DELAY = Duration.from({seconds: 10}); +/** + * Maximum timeout for an individual hedged publish attempt (10s). + */ +const MAX_HEDGE_ATTEMPT_TIMEOUT = Duration.from({seconds: 10}); + +/** + * Zero duration constant for non-negative clamping and comparisons. + */ +const ZERO_DURATION = Duration.from({milliseconds: 0}); + /** * Minimum allowed maxTokens value. */ @@ -248,3 +259,286 @@ export class HedgingTokenBucket { return this.hedgeTokenBucket / HEDGE_TOKEN_SCALE; } } + +/** + * Outcome of a hedged publish batch. + */ +export interface HedgedPublishResult { + readonly response?: google.pubsub.v1.IPublishResponse | null; + readonly wasHedged: boolean; + readonly successfulAttempt: number; +} + +/** + * Callback invoked by HedgingScheduler when a hedged attempt should be started. + */ +export type StartHedgedAttemptCallback = ( + attemptNumber: number, + attemptTimeout: Duration, +) => void; + +/** + * Coordinates multiple publish attempts (attempt 0 = original, 1..N = hedged) + * for a single batch of messages. + */ +export class CancellationSharer { + readonly absoluteDeadline: Duration; + readonly promise: Promise; + + private readonly onSuccess: () => void; + private readonly runningAttempts = new Map(); + private done_ = false; + private successfulAttempt_?: number; + private lastError: unknown; + private resolvePromise?: (result: HedgedPublishResult) => void; + private rejectPromise?: (err: unknown) => void; + + constructor(absoluteDeadline: Duration, onSuccess: () => void) { + this.absoluteDeadline = absoluteDeadline; + this.onSuccess = onSuccess; + this.promise = new Promise((resolve, reject) => { + this.resolvePromise = resolve; + this.rejectPromise = reject; + }); + } + + /** + * Registers an in-flight publish attempt with this coordinator. + */ + addAttempt( + attemptNumber: number, + abortController: AbortController, + attemptPromise: Promise< + google.pubsub.v1.IPublishResponse | null | undefined + >, + ): void { + if (this.done_) { + abortController.abort(); + return; + } + + this.runningAttempts.set(attemptNumber, abortController); + + attemptPromise + .then(response => { + this.handleAttemptSuccess(attemptNumber, response); + return undefined; + }) + .catch((err: unknown) => { + this.handleAttemptFailure(attemptNumber, err); + }); + } + + /** + * Returns true if this batch has already completed (succeeded, failed, or cancelled). + */ + get isDone(): boolean { + return this.done_; + } + + /** + * Returns the attempt number that succeeded, or undefined if none has succeeded. + */ + get successfulAttempt(): number | undefined { + return this.successfulAttempt_; + } + + /** + * Cancels all running attempts and marks the coordinator as done. + */ + cancel(): void { + if (this.done_) { + return; + } + this.done_ = true; + this.cancelAll(); + this.rejectPromise?.(new Error('Cancelled')); + } + + private handleAttemptSuccess( + attemptNumber: number, + response?: google.pubsub.v1.IPublishResponse | null, + ): void { + if (this.done_) { + return; + } + this.done_ = true; + this.successfulAttempt_ = attemptNumber; + this.onSuccess(); + this.cancelAllExcept(attemptNumber); + this.resolvePromise?.({ + response, + wasHedged: attemptNumber > 0, + successfulAttempt: attemptNumber, + }); + } + + private handleAttemptFailure(attemptNumber: number, err: unknown): void { + if (this.done_) { + return; + } + this.runningAttempts.delete(attemptNumber); + this.lastError = err; + + if (attemptNumber === 0 || this.runningAttempts.size === 0) { + this.done_ = true; + this.cancelAll(); + this.rejectPromise?.(this.lastError); + } + } + + private cancelAll(): void { + for (const controller of this.runningAttempts.values()) { + controller.abort(); + } + this.runningAttempts.clear(); + } + + private cancelAllExcept(winningAttempt: number): void { + for (const [attempt, controller] of this.runningAttempts.entries()) { + if (attempt !== winningAttempt) { + controller.abort(); + } + } + this.runningAttempts.clear(); + } +} + +/** + * Represents a pending hedging check in the publisher's hedging queue. + */ +export class HedgedRequest { + readonly coordinator: CancellationSharer; + readonly attemptNumber: number; + readonly sendAfter: Duration; + readonly startHedgedAttempt: StartHedgedAttemptCallback; + + constructor( + coordinator: CancellationSharer, + attemptNumber: number, + sendAfter: Duration, + startHedgedAttempt: StartHedgedAttemptCallback, + ) { + this.coordinator = coordinator; + this.attemptNumber = attemptNumber; + this.sendAfter = sendAfter; + this.startHedgedAttempt = startHedgedAttempt; + } +} + +/** + * Event-driven FIFO queue and single-timer scheduler for publish hedging. + */ +export class HedgingScheduler { + private readonly options: ResolvedHedgingOptions; + private readonly tokenBucket: HedgingTokenBucket; + private readonly onRateLimited?: () => void; + private readonly hedgingQueue: HedgedRequest[] = []; + private queueProcessingTimer?: NodeJS.Timeout; + + constructor( + options: ResolvedHedgingOptions, + tokenBucket: HedgingTokenBucket, + onRateLimited?: () => void, + ) { + this.options = options; + this.tokenBucket = tokenBucket; + this.onRateLimited = onRateLimited; + } + + /** + * Returns the current wall-clock time as a Duration since epoch. + */ + private nowDuration(): Duration { + return Duration.from({milliseconds: Date.now()}); + } + + /** + * Schedules the initial hedge check (attempt 1) for a newly started publish batch. + */ + scheduleFirstHedge( + coordinator: CancellationSharer, + startHedgedAttempt: StartHedgedAttemptCallback, + ): void { + if (coordinator.isDone) { + return; + } + const sendAfter = this.nowDuration().add(this.options.hedgeDelay); + const request = new HedgedRequest( + coordinator, + 1, + sendAfter, + startHedgedAttempt, + ); + this.hedgingQueue.push(request); + this.scheduleQueueProcessing(); + } + + /** + * Clears any scheduled timer and empties the pending hedging queue. + */ + clear(): void { + if (this.queueProcessingTimer) { + clearTimeout(this.queueProcessingTimer); + this.queueProcessingTimer = undefined; + } + this.hedgingQueue.length = 0; + } + + private scheduleQueueProcessing(): void { + if (this.queueProcessingTimer) { + return; + } + const nextItem = this.hedgingQueue[0]; + if (!nextItem) { + return; + } + const delay = atLeast( + nextItem.sendAfter.subtract(this.nowDuration()), + ZERO_DURATION, + ); + this.queueProcessingTimer = setTimeout(() => { + this.processQueue(); + }, delay.milliseconds); + } + + private processQueue(): void { + this.queueProcessingTimer = undefined; + const now = this.nowDuration(); + + while ( + this.hedgingQueue.length > 0 && + Duration.compare(this.hedgingQueue[0].sendAfter, now) <= 0 + ) { + const item = this.hedgingQueue.shift(); + if (!item || item.coordinator.isDone) { + continue; + } + + const remainingTimeout = item.coordinator.absoluteDeadline.subtract(now); + if (Duration.compare(remainingTimeout, ZERO_DURATION) <= 0) { + continue; + } + + const attemptTimeout = atMost( + remainingTimeout, + MAX_HEDGE_ATTEMPT_TIMEOUT, + ); + + if (this.tokenBucket.tryAcquireHedgeToken()) { + const nextItem = new HedgedRequest( + item.coordinator, + item.attemptNumber + 1, + now.add(this.options.hedgeDelay), + item.startHedgedAttempt, + ); + this.hedgingQueue.push(nextItem); + item.startHedgedAttempt(item.attemptNumber, attemptTimeout); + } else { + this.onRateLimited?.(); + } + } + + this.scheduleQueueProcessing(); + } +} diff --git a/handwritten/pubsub/src/publisher/index.ts b/handwritten/pubsub/src/publisher/index.ts index 8de5e8c4198..5801d46cd7d 100644 --- a/handwritten/pubsub/src/publisher/index.ts +++ b/handwritten/pubsub/src/publisher/index.ts @@ -36,9 +36,14 @@ import {promisifySome} from '../util'; import {PubsubMessage, Attributes} from './pubsub-message'; export {PubsubMessage, Attributes} from './pubsub-message'; export { + CancellationSharer, + HedgedPublishResult, + HedgedRequest, HedgingOptions, - ResolvedHedgingOptions, + HedgingScheduler, HedgingTokenBucket, + ResolvedHedgingOptions, + StartHedgedAttemptCallback, } from './hedging'; export type PublishCallback = RequestCallback; diff --git a/handwritten/pubsub/test/publisher/hedging.ts b/handwritten/pubsub/test/publisher/hedging.ts index 265bcdabb07..9d17cc4b767 100644 --- a/handwritten/pubsub/test/publisher/hedging.ts +++ b/handwritten/pubsub/test/publisher/hedging.ts @@ -15,12 +15,20 @@ */ import * as assert from 'assert'; -import {describe, it} from 'mocha'; +import {describe, it, beforeEach, afterEach} from 'mocha'; +import * as sinon from 'sinon'; +import * as defer from 'p-defer'; +import {google} from '../../protos/protos'; import {defaultOptions} from '../../src/default-options'; import {Duration} from '../../src/temporal'; -import {HedgingTokenBucket} from '../../src/publisher/hedging'; +import { + CancellationSharer, + HedgingScheduler, + HedgingTokenBucket, +} from '../../src/publisher/hedging'; +import {TestUtils} from '../test-utils'; -describe('Publisher Hedging (Options & Token Bucket)', () => { +describe('Publisher Hedging', () => { describe('validateAndResolveHedgingOptions', () => { it('returns undefined when hedging options are not provided', () => { assert.strictEqual( @@ -247,4 +255,282 @@ describe('Publisher Hedging (Options & Token Bucket)', () => { assert.strictEqual(bucket.getTokenBalance(), 0); }); }); + + describe('CancellationSharer', () => { + const deadline = Duration.from({seconds: 100}); + + it('resolves when original attempt succeeds and aborts hedged attempts', async () => { + let successCalled = 0; + const coordinator = new CancellationSharer(deadline, () => { + successCalled++; + }); + + const attempt0 = defer(); + const attempt1 = defer(); + const controller0 = new AbortController(); + const controller1 = new AbortController(); + + coordinator.addAttempt(0, controller0, attempt0.promise); + coordinator.addAttempt(1, controller1, attempt1.promise); + + attempt0.resolve({messageIds: ['msg-0']}); + + const result = await coordinator.promise; + assert.deepStrictEqual(result.response, {messageIds: ['msg-0']}); + assert.strictEqual(result.wasHedged, false); + assert.strictEqual(result.successfulAttempt, 0); + assert.strictEqual(coordinator.successfulAttempt, 0); + assert.strictEqual(successCalled, 1); + assert.strictEqual(controller0.signal.aborted, false); + assert.strictEqual(controller1.signal.aborted, true); + }); + + it('resolves when hedged attempt succeeds first and aborts original attempt', async () => { + let successCalled = 0; + const coordinator = new CancellationSharer(deadline, () => { + successCalled++; + }); + + const attempt0 = defer(); + const attempt1 = defer(); + const controller0 = new AbortController(); + const controller1 = new AbortController(); + + coordinator.addAttempt(0, controller0, attempt0.promise); + coordinator.addAttempt(1, controller1, attempt1.promise); + + attempt1.resolve({messageIds: ['msg-1']}); + + const result = await coordinator.promise; + assert.deepStrictEqual(result.response, {messageIds: ['msg-1']}); + assert.strictEqual(result.wasHedged, true); + assert.strictEqual(result.successfulAttempt, 1); + assert.strictEqual(coordinator.successfulAttempt, 1); + assert.strictEqual(successCalled, 1); + assert.strictEqual(controller0.signal.aborted, true); + assert.strictEqual(controller1.signal.aborted, false); + }); + + it('discards hedged attempt failure while original attempt is still running', async () => { + const coordinator = new CancellationSharer(deadline, () => {}); + + const attempt0 = defer(); + const attempt1 = defer(); + const controller0 = new AbortController(); + const controller1 = new AbortController(); + + coordinator.addAttempt(0, controller0, attempt0.promise); + coordinator.addAttempt(1, controller1, attempt1.promise); + + attempt1.reject(new Error('hedged failure')); + await Promise.resolve(); + assert.strictEqual(coordinator.isDone, false); + assert.strictEqual(coordinator.successfulAttempt, undefined); + + attempt0.resolve({messageIds: ['msg-0']}); + const result = await coordinator.promise; + assert.deepStrictEqual(result.response, {messageIds: ['msg-0']}); + assert.strictEqual(result.wasHedged, false); + }); + + it('immediately rejects and aborts running hedged attempts when original attempt fails', async () => { + const coordinator = new CancellationSharer(deadline, () => {}); + + const attempt0 = defer(); + const attempt1 = defer(); + const controller0 = new AbortController(); + const controller1 = new AbortController(); + + coordinator.addAttempt(0, controller0, attempt0.promise); + coordinator.addAttempt(1, controller1, attempt1.promise); + + const originalError = new Error('original permanent failure'); + attempt0.reject(originalError); + + await assert.rejects(coordinator.promise, originalError); + assert.strictEqual(controller1.signal.aborted, true); + }); + + it('immediately aborts an attempt added after coordinator is already done', async () => { + const coordinator = new CancellationSharer(deadline, () => {}); + const controller0 = new AbortController(); + coordinator.addAttempt( + 0, + controller0, + Promise.resolve({messageIds: ['1']}), + ); + await coordinator.promise; + + const lateController = new AbortController(); + const lateAttempt = defer(); + coordinator.addAttempt(1, lateController, lateAttempt.promise); + assert.strictEqual(lateController.signal.aborted, true); + }); + }); + + describe('HedgingScheduler', () => { + interface RecordedAttempt { + attempt: number; + timeoutMs: number; + } + + let sandbox: sinon.SinonSandbox; + let clock: sinon.SinonFakeTimers; + + beforeEach(() => { + sandbox = sinon.createSandbox(); + clock = TestUtils.useFakeTimers(sandbox, 10000); + }); + + afterEach(() => { + clock.restore(); + sandbox.restore(); + }); + + it('fires hedged attempts at hedgeDelay intervals (multiple hedging) while tokens are available', () => { + const options = HedgingTokenBucket.validateAndResolveHedgingOptions({ + hedgeDelay: Duration.from({milliseconds: 200}), + maxTokens: 10, + refillRatio: 0.2, + }); + assert.ok(options); + const bucket = new HedgingTokenBucket(options); + for (let i = 0; i < 10; i++) { + bucket.refillTokenBucket(); + } + assert.strictEqual(bucket.getTokenBalance(), 2.0); + + let rateLimitedCount = 0; + const scheduler = new HedgingScheduler(options, bucket, () => { + rateLimitedCount++; + }); + + const absoluteDeadline = Duration.from({milliseconds: Date.now()}).add( + Duration.from({seconds: 25}), + ); + const coordinator = new CancellationSharer(absoluteDeadline, () => {}); + + const startedAttempts: RecordedAttempt[] = []; + scheduler.scheduleFirstHedge( + coordinator, + (attemptNumber, attemptTimeout) => { + startedAttempts.push({ + attempt: attemptNumber, + timeoutMs: attemptTimeout.milliseconds, + }); + }, + ); + + clock.tick(199); + assert.strictEqual(startedAttempts.length, 0); + + // At +200ms, attempt 1 fires (timeout capped at 10000ms) + clock.tick(1); + assert.deepStrictEqual(startedAttempts, [{attempt: 1, timeoutMs: 10000}]); + assert.strictEqual(bucket.getTokenBalance(), 1.0); + + // At +400ms, attempt 2 fires (consuming the second token) + clock.tick(200); + assert.deepStrictEqual(startedAttempts, [ + {attempt: 1, timeoutMs: 10000}, + {attempt: 2, timeoutMs: 10000}, + ]); + assert.strictEqual(bucket.getTokenBalance(), 0); + + // At +600ms, bucket is empty so attempt 3 is rate-limited and not re-queued + clock.tick(200); + assert.strictEqual(startedAttempts.length, 2); + assert.strictEqual(rateLimitedCount, 1); + + scheduler.clear(); + }); + + it('caps hedged attempt timeout at remainingTimeout when less than 10s', () => { + const options = HedgingTokenBucket.validateAndResolveHedgingOptions({ + hedgeDelay: Duration.from({milliseconds: 500}), + maxTokens: 10, + refillRatio: 0.2, + }); + assert.ok(options); + const bucket = new HedgingTokenBucket(options); + for (let i = 0; i < 5; i++) { + bucket.refillTokenBucket(); + } + + const scheduler = new HedgingScheduler(options, bucket); + const absoluteDeadline = Duration.from({milliseconds: Date.now()}).add( + Duration.from({milliseconds: 3000}), + ); + const coordinator = new CancellationSharer(absoluteDeadline, () => {}); + + let observedTimeout: Duration | undefined; + scheduler.scheduleFirstHedge(coordinator, (_attempt, attemptTimeout) => { + observedTimeout = attemptTimeout; + }); + + clock.tick(500); + assert.ok(observedTimeout); + assert.strictEqual(observedTimeout.milliseconds, 2500); + }); + + it('skips hedging if coordinator completes before hedgeDelay', async () => { + const options = HedgingTokenBucket.validateAndResolveHedgingOptions({ + hedgeDelay: Duration.from({milliseconds: 500}), + maxTokens: 10, + refillRatio: 0.2, + }); + assert.ok(options); + const bucket = new HedgingTokenBucket(options); + for (let i = 0; i < 5; i++) { + bucket.refillTokenBucket(); + } + + const scheduler = new HedgingScheduler(options, bucket); + const absoluteDeadline = Duration.from({milliseconds: Date.now()}).add( + Duration.from({seconds: 60}), + ); + const coordinator = new CancellationSharer(absoluteDeadline, () => {}); + + let hedgeFired = false; + scheduler.scheduleFirstHedge(coordinator, () => { + hedgeFired = true; + }); + + coordinator.cancel(); + await assert.rejects(coordinator.promise, /Cancelled/); + clock.tick(500); + assert.strictEqual(hedgeFired, false); + assert.strictEqual(bucket.getTokenBalance(), 1.0); + }); + + it('does not schedule hedge if coordinator is already done', async () => { + const options = HedgingTokenBucket.validateAndResolveHedgingOptions({ + hedgeDelay: Duration.from({milliseconds: 500}), + maxTokens: 10, + refillRatio: 0.2, + }); + assert.ok(options); + const bucket = new HedgingTokenBucket(options); + for (let i = 0; i < 5; i++) { + bucket.refillTokenBucket(); + } + + const scheduler = new HedgingScheduler(options, bucket); + const absoluteDeadline = Duration.from({milliseconds: Date.now()}).add( + Duration.from({seconds: 60}), + ); + const coordinator = new CancellationSharer(absoluteDeadline, () => {}); + coordinator.cancel(); + await assert.rejects(coordinator.promise, /Cancelled/); + + let hedgeFired = false; + scheduler.scheduleFirstHedge(coordinator, () => { + hedgeFired = true; + }); + + clock.tick(500); + assert.strictEqual(hedgeFired, false); + assert.strictEqual(bucket.getTokenBalance(), 1.0); + }); + }); });