feat(pubsub): integrate publish hedging into MessageQueue

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: {

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<void> {
-    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<void> {
+    const {topic, settings} = this.publisher;
     const requestCallback = topic.request<google.pubsub.v1.IPublishResponse>;
     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<void> {
+    const {topic, settings} = this.publisher;
+    const requestCallback = topic.request<google.pubsub.v1.IPublishResponse>;
+    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();

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);
     });
   }