Removing TestFirestore and TestQueue in favor of drain()

Author: schmidt-sebastian

packages/firestore/src/api/database.ts

@@ -188,6 +188,7 @@ export class Firestore implements firestore.Firestore, FirebaseService {
   // Operations on the _firestoreClient don't block on _firestoreReady. Those
   // are already set to synchronize on the async queue.
   private _firestoreClient: FirestoreClient | undefined;
+  private _queue = new AsyncQueue();
   public _dataConverter: UserDataConverter;
 
   constructor(databaseIdOrApp: FirestoreDatabase | FirebaseApp) {
@@ -310,7 +311,7 @@ export class Firestore implements firestore.Firestore, FirebaseService {
       PlatformSupport.getPlatform(),
       databaseInfo,
       this._config.credentials,
-      this.initializeAsyncQueue()
+      this._queue,
     );
     return this._firestoreClient.start(persistence);
   }
@@ -374,7 +375,8 @@ export class Firestore implements firestore.Firestore, FirebaseService {
     },
     // Exposed via INTERNAL for use in tests.
     disableNetwork: () => this._firestoreClient.disableNetwork(),
-    enableNetwork: () => this._firestoreClient.enableNetwork()
+    enableNetwork: () => this._firestoreClient.enableNetwork(),
+    queue :  this._queue,
   };
 
   collection(pathString: string): firestore.CollectionReference {
@@ -455,11 +457,6 @@ export class Firestore implements firestore.Firestore, FirebaseService {
         );
     }
   }
-
-  /** Creates a new AsyncQueue. Can be overridden to provide a custom queue. */
-  protected initializeAsyncQueue(): AsyncQueue {
-    return new AsyncQueue();
-  }
 }
 
 /**

packages/firestore/src/util/async_queue.ts

@@ -20,6 +20,13 @@ import { AnyDuringMigration, AnyJs } from './misc';
 import { Deferred } from './promise';
 import { Code, FirestoreError } from './error';
 
+type DelayedOperation <T> = {
+  // tslint:disable-next-line:no-any Accept any return type from setTimeout().
+  handle: any;
+  op:() => Promise<T>,
+  deferred: Deferred<T>;
+}
+
 export class AsyncQueue {
   // The last promise in the queue.
   private tail: Promise<AnyJs | void> = Promise.resolve();
@@ -29,16 +36,16 @@ export class AsyncQueue {
   // has a delay that has not yet elapsed). Prior to cleanup, this list may also
   // contain entries that have already been run (in which case `handle` is
   // null).
-  private delayedOperations: {
-    // tslint:disable-next-line:no-any Accept any return type from setTimeout().
-    handle: any;
-    deferred: Deferred<AnyJs | void>;
-  }[] = [];
+  //
+  // tslint:disable-next-line:no-any Accept any type of delayed operation.
+  private delayedOperations: DelayedOperation<any>[] = [];
 
   // The number of operations that are queued to be run in the future (i.e. they
   // have a delay that has not yet elapsed). Unlike `delayedOperations`, this
   // is guaranteed to only contain operations that have not yet been run.
-  private delayedOperationsCount = 0;
+  //
+  // Visible for testing.
+  delayedOperationsCount = 0;
 
   // visible for testing
   failure: Error;
@@ -61,23 +68,27 @@ export class AsyncQueue {
 
     if ((delay || 0) > 0) {
       this.delayedOperationsCount++;
-      const deferred = new Deferred<T>();
-      const opIndex = this.delayedOperations.length;
-      const handle = setTimeout(() => {
+      const nextIndex = this.delayedOperations.length;
+      const delayedOp : DelayedOperation<T> = {
+        handle: null,
+        op: op,
+        deferred: new Deferred<T>()
+      };
+      delayedOp.handle = setTimeout(() => {
         this.scheduleInternal(() => {
-          return op().then(result => {
-            deferred.resolve(result);
+          return delayedOp.op().then(result => {
+            delayedOp.deferred.resolve(result);
           });
         });
         this.delayedOperationsCount--;
         if (this.delayedOperationsCount > 0) {
-          this.delayedOperations[opIndex].handle = null;
+          delayedOp.handle = null;
         } else {
           this.delayedOperations = [];
         }
       }, delay);
-      this.delayedOperations[opIndex] = { handle, deferred };
-      return deferred.promise;
+      this.delayedOperations[nextIndex] = delayedOp;
+      return delayedOp.deferred.promise;
     } else {
       return this.scheduleInternal(op);
     }
@@ -114,16 +125,22 @@ export class AsyncQueue {
   }
 
   /**
-   * Waits for tasks that are scheduled for immediate execution and rejects any
-   * tasks that are pending with a non-zero delay.
+   * Waits until all currently scheduled tasks are finished executing. Tasks
+   * schedule with a delay can be rejected or queued for immediate execution.
+   *
+   * @param executeDelayedTasks
    */
