From 1231c0c42ed6e98211c4994edb7033ab618883ab Mon Sep 17 00:00:00 2001 From: feywind <57276408+feywind@users.noreply.github.com> Date: Thu, 15 Oct 2026 18:59:32 +0000 Subject: [PATCH] feat(pubsub): integrate publish hedging into MessageQueue --- handwritten/pubsub/src/publisher/index.ts | 30 +++- .../pubsub/src/publisher/message-queues.ts | 164 +++++++++++++++++- handwritten/pubsub/src/pubsub.ts | 28 ++- 3 files changed, 209 insertions(+), 13 deletions(-) diff --git a/handwritten/pubsub/src/publisher/index.ts b/handwritten/pubsub/src/publisher/index.ts index 5801d46cd7df..bbb5181ac474 100644 --- a/handwritten/pubsub/src/publisher/index.ts +++ b/handwritten/pubsub/src/publisher/index.ts @@ -19,7 +19,7 @@ import {CallOptions} from 'google-gax'; import {isSpanContextValid, Span} from '@opentelemetry/api'; import {BatchPublishOptions} from './message-batch'; -import {Queue, OrderedQueue} from './message-queues'; +import {Queue, OrderedQueue, logs} from './message-queues'; import {Topic} from '../topic'; import {RequestCallback, EmptyCallback} from '../pubsub'; import {defaultOptions} from '../default-options'; @@ -28,6 +28,7 @@ import * as tracing from '../telemetry-tracing'; import {FlowControl, FlowControlOptions} from './flow-control'; import { HedgingOptions, + HedgingScheduler, ResolvedHedgingOptions, HedgingTokenBucket, } from './hedging'; @@ -100,6 +101,7 @@ export class Publisher { flowControl: FlowControl; hedgingOptions?: ResolvedHedgingOptions; private tokenBucket_?: HedgingTokenBucket; + hedgingScheduler?: HedgingScheduler; constructor(topic: Topic, options?: PublishOptions) { this.flowControl = new FlowControl( @@ -149,10 +151,14 @@ export class Publisher { allPublishes .then(() => allDrains) .then(() => { + this.hedgingScheduler?.clear(); definedCallback(null); return undefined; }) - .catch(definedCallback); + .catch(err => { + this.hedgingScheduler?.clear(); + definedCallback(err); + }); } /** @@ -317,10 +323,24 @@ export class Publisher { messageOrdering, gaxOpts, ); + this.hedgingScheduler?.clear(); this.hedgingOptions = resolvedHedging; - this.tokenBucket_ = resolvedHedging - ? new HedgingTokenBucket(resolvedHedging) - : undefined; + if (resolvedHedging) { + const tokenBucket = new HedgingTokenBucket(resolvedHedging); + this.tokenBucket_ = tokenBucket; + this.hedgingScheduler = new HedgingScheduler( + resolvedHedging, + tokenBucket, + () => { + logs.publishHedged.debug( + 'Hedging rate limited due to lack of tokens.', + ); + }, + ); + } else { + this.tokenBucket_ = undefined; + this.hedgingScheduler = undefined; + } this.settings = { batching: { diff --git a/handwritten/pubsub/src/publisher/message-queues.ts b/handwritten/pubsub/src/publisher/message-queues.ts index ee61ae6f06ee..27855e9e3306 100644 --- a/handwritten/pubsub/src/publisher/message-queues.ts +++ b/handwritten/pubsub/src/publisher/message-queues.ts @@ -14,13 +14,21 @@ * limitations under the License. */ -import {ServiceError} from 'google-gax'; +import {CallOptions, ServiceError} from 'google-gax'; +import {Span} from '@opentelemetry/api'; import {EventEmitter} from 'events'; +import * as extend from 'extend'; import {BatchPublishOptions, MessageBatch} from './message-batch'; import {PublishError} from './publish-error'; import {Publisher, PubsubMessage, PublishCallback} from './'; +import { + CancellationSharer, + HedgingScheduler, + HedgingTokenBucket, +} from './hedging'; import {google} from '../../protos/protos'; +import {Duration} from '../temporal'; import * as tracing from '../telemetry-tracing'; import {filterMessage} from './pubsub-message'; import {promisify} from 'util'; @@ -51,7 +59,7 @@ export abstract class MessageQueue extends EventEmitter { constructor(publisher: Publisher) { super(); this.publisher = publisher; - this.batchOptions = publisher.settings.batching!; + this.batchOptions = publisher.settings.batching ?? {}; } /** @@ -64,7 +72,17 @@ export abstract class MessageQueue extends EventEmitter { * @private */ updateOptions() { - this.batchOptions = this.publisher.settings.batching!; + this.batchOptions = this.publisher.settings.batching ?? {}; + } + + /** + * Indicates whether this queue implementation supports publish hedging. + * Ordered queues override this to return false. + * + * @protected + */ + protected get allowsHedging(): boolean { + return true; } /** @@ -120,8 +138,8 @@ export abstract class MessageQueue extends EventEmitter { bytes: number, reason?: string, ): Promise { - const {topic, settings} = this.publisher; - const reqOpts = { + const {topic} = this.publisher; + const reqOpts: google.pubsub.v1.IPublishRequest = { topic: topic.name, messages: messages.map(filterMessage), }; @@ -138,7 +156,12 @@ export abstract class MessageQueue extends EventEmitter { await topic.pubsub.getClientConfig(); } spanMessages.forEach(m => { - tracing.PubsubSpans.updatePublisherTopicName(m.parentSpan!, topic.name); + if (m.parentSpan) { + tracing.PubsubSpans.updatePublisherTopicName( + m.parentSpan, + topic.name, + ); + } tracing.PubsubEvents.publishStart(m); }); } @@ -153,14 +176,45 @@ export abstract class MessageQueue extends EventEmitter { this.logBatch(reason, messages.length, bytes); } + const hedgingScheduler = this.allowsHedging + ? this.publisher.hedgingScheduler + : undefined; + + if (hedgingScheduler && this.publisher.hedgingOptions) { + await this._publishHedged( + reqOpts, + messages, + spanMessages, + callbacks, + hedgingScheduler, + rpcSpan, + ); + } else { + await this._publishStandard(reqOpts, messages, callbacks, rpcSpan); + } + } + + /** + * Publishes a batch of messages without hedging. + * + * @private + */ + private async _publishStandard( + reqOpts: google.pubsub.v1.IPublishRequest, + messages: PubsubMessage[], + callbacks: PublishCallback[], + rpcSpan?: Span, + ): Promise { + const {topic, settings} = this.publisher; const requestCallback = topic.request; const request = promisify(requestCallback.bind(topic)); + try { const resp = await request({ client: 'PublisherClient', method: 'publish', reqOpts, - gaxOpts: settings.gaxOpts!, + gaxOpts: settings.gaxOpts, }); if (resp) { @@ -182,6 +236,98 @@ export abstract class MessageQueue extends EventEmitter { }); } } + + /** + * Publishes a batch of messages with hedging coordination and scheduling. + * + * @private + */ + private async _publishHedged( + reqOpts: google.pubsub.v1.IPublishRequest, + messages: PubsubMessage[], + spanMessages: PubsubMessage[], + callbacks: PublishCallback[], + hedgingScheduler: HedgingScheduler, + rpcSpan?: Span, + ): Promise { + const {topic, settings} = this.publisher; + const requestCallback = topic.request; + const request = promisify(requestCallback.bind(topic)); + + const totalTimeout = HedgingTokenBucket.resolveTotalTimeout( + settings.gaxOpts, + ); + const absoluteDeadline = Duration.from({milliseconds: Date.now()}).add( + totalTimeout, + ); + const coordinator = new CancellationSharer(absoluteDeadline, () => { + this.publisher.tokenBucket?.refillTokenBucket(); + }); + + const originalAbortController = new AbortController(); + const originalGaxOpts: CallOptions = extend(true, {}, settings.gaxOpts, { + signal: originalAbortController.signal, + }); + const originalAttemptPromise = request({ + client: 'PublisherClient', + method: 'publish', + reqOpts, + gaxOpts: originalGaxOpts, + }); + coordinator.addAttempt(0, originalAbortController, originalAttemptPromise); + + hedgingScheduler.scheduleFirstHedge( + coordinator, + (attemptNumber, attemptTimeout) => { + logs.publishHedged.debug('Publishing hedged attempt %i', attemptNumber); + spanMessages.forEach(m => { + tracing.PubsubEvents.publishStartHedged(m); + }); + const hedgedAbortController = new AbortController(); + const hedgedGaxOpts: CallOptions = extend(true, {}, settings.gaxOpts, { + timeout: attemptTimeout.milliseconds, + retry: null, + signal: hedgedAbortController.signal, + }); + const hedgedAttemptPromise = request({ + client: 'PublisherClient', + method: 'publish', + reqOpts, + gaxOpts: hedgedGaxOpts, + }); + coordinator.addAttempt( + attemptNumber, + hedgedAbortController, + hedgedAttemptPromise, + ); + }, + ); + + let wasHedged = false; + try { + const result = await coordinator.promise; + wasHedged = result.wasHedged; + if (result.response) { + const messageIds = result.response.messageIds || []; + callbacks.forEach((callback, i) => callback(null, messageIds[i])); + } + } catch (e) { + const err = e as ServiceError; + callbacks.forEach(callback => callback(err)); + + throw e; + } finally { + rpcSpan?.end(); + messages.forEach(m => { + if (wasHedged) { + tracing.PubsubEvents.publishEndHedged(m); + } else { + tracing.PubsubEvents.publishEnd(m); + } + m.parentSpan?.end(); + }); + } + } } /** @@ -314,6 +460,10 @@ export class OrderedQueue extends MessageQueue { this.key = key; } + protected get allowsHedging(): boolean { + return false; + } + // This needs to update our existing message batches. updateOptions() { super.updateOptions(); diff --git a/handwritten/pubsub/src/pubsub.ts b/handwritten/pubsub/src/pubsub.ts index 6bd991daff0b..df40a4a88820 100644 --- a/handwritten/pubsub/src/pubsub.ts +++ b/handwritten/pubsub/src/pubsub.ts @@ -19,6 +19,13 @@ import {replaceProjectIdToken} from '@google-cloud/projectify'; import * as extend from 'extend'; import {AuthClient, GoogleAuth, GoogleAuthOptions} from 'google-auth-library'; import * as gax from 'google-gax'; + +// Module augmentation for google-gax CallOptions until the updated google-gax package is published. +declare module 'google-gax' { + export interface CallOptions { + signal?: AbortSignal; + } +} import {v1} from './'; // eslint-disable-next-line @typescript-eslint/no-var-requires @@ -1330,6 +1337,16 @@ export class PubSub { await Promise.all(promises); } + private createCancelledError_(): gax.grpc.ServiceError { + const statusObject = { + code: gax.Status.CANCELLED, + details: 'Call cancelled.', + metadata: new gax.grpc.Metadata(), + }; + const err = new Error(statusObject.details); + return Object.assign(err, statusObject) as gax.grpc.ServiceError; + } + /** * Funnel all API requests through this method, to be sure we have a project * ID. @@ -1356,14 +1373,23 @@ export class PubSub { return; } + if (config.gaxOpts?.signal?.aborted) { + callback(this.createCancelledError_()); + return; + } + this.getClient_(config, (err, client) => { if (err) { callback(err as gax.grpc.ServiceError); return; } + if (config.gaxOpts?.signal?.aborted) { + callback(this.createCancelledError_()); + return; + } let reqOpts = extend(true, {}, config.reqOpts); reqOpts = replaceProjectIdToken(reqOpts, this.projectId); - client![config.method](reqOpts, config.gaxOpts, callback); + client?.[config.method](reqOpts, config.gaxOpts, callback); }); }