Abort onSnapshotListeners on terminate.

Author: tom-andersen

packages/firestore/src/core/component_provider.ts

@@ -485,5 +485,6 @@ export class OnlineComponentProvider {
   async terminate(): Promise<void> {
     await remoteStoreShutdown(this.remoteStore);
     this.datastore?.terminate();
+    this.eventManager?.terminate();
   }
 }

packages/firestore/src/core/event_manager.ts

@@ -17,7 +17,7 @@
 
 import { debugAssert, debugCast } from '../util/assert';
 import { wrapInUserErrorIfRecoverable } from '../util/async_queue';
-import { FirestoreError } from '../util/error';
+import { Code, FirestoreError } from '../util/error';
 import { EventHandler } from '../util/misc';
 import { ObjectMap } from '../util/obj_map';
 
@@ -64,19 +64,17 @@ export interface EventManager {
   onUnlisten?: (query: Query, disableRemoteListen: boolean) => Promise<void>;
   onFirstRemoteStoreListen?: (query: Query) => Promise<void>;
   onLastRemoteStoreUnlisten?: (query: Query) => Promise<void>;
+  terminate(): void;
 }
 
 export function newEventManager(): EventManager {
   return new EventManagerImpl();
 }
 
 export class EventManagerImpl implements EventManager {
-  queries = new ObjectMap<Query, QueryListenersInfo>(
-    q => canonifyQuery(q),
-    queryEquals
-  );
+  queries: ObjectMap<Query, QueryListenersInfo> = newQueriesObjectMap();
 
-  onlineState = OnlineState.Unknown;
+  onlineState: OnlineState = OnlineState.Unknown;
 
   snapshotsInSyncListeners: Set<Observer<void>> = new Set();
 
@@ -98,6 +96,20 @@ export class EventManagerImpl implements EventManager {
    * still listening to the cache.
    */
   onLastRemoteStoreUnlisten?: (query: Query) => Promise<void>;
+
+  terminate(): void {
+    errorAllTargets(
+      this,
+      new FirestoreError(Code.ABORTED, 'Firestore shutting down')
+    );
+  }
+}
+
+function newQueriesObjectMap(): ObjectMap<Query, QueryListenersInfo> {
+  return new ObjectMap<Query, QueryListenersInfo>(
+    q => canonifyQuery(q),
+    queryEquals
+  );
 }
 
 function validateEventManager(eventManagerImpl: EventManagerImpl): void {
@@ -334,6 +346,20 @@ export function removeSnapshotsInSyncListener(
   eventManagerImpl.snapshotsInSyncListeners.delete(observer);
 }
 
+function errorAllTargets(eventManager: EventManager, error: FirestoreError) {
+  const eventManagerImpl = debugCast(eventManager, EventManagerImpl);
+  const queries = eventManagerImpl.queries;
+
+  // Prevent further access by clearing ObjectMap.
+  eventManagerImpl.queries = newQueriesObjectMap();
+
+  queries.forEach((_, queryInfo) => {
+    for (const listener of queryInfo.listeners) {
+      listener.onError(error);
+    }
+  });
+}
+
 // Call all global snapshot listeners that have been set.
 function raiseSnapshotsInSyncEvent(eventManagerImpl: EventManagerImpl): void {
   eventManagerImpl.snapshotsInSyncListeners.forEach(observer => {

packages/firestore/src/core/firestore_client.ts

@@ -118,6 +118,8 @@ export class FirestoreClient {
     appCheckToken: string,
     user: User
   ) => Promise<void> = () => Promise.resolve();
+  private _isShutdown: boolean = false;
+  private onTerminatedListener: (() => void) | undefined;
   _uninitializedComponentsProvider?: {
     _offline: OfflineComponentProvider;
     _offlineKind: 'memory' | 'persistent';
@@ -174,6 +176,13 @@ export class FirestoreClient {
     this.appCheckCredentialListener = listener;
   }
 
+  setOnTermintatedListener(listener: () => void): void {
+    this.onTerminatedListener = listener;
+    if (this._isShutdown) {
+      listener();
+    }
+  }
+
   /**
    * Checks that the client has not been terminated. Ensures that other methods on //
    * this class cannot be called after the client is terminated. //
@@ -188,9 +197,16 @@ export class FirestoreClient {
   }
 
   terminate(): Promise<void> {
+    if (this._isShutdown) {
+      return Promise.resolve();
+    }
     this.asyncQueue.enterRestrictedMode();
     const deferred = new Deferred();
     this.asyncQueue.enqueueAndForgetEvenWhileRestricted(async () => {
+      if (this._isShutdown) {
+        deferred.resolve();
+        return;
+      }
       try {
         if (this._onlineComponents) {
           await this._onlineComponents.terminate();
@@ -204,7 +220,11 @@ export class FirestoreClient {
         // tokens.
         this.authCredentials.shutdown();
         this.appCheckCredentials.shutdown();
+        this._isShutdown = true;
         deferred.resolve();
+        if (this.onTerminatedListener) {
+          this.onTerminatedListener();
+        }
       } catch (e) {
         const firestoreError = wrapInUserErrorIfRecoverable(
           e as Error,
@@ -239,7 +259,7 @@ export async function setOfflineComponentProvider(
   });
 
   // When a user calls clearPersistence() in one client, all other clients
-  // need to be terminated to allow the delete to succeed.
+  // need to be restarted to allow the delete to succeed.
   offlineComponentProvider.persistence.setDatabaseDeletedListener(() =>
     client.terminate()
   );

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

@@ -63,7 +63,8 @@ import {
   FieldPath,
   newTestFirestore,
   SnapshotOptions,
-  newTestApp
+  newTestApp,
+  FirestoreError
 } from '../util/firebase_export';
 import {
   apiDescribe,
@@ -1442,6 +1443,22 @@ apiDescribe('Database', persistence => {
     });
   });
 
+  it('query listener throws error on termination', async () => {
+    return withTestDoc(persistence, async (docRef, firestore) => {
+      const deferred: Deferred<FirestoreError> = new Deferred();
+      const unsubscribe = onSnapshot(docRef, snapshot => {}, deferred.resolve);
+
+      await terminate(firestore);
+
+      await expect(deferred.promise)
+        .to.eventually.haveOwnProperty('message')
+        .equal('Firestore shutting down');
+
+      // Call should proceed without error.
+      unsubscribe();
+    });
+  });
+
   it('can wait for pending writes', async () => {
     await withTestDoc(persistence, async (docRef, firestore) => {
       // Prevent pending writes receiving acknowledgement.