fix: update retry policy to not retry streams that have not made progress (#1946)

MarkDuckworth · web-flow · commit 170d05b1fa6c · 2023-11-17T21:06:14.000Z
fix: update retry policy to not retry streams with retryable error that have not made progress receiving documents
diff --git a/dev/src/index.ts b/dev/src/index.ts
@@ -1637,7 +1637,7 @@ export class Firestore implements firestore.Firestore {
       function streamReady(): void {
         if (!streamInitialized) {
           streamInitialized = true;
-          logger('Firestore._initializeStream', requestTag, 'Releasing stream');
+          logger('Firestore._initializeStream', requestTag, 'Stream ready');
           resolve(resultStream);
         }
       }
diff --git a/dev/src/reference.ts b/dev/src/reference.ts
@@ -44,6 +44,7 @@ import {defaultConverter} from './types';
 import {
   autoId,
   Deferred,
+  getTotalTimeout,
   isPermanentRpcError,
   mapToArray,
   requestTag,
@@ -2569,6 +2570,15 @@ export class Query<
     return isPermanentRpcError(err, methodName);
   }
 
+  _hasRetryTimedOut(methodName: string, startTime: number): boolean {
+    const totalTimeout = getTotalTimeout(methodName);
+    if (totalTimeout === 0) {
+      return false;
+    }
+
+    return Date.now() - startTime >= totalTimeout;
+  }
+
   /**
    * Internal streaming method that accepts an optional transaction ID.
    *
@@ -2579,6 +2589,7 @@ export class Query<
    */
   _stream(transactionId?: Uint8Array): NodeJS.ReadableStream {
     const tag = requestTag();
+    const startTime = Date.now();
 
     let lastReceivedDocument: QueryDocumentSnapshot<
       AppModelType,
@@ -2638,8 +2649,9 @@ export class Query<
         let streamActive: Deferred<boolean>;
         do {
           streamActive = new Deferred<boolean>();
+          const methodName = 'runQuery';
           backendStream = await this._firestore.requestStream(
-            'runQuery',
+            methodName,
             /* bidirectional= */ false,
             request,
             tag
@@ -2656,12 +2668,28 @@ export class Query<
                 'Query failed with retryable stream error:',
                 err
               );
-              // 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.
+
+              // Enqueue a "no-op" write into the stream and wait for it to be
+              // read by the downstream consumer. This ensures that all enqueued
+              // results in the stream are consumed, which will give us an accurate
+              // value for `lastReceivedDocument`.
               stream.write(NOOP_MESSAGE, () => {
-                if (lastReceivedDocument) {
+                if (this._hasRetryTimedOut(methodName, startTime)) {
+                  logger(
+                    'Query._stream',
+                    tag,
+                    'Query failed with retryable stream error but the total retry timeout has exceeded.'
+                  );
+                  stream.destroy(err);
+                  streamActive.resolve(/* active= */ false);
+                } else if (lastReceivedDocument) {
+                  logger(
+                    'Query._stream',
+                    tag,
+                    'Query failed with retryable stream error and progress was made receiving ' +
+                      'documents, so the stream is being retried.'
+                  );
+
                   // 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
@@ -2673,8 +2701,21 @@ export class Query<
                   } else {
                     request = this.startAfter(lastReceivedDocument).toProto();
                   }
+
+                  // Set lastReceivedDocument to null before each retry attempt to ensure the retry makes progress
+                  lastReceivedDocument = null;
+
+                  streamActive.resolve(/* active= */ true);
+                } else {
+                  logger(
+                    'Query._stream',
+                    tag,
+                    'Query failed with retryable stream error however no progress was made receiving ' +
+                      'documents, so the stream is being closed.'
+                  );
+                  stream.destroy(err);
+                  streamActive.resolve(/* active= */ false);
                 }
-                streamActive.resolve(/* active= */ true);
               });
             } else {
               logger(
@@ -3320,48 +3361,33 @@ export class AggregateQuery<
         // catch below.
         const request = this.toProto(transactionId);
 
-        let streamActive: Deferred<boolean>;
-        do {
-          streamActive = new Deferred<boolean>();
-          const backendStream = await firestore.requestStream(
-            'runAggregationQuery',
-            /* bidirectional= */ false,
-            request,
-            tag
-          );
-          stream.on('close', () => {
-            backendStream.resume();
-            backendStream.end();
-          });
-          backendStream.on('error', err => {
-            backendStream.unpipe(stream);
-            // If a non-transactional query failed, attempt to restart.
-            // Transactional queries are retried via the transaction runner.
-            if (
-              !transactionId &&
-              !isPermanentRpcError(err, 'runAggregationQuery')
-            ) {
-              logger(
-                'AggregateQuery._stream',
-                tag,
-                'AggregateQuery failed with retryable stream error:',
-                err
-              );
-              streamActive.resolve(/* active= */ true);
-            } else {
-              logger(
-                'AggregateQuery._stream',
-                tag,
-                'AggregateQuery failed with stream error:',
-                err
-              );
-              stream.destroy(err);
-              streamActive.resolve(/* active= */ false);
-            }
-          });
+        const backendStream = await firestore.requestStream(
+          'runAggregationQuery',
+          /* bidirectional= */ false,
+          request,
+          tag
+        );
+        stream.on('close', () => {
           backendStream.resume();
-          backendStream.pipe(stream);
-        } while (await streamActive.promise);
+          backendStream.end();
+        });
+        backendStream.on('error', err => {
+          // TODO(group-by) When group-by queries are supported for aggregates
+          // consider implementing retries if the stream is making progress
+          // receiving results for groups. See the use of lastReceivedDocument
+          // in the retry strategy for runQuery.
+
+          backendStream.unpipe(stream);
+          logger(
+            'AggregateQuery._stream',
+            tag,
+            'AggregateQuery failed with stream error:',
+            err
+          );
+          stream.destroy(err);
+        });
+        backendStream.resume();
+        backendStream.pipe(stream);
       })
       .catch(e => stream.destroy(e));
 
diff --git a/dev/src/util.ts b/dev/src/util.ts
@@ -178,6 +178,21 @@ export function getRetryCodes(methodName: string): number[] {
   return getServiceConfig(methodName)?.retry?.retryCodes ?? [];
 }
 
+/**
+ * Gets the total timeout in milliseconds from the retry settings in
+ * the service config for the given RPC. If the total timeout is not
+ * set, then `0` is returned.
+ *
+ * @private
+ * @internal
+ */
+export function getTotalTimeout(methodName: string): number {
+  return (
+    getServiceConfig(methodName)?.retry?.backoffSettings?.totalTimeoutMillis ??
+    0
+  );
+}
+
 /**
  * Returns the backoff setting from the service configuration.
  * @private
diff --git a/dev/test/aggregateQuery.ts b/dev/test/aggregateQuery.ts
@@ -126,19 +126,31 @@ describe('aggregate query interface', () => {
     });
   });
 
-  it('handles stream exception at initialization', () => {
+  it('handles stream exception at initialization', async () => {
+    let attempts = 0;
     const query = firestore.collection('collectionId').count();
 
     query._stream = () => {
+      ++attempts;
       throw new Error('Expected error');
     };
 
-    return expect(query.get()).to.eventually.rejectedWith('Expected error');
+    await query
+      .get()
+      .then(() => {
+        throw new Error('Unexpected success in Promise');
+      })
+      .catch(err => {
+        expect(err.message).to.equal('Expected error');
+        expect(attempts).to.equal(1);
+      });
   });
 
   it('handles stream exception during initialization', async () => {
+    let attempts = 0;
     const overrides: ApiOverride = {
       runAggregationQuery: () => {
+        ++attempts;
         return stream(new Error('Expected error'));
       },
     };
@@ -152,6 +164,31 @@ describe('aggregate query interface', () => {
       })
       .catch(err => {
         expect(err.message).to.equal('Expected error');
+        expect(attempts).to.equal(5);
+      });
+  });
+
+  it('handles message without result during initialization', async () => {
+    let attempts = 0;
+    const overrides: ApiOverride = {
+      runAggregationQuery: () => {
+        ++attempts;
+        return stream({readTime: {seconds: 5, nanos: 6}});
+      },
+    };
+    firestore = await createInstance(overrides);
+
+    const query = firestore.collection('collectionId').count();
+    await query
+      .get()
+      .then(() => {
+        throw new Error('Unexpected success in Promise');
+      })
+      .catch(err => {
+        expect(err.message).to.equal(
+          'RunAggregationQueryResponse is missing result'
+        );
+        expect(attempts).to.equal(1);
       });
   });
 });
diff --git a/dev/test/query.ts b/dev/test/query.ts

Original file line number	Diff line number	Diff line change
`@@ -1637,7 +1637,7 @@ export class Firestore implements firestore.Firestore {`
`1637`	`1637`	`function streamReady(): void {`
`1638`	`1638`	`if (!streamInitialized) {`
`1639`	`1639`	`streamInitialized = true;`
`1640`		`- logger('Firestore._initializeStream', requestTag, 'Releasing stream');`
	`1640`	`+ logger('Firestore._initializeStream', requestTag, 'Stream ready');`
`1641`	`1641`	`resolve(resultStream);`
`1642`	`1642`	`}`
`1643`	`1643`	`}`