-  drain(): Promise<void> {
+  drain(executeDelayedTasks: boolean): Promise<void> {
     this.delayedOperations.forEach(entry => {
       if (entry.handle) {
         clearTimeout(entry.handle);
-        entry.deferred.reject(
-          new FirestoreError(Code.CANCELLED, 'Operation cancelled by shutdown')
-        );
+        if (executeDelayedTasks) {
+          this.scheduleInternal(entry.op).then(entry.deferred.resolve, entry.deferred.reject);
+        } else {
+          entry.deferred.reject(
+              new FirestoreError(Code.CANCELLED, 'Operation cancelled by shutdown')
+          );
+        }
       }
     });
     this.delayedOperations = [];

packages/firestore/test/integration/api/database.test.ts

@@ -18,7 +18,7 @@ import { expect } from 'chai';
 import * as firestore from 'firestore';
 
 import { Deferred } from '../../../src/util/promise';
-import { asyncIt } from '../../util/helpers';
+import {asyncIt, fasyncIt} from '../../util/helpers';
 import firebase from '../util/firebase_export';
 import {
   apiDescribe,
@@ -629,11 +629,11 @@ apiDescribe('Database', persistence => {
   asyncIt('can write document after idle timeout', () => {
     return withTestDb(persistence, (db, queue) => {
       const docRef = db.collection('test-collection').doc();
-      return queue
-        .awaitIdleTimeout(() => {
-          return docRef.set({ foo: 'bar' });
-        })
-        .then(() => docRef.set({ foo: 'bar' }));
+      return docRef.set({ foo: 'bar' }).then(() => {
+        console.log(JSON.stringify(queue));
+        expect(queue.delayedOperationsCount).to.be.equal(1);
+        return queue.drain(/* executeDelayedTasks= */ true);
+      }).then(() => docRef.set({ foo: 'bar' }));
     });
   });
 
@@ -653,11 +653,10 @@ apiDescribe('Database', persistence => {
         return deferred.promise.then(unregister);
       };
 
-      return queue
-        .awaitIdleTimeout(() => {
-          return awaitOnlineSnapshot();
-        })
-        .then(() => awaitOnlineSnapshot());
+      return awaitOnlineSnapshot().then(() => {
+        expect(queue.delayedOperationsCount).to.be.equal(1);
+        return queue.drain(/* executeDelayedTasks= */ true);
+      }).then(() => awaitOnlineSnapshot());
     });
   });
 });

packages/firestore/test/integration/remote/stream.test.ts

@@ -30,7 +30,7 @@ import {
 } from '../../../src/remote/watch_change';
 import { AsyncQueue } from '../../../src/util/async_queue';
 import { Deferred } from '../../../src/util/promise';
-import { asyncIt, setMutation } from '../../util/helpers';
+import { asyncIt, fasyncIt, setMutation } from '../../util/helpers';
 import { withTestDatastore } from '../util/helpers';
 import { TestQueue } from '../util/test_queue';
 
