fix: wait for 'drain' as well (#1636)

It's possible to end up in a situation where more publishes are pending when publish() is called, but they aren't actually waited for. This adds 'drain' event support to the main publisher queue and waits for 'drain' on all queues before calling it done.

Maybe fixes #1620 🦕

src/publisher/index.ts

@@ -110,13 +110,30 @@ export class Publisher {
   flush(callback?: EmptyCallback): Promise<void> | void {
     const definedCallback = callback ? callback : () => {};
 
-    const publishes = [promisify(this.queue.publish).bind(this.queue)()];
-    Array.from(this.orderedQueues.values()).forEach(q =>
-      publishes.push(promisify(q.publish).bind(q)())
+    const toDrain = [this.queue, ...Array.from(this.orderedQueues.values())];
+
+    const allDrains = Promise.all(
+      toDrain.map(
+        q =>
+          new Promise<void>(resolve => {
+            const flushResolver = () => {
+              resolve();
+
+              // flush() maybe called more than once, so remove these
+              // event listeners after we've completed flush().
+              q.removeListener('drain', flushResolver);
+            };
+            return q.on('drain', flushResolver);
+          })
+      )
+    );
+
+    const allPublishes = Promise.all(
+      toDrain.map(q => promisify(q.publish).bind(q)())
     );
-    const allPublishes = Promise.all(publishes);
 
     allPublishes
+      .then(() => allDrains)
       .then(() => {
         definedCallback(null);
       })

src/publisher/message-queues.ts

@@ -158,8 +158,11 @@ export class Queue extends MessageQueue {
   }
   /**
    * Cancels any pending publishes and calls _publish immediately.
+   *
+   * @emits Queue#drain when all messages are sent.
    */
   publish(callback?: PublishDone): void {
+    const definedCallback = callback || (() => {});
     const {messages, callbacks} = this.batch;
 
     this.batch = new MessageBatch(this.batchOptions);
@@ -169,7 +172,17 @@ export class Queue extends MessageQueue {
       delete this.pending;
     }
 
-    this._publish(messages, callbacks, callback);
+    this._publish(messages, callbacks, (err: null | ServiceError) => {
+      if (err) {
+        definedCallback(err);
+      } else if (this.batch.messages.length) {
+        // Make another go-around, we're trying to drain the queues fully.
+        this.publish(callback);
+      } else {
+        this.emit('drain');
+        definedCallback(null);
+      }
+    });
   }
 }
 

test/publisher/index.ts

@@ -347,6 +347,11 @@ describe('Publisher', () => {
         sandbox
           .stub(FakeQueue.prototype, '_publish')
           .callsFake((messages, callbacks, callback) => {
+            // Simulate the drain taking longer than the publishes. This can
+            // happen if more messages are queued during the publish().
+            process.nextTick(() => {
+              publisher.queue.emit('drain');
+            });
             if (typeof callback === 'function') callback(null);
           });
 
@@ -356,7 +361,12 @@ describe('Publisher', () => {
             const queue = publisher.orderedQueues.get(
               orderingKey
             ) as unknown as FakeOrderedQueue;
-            queue.emit('drain');
+            // Simulate the drain taking longer than the publishes. This can
+            // happen on some ordered queue scenarios, especially if we have more
+            // than one queue to empty.
+            process.nextTick(() => {
+              queue.emit('drain');
+            });
             if (typeof callback === 'function') callback(null);
           });
 
@@ -495,6 +505,11 @@ describe('Publisher', () => {
       sandbox
         .stub(publisher.queue, '_publish')
         .callsFake((messages, callbacks, callback) => {
+          // Simulate the drain taking longer than the publishes. This can
+          // happen if more messages are queued during the publish().
+          process.nextTick(() => {
+            publisher.queue.emit('drain');
+          });
           if (typeof callback === 'function') callback(null);
         });