fix: wait for 'drain' as well

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().
+            setTimeout(() => {
+              publisher.queue.emit('drain');
+            }, 50);
             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.
+            setTimeout(() => {
+              queue.emit('drain');
+            }, 50);
             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().
+          setTimeout(() => {
+            publisher.queue.emit('drain');
+          }, 50);
           if (typeof callback === 'function') callback(null);
         });