fix: fix duplicates in Query.stream() with back pressure (#1806)

dconeybe · web-flow · commit a5b680da034d · 2023-01-03T10:01:08.000-05:00
diff --git a/dev/src/reference.ts b/dev/src/reference.ts
@@ -17,6 +17,7 @@
 import * as firestore from '@google-cloud/firestore';
 import {Duplex, Readable, Transform} from 'stream';
 import * as deepEqual from 'fast-deep-equal';
+import {GoogleError} from 'google-gax';
 
 import * as protos from '../protos/firestore_v1_proto_api';
 
@@ -92,6 +93,8 @@ const comparisonOperators: {
   'array-contains-any': 'ARRAY_CONTAINS_ANY',
 };
 
+const NOOP_MESSAGE = Symbol('a noop message');
+
 /**
  * onSnapshot() callback that receives a QuerySnapshot.
  *
@@ -2268,6 +2271,11 @@ export class Query<T = firestore.DocumentData> implements firestore.Query<T> {
     return structuredQuery;
   }
 
+  // This method exists solely to enable unit tests to mock it.
+  _isPermanentRpcError(err: GoogleError, methodName: string): boolean {
+    return isPermanentRpcError(err, methodName);
+  }
+
   /**
    * Internal streaming method that accepts an optional transaction ID.
    *
@@ -2285,6 +2293,10 @@ export class Query<T = firestore.DocumentData> implements firestore.Query<T> {
     const stream = new Transform({
       objectMode: true,
       transform: (proto, enc, callback) => {
+        if (proto === NOOP_MESSAGE) {
+          callback(undefined);
+          return;
+        }
         const readTime = Timestamp.fromProto(proto.readTime);
         if (proto.document) {
           const document = this.firestore.snapshot_(
@@ -2337,27 +2349,33 @@ export class Query<T = firestore.DocumentData> implements firestore.Query<T> {
 
             // If a non-transactional query failed, attempt to restart.
             // Transactional queries are retried via the transaction runner.
-            if (!transactionId && !isPermanentRpcError(err, 'runQuery')) {
+            if (!transactionId && !this._isPermanentRpcError(err, 'runQuery')) {
               logger(
                 'Query._stream',
                 tag,
                 'Query failed with retryable stream error:',
                 err
               );
-              if (lastReceivedDocument) {
-                // Restart the query but use the last document we received as the
-                // query cursor. Note that we do not use backoff here. The call to
-                // `requestStream()` will backoff should the restart fail before
-                // delivering any results.
-                if (this._queryOptions.requireConsistency) {
-                  request = this.startAfter(lastReceivedDocument).toProto(
-                    lastReceivedDocument.readTime
-                  );
-                } else {
-                  request = this.startAfter(lastReceivedDocument).toProto();
+              // Enqueue a "no-op" write into the stream and resume the query
+              // once it is processed. This allows any enqueued results to be
+              // consumed before resuming the query so that the query resumption
+              // can start at the correct document.
+              stream.write(NOOP_MESSAGE, () => {
+                if (lastReceivedDocument) {
+                  // Restart the query but use the last document we received as
+                  // the query cursor. Note that we do not use backoff here. The
+                  // call to `requestStream()` will backoff should the restart
+                  // fail before delivering any results.
+                  if (this._queryOptions.requireConsistency) {
+                    request = this.startAfter(lastReceivedDocument).toProto(
+                      lastReceivedDocument.readTime
+                    );
+                  } else {
+                    request = this.startAfter(lastReceivedDocument).toProto();
+                  }
                 }
-              }
-              streamActive.resolve(/* active= */ true);
+                streamActive.resolve(/* active= */ true);
+              });
             } else {
               logger(
                 'Query._stream',
diff --git a/dev/test/query.ts b/dev/test/query.ts
@@ -35,6 +35,7 @@ import {DocumentSnapshot, DocumentSnapshotBuilder} from '../src/document';
 import {QualifiedResourcePath} from '../src/path';
 import {
   ApiOverride,
+  collect,
   createInstance,
   document,
   InvalidApiUsage,
@@ -49,8 +50,10 @@ import {
   writeResult,
 } from './util/helpers';
 
+import {GoogleError} from 'google-gax';
 import api = google.firestore.v1;
 import protobuf = google.protobuf;
+import {Deferred} from '../src/util';
 
 const PROJECT_ID = 'test-project';
 const DATABASE_ROOT = `projects/${PROJECT_ID}/databases/(default)`;
@@ -2507,3 +2510,120 @@ describe('collectionGroup queries', () => {
     });
   });
 });