@@ -214,10 +214,10 @@ describe('Write Stream', () => {
 
   asyncIt('closes when idle', () => {
     let writeStream: PersistentWriteStream;
-    let queue: TestQueue;
+    let queue: AsyncQueue;
 
-    return withTestDatastore((ds, testQueue) => {
-      queue = testQueue;
+    return withTestDatastore((ds, q) => {
+      queue = q;
       writeStream = ds.newPersistentWriteStream();
       writeStream.start(streamListener);
       return streamListener.awaitCallback('open');
@@ -226,24 +226,25 @@ describe('Write Stream', () => {
         writeStream.writeHandshake();
         return streamListener.awaitCallback('handshakeComplete');
       })
-      .then(() =>
-        queue
-          .awaitIdleTimeout(() => {
-            writeStream.markIdle();
-            return streamListener.awaitCallback('close');
-          })
-          .then(() => {
-            expect(writeStream.isOpen()).to.be.false;
-          })
-      );
+      .then(() => {
+        writeStream.markIdle();
+        expect(queue.delayedOperationsCount).to.be.equal(1);
+          return Promise.all([
+           streamListener.awaitCallback('close'),
+              queue.drain(/*executeDelayedTasks=*/ true)
+              ]);
+        })
+      .then(() => {
+        expect(writeStream.isOpen()).to.be.false;
+      });
   });
 
   asyncIt('cancels idle on write', () => {
     let writeStream: PersistentWriteStream;
-    let queue: TestQueue;
+    let queue: AsyncQueue;
 
-    return withTestDatastore((ds, testQueue) => {
-      queue = testQueue;
+    return withTestDatastore((ds, q) => {
+      queue = q;
       writeStream = ds.newPersistentWriteStream();
       writeStream.start(streamListener);
       return streamListener.awaitCallback('open');
@@ -252,17 +253,15 @@ describe('Write Stream', () => {
         writeStream.writeHandshake();
         return streamListener.awaitCallback('handshakeComplete');
       })
-      .then(() =>
-        queue
-          .awaitIdleTimeout(() => {
-            // Mark the stream idle, but immediately cancel the idle timer by issuing another write.
-            writeStream.markIdle();
-            writeStream.writeMutations(SINGLE_MUTATION);
-            return streamListener.awaitCallback('mutationResult');
-          })
-          .then(() => {
-            expect(writeStream.isOpen()).to.be.true;
-          })
-      );
+      .then(() => {
+        // Mark the stream idle, but immediately cancel the idle timer by issuing another write.
+        writeStream.markIdle();
+        expect(queue.delayedOperationsCount).to.be.equal(1);
+        writeStream.writeMutations(SINGLE_MUTATION);
+        return streamListener.awaitCallback('mutationResult');
+      }).then(() => queue.drain(/*executeDelayedTasks=*/ true))
+        .then(() => {
+          expect(writeStream.isOpen()).to.be.true;
+    });
   });
 });

packages/firestore/test/integration/util/helpers.ts

@@ -16,17 +16,13 @@
 
 import * as firestore from 'firestore';
 
-import { FirebaseApp } from '@firebase/app';
-
 import { DatabaseId, DatabaseInfo } from '../../../src/core/database_info';
 import { Datastore } from '../../../src/remote/datastore';
 
 import firebase from './firebase_export';
 import { EmptyCredentialsProvider } from '../../../src/api/credentials';
 import { PlatformSupport } from '../../../src/platform/platform';
 import { AsyncQueue } from '../../../src/util/async_queue';
-import { Firestore, FirestoreDatabase } from '../../../src/api/database';
-import { TestQueue } from './test_queue';
 
 // tslint:disable-next-line:no-any __karma__ is an untyped global
 declare const __karma__: any;
@@ -59,23 +55,6 @@ function isIeOrEdge(): boolean {
   );
 }
 
-/** Implementation of Firestore that offers access to the underlying queue. */
-class TestFirestore extends Firestore {
-  _queue = new TestQueue();
-
-  constructor(databaseIdOrApp: FirestoreDatabase | FirebaseApp) {
-    super(databaseIdOrApp);
-  }
-
-  get queue() {
-    return this._queue;
-  }
-
-  protected initializeAsyncQueue(): AsyncQueue {
-    return this._queue;
-  }
-}
-
 export function isPersistenceAvailable(): boolean {
   return !isIeOrEdge();
 }
@@ -108,10 +87,10 @@ export function getDefaultDatabaseInfo(): DatabaseInfo {
 }
 
 export function withTestDatastore(
-  fn: (datastore: Datastore, queue: TestQueue) => Promise<void>
+  fn: (datastore: Datastore, queue: AsyncQueue) => Promise<void>
 ): Promise<void> {
   const databaseInfo = getDefaultDatabaseInfo();
-  const queue = new TestQueue();
+  const queue = new AsyncQueue();
   return PlatformSupport.getPlatform()
     .loadConnection(databaseInfo)
     .then(conn => {
@@ -132,25 +111,29 @@ export function withTestDatastore(
 
 export function withTestDb(
   persistence: boolean,
-  fn: (db: firestore.Firestore, queue: TestQueue) => Promise<void>
+  fn: (db: firestore.Firestore, queue: AsyncQueue) => Promise<void>
 ): Promise<void> {
   return withTestDbs(persistence, 1, ([db]) => {
-    return fn(db, (db as TestFirestore).queue);
+    // tslint:disable-next-line:no-any queue isn't exposed via d.ts
+    const firestoreInternal = db.INTERNAL as any;
+    return fn(db, firestoreInternal.queue);
   });
 }
 
 /** Runs provided fn with a db for an alternate project id. */
 export function withAlternateTestDb(
   persistence: boolean,
-  fn: (db: firestore.Firestore, queue: TestQueue) => Promise<void>
+  fn: (db: firestore.Firestore, queue: AsyncQueue) => Promise<void>
 ): Promise<void> {
   return withTestDbsSettings(
     persistence,
     ALT_PROJECT_ID,
     DEFAULT_SETTINGS,
     1,
     ([db]) => {
-      return fn(db, (db as TestFirestore).queue);
+      // tslint:disable-next-line:no-any queue isn't exposed via d.ts
+      const firestoreInternal = db.INTERNAL as any;
+      return fn(db, firestoreInternal.queue);
     }
   );
 }
@@ -189,10 +172,13 @@ export function withTestDbsSettings(
       { apiKey: 'fake-api-key', projectId },
       'test-app-' + appCount++
     );
-    const firestore = new TestFirestore(app);
+
+    // tslint:disable-next-line:no-any Firestore is not exposed in firebase.d.ts
+    const firebaseAny = firebase as any;
+    const firestore = firebaseAny.firestore(app);
     firestore.settings(settings);
 
-    let ready: Promise<TestFirestore>;
+    let ready: Promise<firestore.Firestore>;
     if (persistence) {
       ready = firestore.enablePersistence().then(() => firestore);
     } else {
@@ -202,7 +188,7 @@ export function withTestDbsSettings(
     promises.push(ready);
   }
 
-  return Promise.all(promises).then((dbs: TestFirestore[]) => {
+  return Promise.all(promises).then((dbs: firestore.Firestore[]) => {
     return fn(dbs)
       .then(wipeDb.bind(null, dbs[0]), error => {
         return wipeDb(dbs[0]).then(() => {

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

@@ -422,7 +422,7 @@ abstract class TestRunner {
     console.log('Running spec: ' + this.name);
     return sequence(steps, async step => {
       await this.doStep(step);
-      await this.queue.drain();
+      await this.queue.drain(/* executeDelayedTasks */ false);
       this.validateStepExpectations(step.expect!);
       this.validateStateExpectations(step.stateExpect!);
       this.eventList = [];