fix: Support more than 100 long-lived streams

dev/src/index.ts

@@ -58,6 +58,7 @@ import {
 import {WriteBatch} from './write-batch';
 
 import api = google.firestore.v1;
+import {Deferred} from '../test/util/helpers';
 
 export {
   CollectionReference,
@@ -1144,22 +1145,19 @@ export class Firestore {
    * @returns The given Stream once it is considered healthy.
    */
   private _initializeStream(
-    releaser: () => void,
     resultStream: NodeJS.ReadableStream,
     requestTag: string
-  ): Promise<NodeJS.ReadableStream>;
+  ): Promise<void>;
   private _initializeStream(
-    releaser: () => void,
     resultStream: NodeJS.ReadWriteStream,
     requestTag: string,
     request: {}
-  ): Promise<NodeJS.ReadWriteStream>;
+  ): Promise<void>;
   private _initializeStream(
-    releaser: () => void,
     resultStream: NodeJS.ReadableStream | NodeJS.ReadWriteStream,
     requestTag: string,
     request?: {}
-  ): Promise<NodeJS.ReadableStream | NodeJS.ReadWriteStream> {
+  ): Promise<void> {
     /** The last error we received and have not forwarded yet. */
     let errorReceived: Error | null = null;
 
@@ -1185,7 +1183,6 @@ export class Firestore {
             errorReceived
           );
           resultStream.emit('error', errorReceived);
-          releaser();
           errorReceived = null;
         } else if (!streamInitialized) {
           logger('Firestore._initializeStream', requestTag, 'Releasing stream');
@@ -1196,7 +1193,7 @@ export class Firestore {
           // 'end' event we intend to forward here. We therefore need to wait
           // until the API consumer registers their listeners (in the .then()
           // call) before emitting any further events.
-          resolve(resultStream);
+          resolve();
 
           // We execute the forwarding of the 'end' event via setTimeout() as
           // V8 guarantees that the above the Promise chain is resolved before
@@ -1209,7 +1206,6 @@ export class Firestore {
                 'Forwarding stream close'
               );
               resultStream.emit('end');
-              releaser();
             }
           }, 0);
         }
@@ -1251,7 +1247,6 @@ export class Firestore {
           );
           streamInitialized = true;
           reject(err);
-          releaser();
         } else {
           errorReceived = err;
         }