+
+describe('query resumption', () => {
+  let firestore: Firestore;
+
+  beforeEach(() => {
+    setTimeoutHandler(setImmediate);
+    return createInstance().then(firestoreInstance => {
+      firestore = firestoreInstance;
+    });
+  });
+
+  afterEach(async () => {
+    await verifyInstance(firestore);
+    setTimeoutHandler(setTimeout);
+  });
+
+  // Prevent regression of
+  // https://github.com/googleapis/nodejs-firestore/issues/1790
+  it('results should not be double produced on retryable error with back pressure', async () => {
+    // Generate the IDs of the documents that will match the query.
+    const documentIds = Array.from(new Array(500), (_, index) => `doc${index}`);
+
+    // Finds the index in `documentIds` of the document referred to in the
+    // "startAt" of the given request.
+    function getStartAtDocumentIndex(
+      request: api.IRunQueryRequest
+    ): number | null {
+      const startAt = request.structuredQuery?.startAt;
+      const startAtValue = startAt?.values?.[0]?.referenceValue;
+      const startAtBefore = startAt?.before;
+      if (typeof startAtValue !== 'string') {
+        return null;
+      }
+      const docId = startAtValue.split('/').pop()!;
+      const docIdIndex = documentIds.indexOf(docId);
+      if (docIdIndex < 0) {
+        return null;
+      }
+      return startAtBefore ? docIdIndex : docIdIndex + 1;
+    }
+
+    const RETRYABLE_ERROR_DOMAIN = 'RETRYABLE_ERROR_DOMAIN';
+
+    // A mock replacement for Query._isPermanentRpcError which (a) resolves
+    // a promise once invoked and (b) treats a specific error "domain" as
+    // non-retryable.
+    function mockIsPermanentRpcError(err: GoogleError): boolean {
+      mockIsPermanentRpcError.invoked.resolve(true);
+      return err?.domain !== RETRYABLE_ERROR_DOMAIN;
+    }
+    mockIsPermanentRpcError.invoked = new Deferred();
+
+    // Return the first half of the documents, followed by a retryable error.
+    function* getRequest1Responses(): Generator<api.IRunQueryResponse | Error> {
+      const runQueryResponses = documentIds
+        .slice(0, documentIds.length / 2)
+        .map(documentId => result(documentId));
+      for (const runQueryResponse of runQueryResponses) {
+        yield runQueryResponse;
+      }
+      const retryableError = new GoogleError('simulated retryable error');
+      retryableError.domain = RETRYABLE_ERROR_DOMAIN;
+      yield retryableError;
+    }
+
+    // Return the remaining documents.
+    function* getRequest2Responses(
+      request: api.IRunQueryRequest
+    ): Generator<api.IRunQueryResponse> {
+      const startAtDocumentIndex = getStartAtDocumentIndex(request);
+      if (startAtDocumentIndex === null) {
+        throw new Error('request #2 should specify a valid startAt');
+      }
+      const runQueryResponses = documentIds
+        .slice(startAtDocumentIndex)
+        .map(documentId => result(documentId));
+      for (const runQueryResponse of runQueryResponses) {
+        yield runQueryResponse;
+      }
+    }
+
+    // Set up the mocked responses from Watch.
+    let requestNum = 0;
+    const overrides: ApiOverride = {
+      runQuery: request => {
+        requestNum++;
+        switch (requestNum) {
+          case 1:
+            return stream(...getRequest1Responses());
+          case 2:
+            return stream(...getRequest2Responses(request!));
+          default:
+            throw new Error(`should never get here (requestNum=${requestNum})`);
+        }
+      },
+    };
+
+    // Create an async iterator to get the result set but DO NOT iterate over
+    // it immediately. Instead, allow the responses to pile up and fill the
+    // buffers. Once isPermanentError() is invoked, indicating that the first
+    // request has failed and is about to be retried, collect the results from
+    // the async iterator into an array.
+    firestore = await createInstance(overrides);
+    const query = firestore.collection('collectionId');
+    query._isPermanentRpcError = mockIsPermanentRpcError;
+    const iterator = query
+      .stream()
+      [Symbol.asyncIterator]() as AsyncIterator<QueryDocumentSnapshot>;
+    await mockIsPermanentRpcError.invoked.promise;
+    const snapshots = await collect(iterator);
+
+    // Verify that the async iterator returned the correct documents and,
+    // especially, does not have duplicate results.
+    const actualDocumentIds = snapshots.map(snapshot => snapshot.id);
+    expect(actualDocumentIds).to.eql(documentIds);
+  });
+});
diff --git a/dev/test/util/helpers.ts b/dev/test/util/helpers.ts
@@ -413,3 +413,34 @@ export async function bundleToElementArray(
   }
   return result;
 }
+
+/**
+ * Reads the elements of an AsyncIterator.
+ *
+ * Example:
+ *
+ * const query = firestore.collection('collectionId');
+ * const iterator = query.stream()[Symbol.asyncIterator]()
+ *   as AsyncIterator<QueryDocumentSnapshot>;
+ * return collect(iterator).then(snapshots => {
+ *   expect(snapshots).to.have.length(2);
+ * });
+ *
+ * @param iterator the iterator whose elements over which to iterate.
+ * @return a Promise that is fulfilled with the elements that were produced, or
+ * is rejected with the cause of the first failed iteration.
+ */
+export async function collect<T, TReturn, TNext>(
+  iterator: AsyncIterator<T, TReturn, TNext>
+): Promise<Array<T>> {
+  const values: Array<T> = [];
+  // eslint-disable-next-line no-constant-condition
+  while (true) {
+    const {done, value} = await iterator.next();
+    if (done) {
+      break;
+    }
+    values.push(value);
+  }
+  return values;
+}
