From 9ef582aa0508a3d02fb036f741c8c51e5ff4307c Mon Sep 17 00:00:00 2001
From: Brian Chen <chenbrian@google.com>
Date: Thu, 19 Dec 2019 15:45:35 -0800
Subject: [PATCH] feat: add ability to close channels (#824)

---
 dev/src/index.ts             |  9 +++++++++
 dev/src/pool.ts              | 17 +++++++++++++++++
 dev/system-test/firestore.ts | 18 ++++++++++++++++++
 dev/test/pool.ts             | 17 +++++++++++++++++
 types/firestore.d.ts         |  7 +++++++
 5 files changed, 68 insertions(+)

diff --git a/dev/src/index.ts b/dev/src/index.ts
index af3498ed0..f5c064137 100644
--- a/dev/src/index.ts
+++ b/dev/src/index.ts
@@ -1007,6 +1007,15 @@ export class Firestore {
       });
   }
 
+  /**
+   * Terminates the Firestore client and closes all open streams.
+   *
+   * @return A Promise that resolves when the client is terminated.
+   */
+  terminate(): Promise<void> {
+    return this._clientPool.terminate();
+  }
+
   /**
    * Initializes the client if it is not already initialized. All methods in the
    * SDK can be used after this method completes.
diff --git a/dev/src/pool.ts b/dev/src/pool.ts
index b7c861b67..4010817fa 100644
--- a/dev/src/pool.ts
+++ b/dev/src/pool.ts
@@ -35,6 +35,12 @@ export class ClientPool<T> {
    */
   private activeClients: Map<T, number> = new Map();
 
+  /**
+   * Whether the Firestore instance has been terminated. Once terminated, the
+   * ClientPool can longer schedule new operations.
+   */
+  private terminated = false;
+
   /**
    * @param concurrentOperationLimit The number of operations that each client
    * can handle.
@@ -148,6 +154,9 @@ export class ClientPool<T> {
    * @private
    */
   run<V>(requestTag: string, op: (client: T) => Promise<V>): Promise<V> {
+    if (this.terminated) {
+      throw new Error('The client has already been terminated');
+    }
     const client = this.acquire(requestTag);
 
     return op(client)
@@ -161,6 +170,14 @@ export class ClientPool<T> {
       });
   }
 
+  async terminate(): Promise<void> {
+    this.terminated = true;
+    for (const [client, _requestCount] of this.activeClients) {
+      this.activeClients.delete(client);
+      await this.clientDestructor(client);
+    }
+  }
+
   /**
    * Deletes clients that are no longer executing operations. Keeps up to one
    * idle client to reduce future initialization costs.
diff --git a/dev/system-test/firestore.ts b/dev/system-test/firestore.ts
index 5fa028b31..e5b45af62 100644
--- a/dev/system-test/firestore.ts
+++ b/dev/system-test/firestore.ts
@@ -127,6 +127,24 @@ describe('Firestore class', () => {
         expect(docs[1].data()).to.deep.equal({f: 'a'});
       });
   });
+
+  it('cannot make calls after the client has been terminated', () => {
+    const ref1 = randomCol.doc('doc1');
+    ref1.onSnapshot(snapshot => {
+      return Promise.reject('onSnapshot() should be called');
+    });
+    return firestore
+      .terminate()
+      .then(() => {
+        return ref1.set({foo: 100});
+      })
+      .then(() => {
+        Promise.reject('set() should have failed');
+      })
+      .catch(err => {
+        expect(err.message).to.equal('The client has already been terminated');
+      });
+  });
 });
 
 describe('CollectionReference class', () => {
diff --git a/dev/test/pool.ts b/dev/test/pool.ts
index b9b2afecf..cd50ff73b 100644
--- a/dev/test/pool.ts
+++ b/dev/test/pool.ts
@@ -191,4 +191,21 @@ describe('Client pool', () => {
     );
     return expect(op).to.eventually.be.rejectedWith('Generated error');
   });
+
+  it('rejects subsequent operations after being terminated', () => {
+    const clientPool = new ClientPool<{}>(1, () => {
+      return {};
+    });
+
+    return clientPool
+      .terminate()
+      .then(() => {
+        clientPool.run(REQUEST_TAG, () =>
+          Promise.reject('Call to run() should have failed')
+        );
+      })
+      .catch((err: Error) => {
+        expect(err.message).to.equal('The client has already been terminated');
+      });
+  });
 });
diff --git a/types/firestore.d.ts b/types/firestore.d.ts
index 2cc5796b3..3a7122922 100644
--- a/types/firestore.d.ts
+++ b/types/firestore.d.ts
@@ -154,6 +154,13 @@ declare namespace FirebaseFirestore {
     getAll(...documentRefsOrReadOptions: Array<DocumentReference|ReadOptions>):
         Promise<DocumentSnapshot[]>;
 
+    /**
+     * Terminates the Firestore client and closes all open streams.
+     *
+     * @return A Promise that resolves when the client is terminated.
+     */
+    terminate(): Promise<void>;
+
     /**
      * Fetches the root collections that are associated with this Firestore
      * database.