add iterative drain on AQ

Author: Brian Chen

packages/firestore/src/util/async_queue.ts

@@ -206,6 +206,9 @@ export class AsyncQueue {
   // visible for testing
   failure: Error | null = null;
 
+  // Number of unexecuted operations remaining on the AsyncQueue.
+  private operationsCount = 0;
+
   // Flag set while there's an outstanding AsyncQueue operation, used for
   // assertion sanity-checks.
   private operationInProgress = false;
@@ -219,6 +222,11 @@ export class AsyncQueue {
     return this._isShuttingDown;
   }
 
+  // Visible for testing
+  isEmpty(): boolean {
+    return this.operationsCount === 0;
+  }
+
   /**
    * Adds a new operation to the queue without waiting for it to complete (i.e.
    * we ignore the Promise result).
@@ -280,6 +288,7 @@ export class AsyncQueue {
   }
 
   private enqueueInternal<T extends unknown>(op: () => Promise<T>): Promise<T> {
+    this.operationsCount += 1;
     const newTail = this.tail.then(() => {
       this.operationInProgress = true;
       return op()
@@ -305,6 +314,7 @@ export class AsyncQueue {
         })
         .then(result => {
           this.operationInProgress = false;
+          this.operationsCount -= 1;
           return result;
         });
     });

packages/firestore/test/unit/specs/spec_test_runner.ts

@@ -549,6 +549,16 @@ abstract class TestRunner {
   async run(step: SpecStep): Promise<void> {
     await this.doStep(step);
     await this.queue.drain();
+    // Wait for any additional operations that the initial operations
+    // scheduled. We skip this iterative drain for steps that involve changing
+    // the primary tab since tabs that become the primary can schedule
+    // multiple operations onto the queue that subsequent steps require.
+    while (
+      step.applyClientState?.primary === undefined &&
+      !this.queue.isEmpty()
+    ) {
+      await this.queue.drain();
+    }
     this.validateExpectedSnapshotEvents(step.expectedSnapshotEvents!);
     await this.validateExpectedState(step.expectedState!);
     this.validateSnapshotsInSyncEvents(step.expectedSnapshotsInSyncEvents);