@@ -1359,36 +1354,47 @@ export class Firestore {
     const attempts = allowRetries ? MAX_REQUEST_RETRIES : 1;
     const callOptions = this.createCallOptions();
 
-    const gapicClient = this._clientPool.acquire();
-    const releaser = this._clientPool.createReleaser(gapicClient);
+    const result = new Deferred<NodeJS.ReadableStream>();
 
-    return this._retry(attempts, requestTag, () => {
-      return new Promise<NodeJS.ReadableStream>((resolve, reject) => {
-        try {
+    this._clientPool.run(gapicClient => {
+      // While we return the stream to the callee early, we don't want to
+      // release the GAPIC client until the callee has finished processing the
+      // stream.
+      const lifetime = new Deferred<void>();
+
+      this._retry(attempts, requestTag, async () => {
+        logger(
+          'Firestore.readStream',
+          requestTag,
+          'Sending request: %j',
+          request
+        );
+        const stream = gapicClient[methodName](request, callOptions);
+        const logStream = through2.obj(function(this, chunk, enc, callback) {
           logger(
             'Firestore.readStream',
             requestTag,
-            'Sending request: %j',
-            request
+            'Received response: %j',
+            chunk
           );
-          const stream = gapicClient[methodName](request, callOptions);
-          const logStream = through2.obj(function(this, chunk, enc, callback) {
-            logger(
-              'Firestore.readStream',
-              requestTag,
-              'Received response: %j',
-              chunk
-            );
-            this.push(chunk);
-            callback();
-          });
-          resolve(bun([stream, logStream]));
-        } catch (err) {
-          logger('Firestore.readStream', requestTag, 'Received error:', err);
-          reject(err);
-        }
-      }).then(stream => this._initializeStream(releaser, stream, requestTag));
+          this.push(chunk);
+          callback();
+        });
+
+        const resultStream = bun([stream, logStream]);
+        resultStream.on('close', lifetime.resolve);
+
+        await this._initializeStream(resultStream, requestTag);
+        result.resolve(resultStream);
+      }).catch(err => {
+        lifetime.resolve();
+        result.reject(err);
+      });
+
+      return lifetime.promise;
     });
+
+    return result.promise;
   }
 
   /**
@@ -1416,11 +1422,15 @@ export class Firestore {
     const attempts = allowRetries ? MAX_REQUEST_RETRIES : 1;
     const callOptions = this.createCallOptions();
 
-    const gapicClient = this._clientPool.acquire();
-    const releaser = this._clientPool.createReleaser(gapicClient);
+    const result = new Deferred<NodeJS.ReadWriteStream>();
+
+    this._clientPool.run(gapicClient => {
+      // While we return the stream to the callee early, we don't want to
+      // release the GAPIC client until the callee has finished processing the
+      // stream.
+      const lifetime = new Deferred<void>();
 
-    return this._retry(attempts, requestTag, () => {
-      return Promise.resolve().then(() => {
+      this._retry(attempts, requestTag, async () => {
         logger('Firestore.readWriteStream', requestTag, 'Opening stream');
         const requestStream = gapicClient[methodName](callOptions);
 
@@ -1436,14 +1446,19 @@ export class Firestore {
         });
 
         const resultStream = bun([requestStream, logStream]);
-        return this._initializeStream(
-          releaser,
-          resultStream,
-          requestTag,
-          request
-        );
+        resultStream.on('close', lifetime.resolve);
+        await this._initializeStream(resultStream, requestTag, request);
+
+        result.resolve(resultStream);
+      }).catch(err => {
+        lifetime.resolve();
+        result.reject(err);
       });
+
+      return lifetime.promise;
     });
+
+    return result.promise;
   }
 }
 

dev/src/pool.ts

@@ -50,7 +50,7 @@ export class ClientPool<T> {
    *
    * @private
    */
-  acquire(): T {
+  private acquire(): T {
     let selectedClient: T | null = null;
     let selectedRequestCount = 0;
 
@@ -79,7 +79,7 @@ export class ClientPool<T> {
    * removing it from the pool of active clients.
    * @private
    */
-  release(client: T): void {
+  private release(client: T): void {
     let requestCount = this.activeClients.get(client) || 0;
     assert(requestCount > 0, 'No active request');
 
@@ -91,26 +91,6 @@ export class ClientPool<T> {
     }
   }
 
-  /**
-   * Creates a new function that will release the given client, when called.
-   *
-   * This guarantees that the given client can only be released once.
-   *
-   * @private
-   */
-  createReleaser(client: T): () => void {
-    // Unfortunately, once the release() call is disconnected from the Promise
-    // returned from _initializeStream, there's no single callback in which the
-    // releaser can be guaranteed to be called once.
-    let released = false;
-    return () => {
-      if (!released) {
-        released = true;
-        this.release(client);
-      }
-    };
-  }
-
   /**
    * The number of currently registered clients.
    *

dev/system-test/firestore.ts

@@ -31,6 +31,7 @@ import {
   Timestamp,
 } from '../src';
 import {autoId} from '../src/util';
+import {Deferred} from '../test/util/helpers';
 
 const version = require('../../package.json').version;
 
@@ -913,6 +914,37 @@ describe('DocumentReference class', () => {
         maybeRun();
       });
     });
+
+    it('handles more than 100 concurrent listeners', async () => {
+      const ref = randomCol.doc('doc');
+
+      const emptyResults: Array<Deferred<void>> = [];
+      const documentResults: Array<Deferred<void>> = [];
+      const unsubscribeCallbacks: Array<() => void> = [];
+
+      // A single GAPIC client can only handle 100 concurrent streams. We set
+      // up 100+ long-lived listeners to verify that Firestore pools requests
+      // across multiple clients.
+      for (let i = 0; i < 150; ++i) {
+        emptyResults[i] = new Deferred<void>();
+        documentResults[i] = new Deferred<void>();
+
+        unsubscribeCallbacks[i] = randomCol
+          .where('i', '>', i)
+          .onSnapshot(snapshot => {
+            if (snapshot.size === 0) {
+              emptyResults[i].resolve();
+            } else if (snapshot.size === 1) {
+              documentResults[i].resolve();
+            }
+          });
+      }
+
+      await Promise.all(emptyResults.map(d => d.promise));
+      ref.set({i: 1337});
+      await Promise.all(documentResults.map(d => d.promise));
+      unsubscribeCallbacks.forEach(c => c());
+    });
   });
 });