[Multi-Tab] Recover from unexpected primary lease loss

packages/firestore/src/core/view.ts

@@ -383,9 +383,34 @@ export class View {
     return changes;
   }
 
+  /**
+   * Update the in-memory state of the current view with the state read from
+   * persistence.
+   *
+   * We update the query view whenever a client's primary status changes:
+   * - When a client transitions from primary to secondary, it can miss
+   *   LocalStorage updates and its query views may temporarily not be
+   *   synchronized with the state on disk.
+   * - For secondary to primary transitions, the client needs to update the list
+   *   of `syncedDocuments` since secondary clients update their query views
+   *   based purely on synthesized RemoteEvents.
+   *
+   * @param localDocs - The documents that match the query according to the
+   * LocalStore.
+   * @param remoteKeys - The keys of the documents that match the query
+   * according to the backend.
+   *
+   * @return The ViewChange that resulted from this synchronization.
+   */
   // PORTING NOTE: Multi-tab only.
-  synchronizeWithRemoteKeys(remoteKeys: DocumentKeySet): void {
+  synchronizeWithPersistedState(
+    localDocs: MaybeDocumentMap,
+    remoteKeys: DocumentKeySet
+  ): ViewChange {
     this._syncedDocuments = remoteKeys;
+    this.limboDocuments = documentKeySet();
+    const docChanges = this.computeDocChanges(localDocs);
+    return this.applyChanges(docChanges, /*updateLimboDocuments=*/ true);
   }
 
   /**

packages/firestore/src/local/indexeddb_persistence.ts

@@ -478,6 +478,8 @@ export class IndexedDbPersistence implements Persistence {
               log.error(
                 `Failed to obtain primary lease for action '${action}'.`
               );
+              this.isPrimary = false;
+              this.queue.enqueue(() => this.primaryStateListener(false));
               throw new FirestoreError(
                 Code.FAILED_PRECONDITION,
                 PRIMARY_LEASE_LOST_ERROR_MSG
@@ -734,6 +736,12 @@ export class IndexedDbPersistence implements Persistence {
   }
 }
 
+export function isPrimaryLeaseLostError(err: FirestoreError): boolean {
+  return (
+    err.code === Code.FAILED_PRECONDITION &&
+    err.message === PRIMARY_LEASE_LOST_ERROR_MSG
+  );
+}
 /**
  * Helper to get a typed SimpleDbStore for the owner object store.
  */

packages/firestore/src/local/local_store.ts

@@ -687,7 +687,7 @@ export class LocalStore {
    * found - used for testing.
    */
   readDocument(key: DocumentKey): Promise<MaybeDocument | null> {
-    return this.persistence.runTransaction('read document', true, txn => {
+    return this.persistence.runTransaction('read document', false, txn => {
       return this.localDocuments.getDocument(txn, key);
     });
   }

packages/firestore/src/remote/remote_store.ts

@@ -51,6 +51,7 @@ import {
 import { OnlineStateTracker } from './online_state_tracker';
 import { AsyncQueue } from '../util/async_queue';
 import { DocumentKeySet } from '../model/collections';
+import { isPrimaryLeaseLostError } from '../local/indexeddb_persistence';
 
 const LOG_TAG = 'RemoteStore';
 
@@ -580,7 +581,24 @@ export class RemoteStore implements TargetMetadataProvider {
         for (const batch of this.writePipeline) {
           this.writeStream.writeMutations(batch.mutations);
         }
-      });
+      })
+      .catch(err => this.ignoreIfPrimaryLeaseLoss(err));
+  }
+
+  /**
+   * Verifies the error thrown by an LocalStore operation. If a LocalStore
+   * operation fails because the primary lease has been taken by another client,
+   * we ignore the error. All other errors are re-thrown.
+   *
+   * @param err An error returned by a LocalStore operation.
+   * @return A Promise that resolves after we recovered, or the original error.
+   */
+  private ignoreIfPrimaryLeaseLoss(err: FirestoreError): void {
+    if (isPrimaryLeaseLostError(err)) {
+      log.debug(LOG_TAG, 'Unexpectedly lost primary lease');
+    } else {
+      throw err;
+    }
   }
 
   private onMutationResult(
@@ -650,7 +668,9 @@ export class RemoteStore implements TargetMetadataProvider {
       );
       this.writeStream.lastStreamToken = emptyByteString();
 
-      return this.localStore.setLastStreamToken(emptyByteString());
+      return this.localStore
+        .setLastStreamToken(emptyByteString())
+        .catch(err => this.ignoreIfPrimaryLeaseLoss(err));
     } else {
       // Some other error, don't reset stream token. Our stream logic will
       // just retry with exponential backoff.

packages/firestore/src/util/obj.ts

@@ -39,6 +39,13 @@ export function size<V>(obj: Dict<V>): number {
   return count;
 }
 
+/** Extracts the numeric indices from a dictionary. */
+export function indices<V>(obj: { [numberKey: number]: V }): number[] {
+  return Object.keys(obj).map(key => {
+    return Number(key);
+  });
+}
+
 /** Returns the given value if it's defined or the defaultValue otherwise. */
 export function defaulted<V>(value: V | undefined, defaultValue: V): V {
   return value !== undefined ? value : defaultValue;

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

@@ -393,7 +393,52 @@ describeSpec('Limbo Documents:', [], () => {
     }
   );
 
-  // TODO(multitab): We need a test case that verifies that a primary client
-  // that loses its primary lease while documents are in limbo correctly handles
-  // these documents even when it picks up its lease again.
+  specTest(
+    'Limbo documents survive primary state transitions',
+    ['multi-client'],
+    () => {
+      const query = Query.atPath(path('collection'));
+      const docA = doc('collection/a', 1000, { key: 'a' });
+      const docB = doc('collection/b', 1001, { key: 'b' });
+      const docC = doc('collection/c', 1002, { key: 'c' });
+      const deletedDocB = deletedDoc('collection/b', 1006);
+      const deletedDocC = deletedDoc('collection/c', 1008);
+
+      return (
+        client(0, false)
+          .expectPrimaryState(true)
+          .userListens(query)
+          .watchAcksFull(query, 1002, docA, docB, docC)
+          .expectEvents(query, { added: [docA, docB, docC] })
+          .watchRemovesDoc(docB.key, query)
+          .watchRemovesDoc(docC.key, query)
+          .watchSnapshots(1003)
+          .expectEvents(query, { fromCache: true })
+          .expectLimboDocs(docB.key, docC.key)
+          .client(1)
+          .stealPrimaryLease()
+          .client(0)
+          .runTimer(TimerId.ClientMetadataRefresh)
+          .expectPrimaryState(false)
+          .expectLimboDocs()
+          .client(1)
+          // TODO(37254270): This should be 'resume-token-1003' from the last
+          // global snapshot.
+          .expectListen(query, 'resume-token-1002')
+          .watchAcksFull(query, 1005)
+          .expectLimboDocs(docB.key, docC.key)
+          .ackLimbo(1006, deletedDocB)
+          .expectLimboDocs(docC.key)
+          .client(0)
+          .expectEvents(query, { removed: [docB], fromCache: true })
+          .stealPrimaryLease()
+          .expectListen(query, 'resume-token-1005')
+          .watchAcksFull(query, 1007)
+          .expectLimboDocs(docC.key)
+          .ackLimbo(1007, deletedDocC)
+          .expectLimboDocs()
+          .expectEvents(query, { removed: [docC] })
+      );
+    }
+  );
 });

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

@@ -939,6 +939,66 @@ describeSpec('Listens:', [], () => {
       .expectEvents(query, { added: [docB] });
   });
 
+  specTest('Query recovers after primary takeover', ['multi-client'], () => {
+    const query = Query.atPath(path('collection'));
+    const docA = doc('collection/a', 1000, { key: 'a' });
+    const docB = doc('collection/b', 2000, { key: 'b' });
+    const docC = doc('collection/c', 3000, { key: 'c' });
+
+    return (
+      client(0)
+        .expectPrimaryState(true)
+        .userListens(query)
+        .watchAcksFull(query, 1000, docA)
+        .expectEvents(query, { added: [docA] })
+        .client(1)
+        .userListens(query)
+        .expectEvents(query, { added: [docA] })
+        .stealPrimaryLease()
+        .expectListen(query, 'resume-token-1000')
+        .watchAcksFull(query, 2000, docB)
+        .expectEvents(query, { added: [docB] })
+        .client(0)
+        // Client 0 ignores all events until it transitions to secondary
+        .client(1)
+        .watchSends({ affects: [query] }, docC)
+        .watchSnapshots(3000)
+        .expectEvents(query, { added: [docC] })
+        .client(0)
+        .runTimer(TimerId.ClientMetadataRefresh)
+        .expectPrimaryState(false)
+        .expectEvents(query, { added: [docB, docC] })
+    );
+  });
+
+  specTest(
+    'Unresponsive primary ignores watch update',
+    ['multi-client'],
+    () => {
+      const query = Query.atPath(path('collection'));
+      const docA = doc('collection/a', 1000, { key: 'a' });
+
+      return (
+        client(0)
+          .expectPrimaryState(true)
+          .client(1)
+          .userListens(query)
+          .client(0)
+          .expectListen(query)
+          .client(1)
+          .stealPrimaryLease()
+          .client(0)
+          // Send a watch update to client 0, who is longer primary (but doesn't
+          // know it yet). The watch update gets ignored.
+          .watchAcksFull(query, 1000, docA)
+          .client(1)
+          .expectListen(query)
+          .watchAcksFull(query, 1000, docA)
+          .expectEvents(query, { added: [docA] })
+      );
+    }
+  );
+
   specTest(
     'Listen is established in newly started primary',
     ['multi-client'],

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

@@ -24,7 +24,11 @@ import {
 } from '../../../src/model/document';
 import { DocumentKey } from '../../../src/model/document_key';
 import { JsonObject } from '../../../src/model/field_value';
-import { mapRpcCodeFromCode } from '../../../src/remote/rpc_error';
+import {
+  isPermanentError,
+  mapCodeFromRpcCode,
+  mapRpcCodeFromCode
+} from '../../../src/remote/rpc_error';
 import { assert } from '../../../src/util/assert';
 import { fail } from '../../../src/util/assert';
 import { Code } from '../../../src/util/error';
@@ -76,12 +80,20 @@ export class ClientMemoryState {
     this.reset();
   }
 
+  /** Reset all internal memory state (as done during a client restart). */
   reset(): void {
     this.queryMapping = {};
     this.limboMapping = {};
     this.activeTargets = {};
     this.limboIdGenerator = TargetIdGenerator.forSyncEngine();
   }
+
+  /**
+   * Reset the internal limbo mapping (as done during a primary lease failover).
+   */
+  resetLimboMapping(): void {
+    this.limboMapping = {};
+  }
 }
 
 /**
@@ -457,12 +469,16 @@ export class SpecBuilder {
   writeAcks(
     doc: string,
     version: TestSnapshotVersion,
-    options?: { expectUserCallback: boolean }
+    options?: { expectUserCallback?: boolean; keepInQueue?: boolean }
   ): this {
     this.nextStep();
-    this.currentStep = { writeAck: { version } };
+    options = options || {};
 
-    if (!options || options.expectUserCallback) {
+    this.currentStep = {
+      writeAck: { version, keepInQueue: !!options.keepInQueue }
+    };
+
+    if (options.expectUserCallback) {
       return this.expectUserCallbacks({ acknowledged: [doc] });
     } else {
       return this;
@@ -476,13 +492,23 @@ export class SpecBuilder {
    */
   failWrite(
     doc: string,
-    err: RpcError,
-    options?: { expectUserCallback: boolean }
+    error: RpcError,
+    options?: { expectUserCallback?: boolean; keepInQueue?: boolean }
   ): this {
     this.nextStep();
-    this.currentStep = { failWrite: { error: err } };
+    options = options || {};
 
-    if (!options || options.expectUserCallback) {
+    // If this is a permanent error, the write is not expected to be sent
+    // again.
+    const isPermanentFailure = isPermanentError(mapCodeFromRpcCode(error.code));
+    const keepInQueue =
+      options.keepInQueue !== undefined
+        ? options.keepInQueue
+        : !isPermanentFailure;
+
+    this.currentStep = { failWrite: { error, keepInQueue } };
+
+    if (options.expectUserCallback) {
       return this.expectUserCallbacks({ rejected: [doc] });
     } else {
       return this;
@@ -904,6 +930,31 @@ export class MultiClientSpecBuilder extends SpecBuilder {
     return this;
   }
 
+  /**
+   * Take the primary lease, even if another client has already obtained the
+   * lease.
+   */
+  stealPrimaryLease(): this {
+    this.nextStep();
+    this.currentStep = {
+      applyClientState: {
+        primary: true
+      },
+      stateExpect: {
+        isPrimary: true
+      }
+    };
+
+    // HACK: SyncEngine resets its limbo mapping when it gains the primary
+    // lease. The SpecTests need to also clear their mapping, but when we parse
+    // the spec tests, we don't know when the primary lease transition happens.
+    // It is likely going to happen right after `stealPrimaryLease`, so we are
+    // clearing the limbo mapping here.
+    this.clientState.resetLimboMapping();
+
+    return this;
+  }
+
   protected nextStep(): void {
     if (this.currentStep !== null) {
       this.currentStep.clientIndex = this.activeClientIndex;