From 0d2f97ee78461273ceddfb211f536f13cbd29523 Mon Sep 17 00:00:00 2001
From: Sebastian Schmidt <mrschmidt@google.com>
Date: Tue, 8 Oct 2019 13:11:47 -0700
Subject: [PATCH] Call GrpcClient.close() when we dispose of clients

---
 dev/src/index.ts               |  5 +++--
 dev/src/pool.ts                |  8 ++++++--
 dev/src/v1/firestore_client.js | 13 +++++++++++--
 dev/test/pool.ts               | 17 +++++++++++++++++
 4 files changed, 37 insertions(+), 6 deletions(-)

diff --git a/dev/src/index.ts b/dev/src/index.ts
index 7c7d66e60..617b1b6b0 100644
--- a/dev/src/index.ts
+++ b/dev/src/index.ts
@@ -388,7 +388,7 @@ export class Firestore {
 
     this._clientPool = new ClientPool(
       MAX_CONCURRENT_REQUESTS_PER_CLIENT,
-      () => {
+      /* clientFactory= */ () => {
         let client: GapicClient;
 
         if (this._settings.ssl === false) {
@@ -401,7 +401,8 @@ export class Firestore {
 
         logger('Firestore', null, 'Initialized Firestore GAPIC Client');
         return client;
-      }
+      },
+        /* clientDestructor= */  (client: GapicClient) => client.close()
     );
 
     logger('Firestore', null, 'Initialized Firestore');
diff --git a/dev/src/pool.ts b/dev/src/pool.ts
index fc229dac2..e37a6f5cd 100644
--- a/dev/src/pool.ts
+++ b/dev/src/pool.ts
@@ -40,11 +40,14 @@ export class ClientPool<T> {
    * can handle.
    * @param clientFactory A factory function called as needed when new clients
    * are required.
+   * @param clientDestructor A cleanup function that is called for each client
+   * when the client is disposed of.
    */
   constructor(
     private readonly concurrentOperationLimit: number,
-    private readonly clientFactory: () => T
-  ) {}
+    private readonly clientFactory: () => T,
+    private readonly clientDestructor: (client:T) => void = () => {})
+   {}
 
   /**
    * Returns an already existing client if it has less than the maximum number
@@ -171,6 +174,7 @@ export class ClientPool<T> {
         ++idleClients;
 
         if (idleClients > 1) {
+          this.clientDestructor(client);
           this.activeClients.delete(client);
         }
       }
diff --git a/dev/src/v1/firestore_client.js b/dev/src/v1/firestore_client.js
index e3ee2630c..21d2066d3 100644
--- a/dev/src/v1/firestore_client.js
+++ b/dev/src/v1/firestore_client.js
@@ -196,7 +196,7 @@ class FirestoreClient {
 
     // Put together the "service stub" for
     // google.firestore.v1.Firestore.
-    const firestoreStub = gaxGrpc.createStub(
+    this._firestoreStub = gaxGrpc.createStub(
       opts.fallback
         ? protos.lookupService('google.firestore.v1.Firestore')
         : protos.google.firestore.v1.Firestore,
@@ -221,7 +221,7 @@ class FirestoreClient {
       'listCollectionIds',
     ];
     for (const methodName of firestoreStubMethods) {
-      const innerCallPromise = firestoreStub.then(
+      const innerCallPromise = this._firestoreStub.then(
         stub => (...args) => {
           return stub[methodName].apply(stub, args);
         },
@@ -271,6 +271,15 @@ class FirestoreClient {
     ];
   }
 
+  /**
+   * Terminate the GRPC channel and close the client.
+   *
+   * The client will no longer be usable and all future behavior is undefined.
+   */
+  close() {
+    this._firestoreStub.then(grpcClient => grpcClient.close());
+  }
+
   /**
    * Return the project ID used by this class.
    * @param {function(Error, string)} callback - the callback to
diff --git a/dev/test/pool.ts b/dev/test/pool.ts
index 15c5e055f..06b33e3a0 100644
--- a/dev/test/pool.ts
+++ b/dev/test/pool.ts
@@ -152,6 +152,23 @@ describe('Client pool', () => {
     );
   });
 
+  it('garbage collection calls destructor', () => {
+    const garbageCollect = new Deferred();
+
+    const clientPool = new ClientPool<{}>(1, () => {
+      return {};
+    }, () => garbageCollect.resolve());
+
+    const operationPromises = deferredPromises(2);
+
+    clientPool.run(REQUEST_TAG, () => operationPromises[0].promise);
+    clientPool.run(REQUEST_TAG, () => operationPromises[1].promise);
+
+    operationPromises.forEach(deferred => deferred.resolve());
+
+    return garbageCollect.promise;
+  });
+
   it('forwards success', () => {
     const clientPool = new ClientPool<{}>(1, () => {
       return {};