Allow failing of individual transactions (#3052)

Author: schmidt-sebastian

packages/firestore/test/unit/specs/recovery_spec.test.ts

@@ -34,7 +34,11 @@ describeSpec('Persistence Recovery', ['no-ios', 'no-android'], () => {
           .client(1)
           .expectPrimaryState(false)
           .userSets('collection/a', { v: 1 })
-          .failDatabase()
+          .failDatabaseTransactions({
+            'Locally write mutations': true,
+            'Synchronize last document change read time': true,
+            'Lookup mutation documents': true
+          })
           .client(0)
           .writeAcks('collection/a', 1, { expectUserCallback: false })
           .client(1)
@@ -64,7 +68,11 @@ describeSpec('Persistence Recovery', ['no-ios', 'no-android'], () => {
         .client(1)
         .expectPrimaryState(false)
         .userListens(query)
-        .failDatabase()
+        .failDatabaseTransactions({
+          'Allocate target': true,
+          'Lookup mutation documents': true,
+          'Get new document changes': true
+        })
         .client(0)
         .expectListen(query)
         .watchAcksFull(query, 1000)
@@ -84,7 +92,10 @@ describeSpec('Persistence Recovery', ['no-ios', 'no-android'], () => {
       return (
         client(0)
           .expectPrimaryState(true)
-          .failDatabase()
+          .failDatabaseTransactions({
+            'Allocate target': true,
+            'Get target data': true
+          })
           .client(1)
           .userListens(query)
           .client(0)
@@ -94,7 +105,10 @@ describeSpec('Persistence Recovery', ['no-ios', 'no-android'], () => {
           .recoverDatabase()
           .runTimer(TimerId.AsyncQueueRetry)
           .expectListen(query)
-          .failDatabase()
+          .failDatabaseTransactions({
+            'Allocate target': true,
+            'Release target': true
+          })
           .client(1)
           .userUnlistens(query)
           .client(0)
@@ -113,7 +127,12 @@ describeSpec('Persistence Recovery', ['no-ios', 'no-android'], () => {
     return spec()
       .userSets('collection/key1', { foo: 'a' })
       .expectNumOutstandingWrites(1)
-      .failDatabase()
+      .failDatabaseTransactions({
+        'Locally write mutations': true,
+        notifyLocalViewChanges: true,
+        'Get next mutation batch': true,
+        'Get last stream token': true
+      })
       .userSets('collection/key2', { bar: 'b' })
       .expectUserCallbacks({ rejected: ['collection/key2'] })
       .recoverDatabase()
@@ -149,7 +168,11 @@ describeSpec('Persistence Recovery', ['no-ios', 'no-android'], () => {
         fromCache: true,
         hasPendingWrites: true
       })
-      .failDatabase()
+      .failDatabaseTransactions({
+        'Locally write mutations': true,
+        notifyLocalViewChanges: true,
+        'Get next mutation batch': true
+      })
       .userSets('collection/key2', { foo: 'b' })
       .expectUserCallbacks({ rejected: ['collection/key2'] })
       .recoverDatabase()
@@ -173,7 +196,7 @@ describeSpec('Persistence Recovery', ['no-ios', 'no-android'], () => {
       .userListens(query1)
       .watchAcksFull(query1, 1)
       .expectEvents(query1, {})
-      .failDatabase()
+      .failDatabaseTransactions({ 'Allocate target': true })
       .userListens(query2)
       .expectEvents(query2, { errorCode: Code.UNAVAILABLE })
       .recoverDatabase()
@@ -189,7 +212,7 @@ describeSpec('Persistence Recovery', ['no-ios', 'no-android'], () => {
       .userListens(query1)
       .watchAcksFull(query1, 1)
       .expectEvents(query1, {})
-      .failDatabase()
+      .failDatabaseTransactions({ 'Allocate target': true })
       .userListens(query2)
       .expectEvents(query2, { errorCode: Code.UNAVAILABLE })
       .recoverDatabase()
@@ -213,7 +236,10 @@ describeSpec('Persistence Recovery', ['no-ios', 'no-android'], () => {
             added: [doc1]
           })
           .watchSends({ affects: [query] }, doc2)
-          .failDatabase()
+          .failDatabaseTransactions({
+            'Get last remote snapshot version': true,
+            'Release target': true
+          })
           .watchSnapshots(1500)
           // `failDatabase()` causes us to go offline.
           .expectActiveTargets()
@@ -250,7 +276,10 @@ describeSpec('Persistence Recovery', ['no-ios', 'no-android'], () => {
           .expectEvents(doc2Query, {
             added: [doc2]
           })
-          .failDatabase()
+          .failDatabaseTransactions({
+            'Get last remote snapshot version': true,
+            'Release target': true
+          })
           .watchRemoves(
             doc1Query,
             new RpcError(Code.PERMISSION_DENIED, 'Simulated target error')

packages/firestore/test/unit/specs/spec_builder.ts

@@ -45,6 +45,7 @@ import {
   parseQuery,
   runSpec,
   SpecConfig,
+  SpecDatabaseFailures,
   SpecDocument,
   SpecQuery,
   SpecQueryFilter,
@@ -435,12 +436,15 @@ export class SpecBuilder {
     return this;
   }
 
-  /** Fails all database operations until `recoverDatabase()` is called. */
-  failDatabase(): this {
+  /**
+   * Fails the specified database transaction until `recoverDatabase()` is
+   * called.
+   */
+  failDatabaseTransactions(failureMode: SpecDatabaseFailures): this {
     this.nextStep();
     this.injectFailures = true;
     this.currentStep = {
-      failDatabase: true
+      failDatabase: failureMode
     };
     return this;
   }

packages/firestore/test/unit/specs/spec_test_components.ts

@@ -36,6 +36,7 @@ import {
   MemoryPersistence
 } from '../../../src/local/memory_persistence';
 import { LruParams } from '../../../src/local/lru_garbage_collector';
+import { PersistenceAction, SpecDatabaseFailures } from './spec_test_runner';
 import { Connection, Stream } from '../../../src/remote/connection';
 import { StreamBridge } from '../../../src/remote/stream_bridge';
 import * as api from '../../../src/protos/firestore_proto_api';
@@ -56,7 +57,7 @@ import { expect } from 'chai';
  * transaction failures.
  */
 export class MockMemoryPersistence extends MemoryPersistence {
-  injectFailures = false;
+  injectFailures?: SpecDatabaseFailures;
 
   runTransaction<T>(
     action: string,
@@ -65,13 +66,8 @@ export class MockMemoryPersistence extends MemoryPersistence {
       transaction: PersistenceTransaction
     ) => PersistencePromise<T>
   ): Promise<T> {
-    if (this.injectFailures) {
-      return Promise.reject(
-        new IndexedDbTransactionError(new Error('Simulated retryable error'))
-      );
-    } else {
-      return super.runTransaction(action, mode, transactionOperation);
-    }
+    failTransactionIfNeeded(this.injectFailures, action);
+    return super.runTransaction(action, mode, transactionOperation);
   }
 }
 
@@ -80,7 +76,7 @@ export class MockMemoryPersistence extends MemoryPersistence {
  * transaction failures.
  */
 export class MockIndexedDbPersistence extends IndexedDbPersistence {
-  injectFailures = false;
+  injectFailures?: SpecDatabaseFailures;
 
   runTransaction<T>(
     action: string,
@@ -89,12 +85,27 @@ export class MockIndexedDbPersistence extends IndexedDbPersistence {
       transaction: PersistenceTransaction
     ) => PersistencePromise<T>
   ): Promise<T> {
-    if (this.injectFailures) {
-      return Promise.reject(
-        new IndexedDbTransactionError(new Error('Simulated retryable error'))
+    failTransactionIfNeeded(this.injectFailures, action);
+    return super.runTransaction(action, mode, transactionOperation);
+  }
+}
+
+/**
+ * Shared failure handler between MockIndexedDbPersistence and
+ * MockMemoryPersistence that can inject transaction failures.
+ */
+function failTransactionIfNeeded(
+  config: SpecDatabaseFailures | undefined,
+  action: string
+): void {
+  if (config) {
+    const shouldFail = config[action as PersistenceAction];
+    if (shouldFail === undefined) {
+      throw fail('Failure mode not specified for action: ' + action);
+    } else if (shouldFail) {
+      throw new IndexedDbTransactionError(
+        new Error('Simulated retryable error: ' + action)
       );
-    } else {
-      return super.runTransaction(action, mode, transactionOperation);
     }
   }
 }

packages/firestore/test/unit/specs/spec_test_runner.ts

@@ -341,8 +341,8 @@ abstract class TestRunner {
     } else if ('changeUser' in step) {
       return this.doChangeUser(step.changeUser!);
     } else if ('failDatabase' in step) {
-      return step.failDatabase!
-        ? this.doFailDatabase()
+      return step.failDatabase
+        ? this.doFailDatabase(step.failDatabase!)
         : this.doRecoverDatabase();
     } else {
       return fail('Unknown step: ' + JSON.stringify(step));
@@ -371,7 +371,7 @@ abstract class TestRunner {
     await this.queue.enqueue(() => this.eventManager.listen(queryListener));
 
     if (targetFailed) {
-      expect(this.persistence.injectFailures).to.be.true;
+      expect(this.persistence.injectFailures?.['Allocate target']).to.be.true;
     } else {
       // Skip the backoff that may have been triggered by a previous call to
       // `watchStreamCloses()`.
@@ -446,7 +446,7 @@ abstract class TestRunner {
       () => this.rejectedDocs.push(...documentKeys)
     );
 
-    if (!this.persistence.injectFailures) {
+    if (this.persistence.injectFailures?.['Locally write mutations'] !== true) {
       this.sharedWrites.push(mutations);
     }
 
@@ -717,12 +717,14 @@ abstract class TestRunner {
     );
   }
 
-  private async doFailDatabase(): Promise<void> {
-    this.persistence.injectFailures = true;
+  private async doFailDatabase(
+    failActions: SpecDatabaseFailures
+  ): Promise<void> {
+    this.persistence.injectFailures = failActions;
   }
 
   private async doRecoverDatabase(): Promise<void> {
-    this.persistence.injectFailures = false;
+    this.persistence.injectFailures = undefined;
   }
 
   private validateExpectedSnapshotEvents(
@@ -1181,6 +1183,42 @@ export interface SpecConfig {
   maxConcurrentLimboResolutions?: number;
 }
 
+/**
+ * The cumulative list of actions run against Persistence. This is used by the
+ * Spec tests to fail specific types of actions.
+ */
+export type PersistenceAction =
+  | 'Get next mutation batch'
+  | 'read document'
+  | 'Allocate target'
+  | 'Release target'
+  | 'Execute query'
+  | 'Handle user change'
+  | 'Locally write mutations'
+  | 'Acknowledge batch'
+  | 'Reject batch'
+  | 'Get highest unacknowledged batch id'
+  | 'Get last stream token'
+  | 'Set last stream token'
+  | 'Get last remote snapshot version'
+  | 'Set last remote snapshot version'
+  | 'Apply remote event'
+  | 'notifyLocalViewChanges'
+  | 'Remote document keys'
+  | 'Collect garbage'
+  | 'maybeGarbageCollectMultiClientState'
+  | 'Lookup mutation documents'
+  | 'Get target data'
+  | 'Get new document changes'
+  | 'Synchronize last document change read time';
+
+/** Specifies failure or success for a list of database actions. */
+export type SpecDatabaseFailures = Partial<
+  {
+    readonly [key in PersistenceAction]: boolean;
+  }
+>;
+
 /**
  * Union type for each step. The step consists of exactly one `field`
  * set and optionally expected events in the `expect` field.
@@ -1225,8 +1263,8 @@ export interface SpecStep {
   /** Fail a write */
   failWrite?: SpecWriteFailure;
 
-  /** Fail all database transactions. */
-  failDatabase?: boolean;
+  /** Fails the listed database actions. */
+  failDatabase?: false | SpecDatabaseFailures;
 
   /**
    * Run a queued timer task (without waiting for the delay to expire). See