Merge

Author: schmidt-sebastian

packages/firestore/src/local/indexeddb_mutation_queue.ts

@@ -26,6 +26,7 @@ import { BATCHID_UNKNOWN, MutationBatch } from '../model/mutation_batch';
 import { ResourcePath } from '../model/path';
 import { debugAssert, fail, hardAssert } from '../util/assert';
 import { primitiveComparator } from '../util/misc';
+import { ByteString } from '../util/byte_string';
 import { SortedMap } from '../util/sorted_map';
 import { SortedSet } from '../util/sorted_set';
 import { decodeResourcePath } from './encoded_resource_path';
@@ -117,6 +118,40 @@ export class IndexedDbMutationQueue implements MutationQueue {
       .next(() => empty);
   }
 
+  acknowledgeBatch(
+    transaction: PersistenceTransaction,
+    batch: MutationBatch,
+    streamToken: ByteString
+  ): PersistencePromise<void> {
+    return this.getMutationQueueMetadata(transaction).next(metadata => {
+      // We can't store the resumeToken as a ByteString in IndexedDB, so we
+      // convert it to a Base64 string for storage.
+      metadata.lastStreamToken = streamToken.toBase64();
+
+      return mutationQueuesStore(transaction).put(metadata);
+    });
+  }
+
+  getLastStreamToken(
+    transaction: PersistenceTransaction
+  ): PersistencePromise<ByteString> {
+    return this.getMutationQueueMetadata(transaction).next<ByteString>(
+      metadata => ByteString.fromBase64String(metadata.lastStreamToken)
+    );
+  }
+
+  setLastStreamToken(
+    transaction: PersistenceTransaction,
+    streamToken: ByteString
+  ): PersistencePromise<void> {
+    return this.getMutationQueueMetadata(transaction).next(metadata => {
+      // We can't store the resumeToken as a ByteString in IndexedDB, so we
+      // convert it to a Base64 string for storage.
+      metadata.lastStreamToken = streamToken.toBase64();
+      return mutationQueuesStore(transaction).put(metadata);
+    });
+  }
+
   addMutationBatch(
     transaction: PersistenceTransaction,
     localWriteTime: Timestamp,

packages/firestore/src/local/indexeddb_schema.ts

@@ -423,8 +423,6 @@ export class DbMutationQueue {
      *
      * After sending this token, earlier tokens may not be used anymore so
      * only a single stream token is retained.
-     *
-     * NOTE: this is deprecated and no longer used by the code.
      */
     public lastStreamToken: string
   ) {}

packages/firestore/src/local/local_store.ts

@@ -60,6 +60,7 @@ import { RemoteDocumentCache } from './remote_document_cache';
 import { RemoteDocumentChangeBuffer } from './remote_document_change_buffer';
 import { ClientId } from './shared_client_state';
 import { TargetData, TargetPurpose } from './target_data';
+import { ByteString } from '../util/byte_string';
 import { IndexedDbPersistence } from './indexeddb_persistence';
 import { IndexedDbMutationQueue } from './indexeddb_mutation_queue';
 import { IndexedDbRemoteDocumentCache } from './indexeddb_remote_document_cache';
@@ -377,11 +378,11 @@ export class LocalStore {
         const documentBuffer = this.remoteDocuments.newChangeBuffer({
           trackRemovals: true // Make sure document removals show up in `getNewDocumentChanges()`
         });
-        return this.applyWriteToRemoteDocuments(
-          txn,
-          batchResult,
-          documentBuffer
-        )
+        return this.mutationQueue
+          .acknowledgeBatch(txn, batchResult.batch, batchResult.streamToken)
+          .next(() =>
+            this.applyWriteToRemoteDocuments(txn, batchResult, documentBuffer)
+          )
           .next(() => documentBuffer.apply(txn))
           .next(() => this.mutationQueue.performConsistencyCheck(txn))
           .next(() => this.localDocuments.getDocuments(txn, affected));
@@ -432,6 +433,32 @@ export class LocalStore {
     );
   }
 
+  /** Returns the last recorded stream token for the current user. */
+  getLastStreamToken(): Promise<ByteString> {
+    return this.persistence.runTransaction(
+      'Get last stream token',
+      'readonly',
+      txn => {
+        return this.mutationQueue.getLastStreamToken(txn);
+      }
+    );
+  }
+
+  /**
+   * Sets the stream token for the current user without acknowledging any
+   * mutation batch. This is usually only useful after a stream handshake or in
+   * response to an error that requires clearing the stream token.
+   */
+  setLastStreamToken(streamToken: ByteString): Promise<void> {
+    return this.persistence.runTransaction(
+      'Set last stream token',
+      'readwrite-primary',
+      txn => {
+        return this.mutationQueue.setLastStreamToken(txn, streamToken);
+      }
+    );
+  }
+
   /**
    * Returns the last consistent snapshot processed (used by the RemoteStore to
    * determine whether to buffer incoming snapshots from the backend).

packages/firestore/src/local/memory_mutation_queue.ts

@@ -23,6 +23,7 @@ import { Mutation } from '../model/mutation';
 import { MutationBatch, BATCHID_UNKNOWN } from '../model/mutation_batch';
 import { debugAssert, hardAssert } from '../util/assert';
 import { primitiveComparator } from '../util/misc';
+import { ByteString } from '../util/byte_string';
 import { SortedMap } from '../util/sorted_map';
 import { SortedSet } from '../util/sorted_set';
 
@@ -42,6 +43,12 @@ export class MemoryMutationQueue implements MutationQueue {
   /** Next value to use when assigning sequential IDs to each mutation batch. */
   private nextBatchId: BatchId = 1;
 
+  /** The last received stream token from the server, used to acknowledge which
+   * responses the client has processed. Stream tokens are opaque checkpoint
+   * markers whose only real value is their inclusion in the next request.
+   */
+  private lastStreamToken: ByteString = ByteString.EMPTY_BYTE_STRING;
+
   /** An ordered mapping between documents and the mutations batch IDs. */
   private batchesByDocumentKey = new SortedSet(DocReference.compareByKey);
 
@@ -54,6 +61,46 @@ export class MemoryMutationQueue implements MutationQueue {
     return PersistencePromise.resolve(this.mutationQueue.length === 0);
   }
 
+  acknowledgeBatch(
+    transaction: PersistenceTransaction,
+    batch: MutationBatch,
+    streamToken: ByteString
+  ): PersistencePromise<void> {
+    const batchId = batch.batchId;
+    const batchIndex = this.indexOfExistingBatchId(batchId, 'acknowledged');
+    hardAssert(
+      batchIndex === 0,
+      'Can only acknowledge the first batch in the mutation queue'
+    );
+
+    // Verify that the batch in the queue is the one to be acknowledged.
+    const check = this.mutationQueue[batchIndex];
+    debugAssert(
+      batchId === check.batchId,
+      'Queue ordering failure: expected batch ' +
+        batchId +
+        ', got batch ' +
+        check.batchId
+    );
+
+    this.lastStreamToken = streamToken;
+    return PersistencePromise.resolve();
+  }
+
+  getLastStreamToken(
+    transaction: PersistenceTransaction
+  ): PersistencePromise<ByteString> {
+    return PersistencePromise.resolve(this.lastStreamToken);
+  }
+
+  setLastStreamToken(
+    transaction: PersistenceTransaction,
+    streamToken: ByteString
+  ): PersistencePromise<void> {
+    this.lastStreamToken = streamToken;
+    return PersistencePromise.resolve();
+  }
+
   addMutationBatch(
     transaction: PersistenceTransaction,
     localWriteTime: Timestamp,

packages/firestore/src/local/mutation_queue.ts

@@ -21,6 +21,7 @@ import { BatchId } from '../core/types';
 import { DocumentKey } from '../model/document_key';
 import { Mutation } from '../model/mutation';
 import { MutationBatch } from '../model/mutation_batch';
+import { ByteString } from '../util/byte_string';
 import { SortedMap } from '../util/sorted_map';
 
 import { PersistenceTransaction } from './persistence';
@@ -31,6 +32,26 @@ export interface MutationQueue {
   /** Returns true if this queue contains no mutation batches. */
   checkEmpty(transaction: PersistenceTransaction): PersistencePromise<boolean>;
 
+  /**
+   * Acknowledges the given batch.
+   */
+  acknowledgeBatch(
+    transaction: PersistenceTransaction,
+    batch: MutationBatch,
+    streamToken: ByteString
+  ): PersistencePromise<void>;
+
+  /** Returns the current stream token for this mutation queue. */
+  getLastStreamToken(
+    transaction: PersistenceTransaction
+  ): PersistencePromise<ByteString>;
+
+  /** Sets the stream token for this mutation queue. */
+  setLastStreamToken(
+    transaction: PersistenceTransaction,
+    streamToken: ByteString
+  ): PersistencePromise<void>;
+
   /**
    * Creates a new mutation batch and adds it to this mutation queue.
    *

packages/firestore/src/model/mutation_batch.ts

@@ -20,6 +20,7 @@ import { SnapshotVersion } from '../core/snapshot_version';
 import { BatchId } from '../core/types';
 import { hardAssert, debugAssert } from '../util/assert';
 import { arrayEquals } from '../util/misc';
+import { ByteString } from '../util/byte_string';
 import {
   documentKeySet,
   DocumentKeySet,
@@ -188,6 +189,7 @@ export class MutationBatchResult {
     readonly batch: MutationBatch,
     readonly commitVersion: SnapshotVersion,
     readonly mutationResults: MutationResult[],
+    readonly streamToken: ByteString,
     /**
      * A pre-computed mapping from each mutated document to the resulting
      * version.
@@ -203,7 +205,8 @@ export class MutationBatchResult {
   static from(
     batch: MutationBatch,
     commitVersion: SnapshotVersion,
-    results: MutationResult[]
+    results: MutationResult[],
+    streamToken: ByteString
   ): MutationBatchResult {
     hardAssert(
       batch.mutations.length === results.length,
@@ -219,6 +222,12 @@ export class MutationBatchResult {
       versionMap = versionMap.insert(mutations[i].key, results[i].version);
     }
 
-    return new MutationBatchResult(batch, commitVersion, results, versionMap);
+    return new MutationBatchResult(
+      batch,
+      commitVersion,
+      results,
+      streamToken,
+      versionMap
+    );
   }
 }

packages/firestore/src/remote/persistent_stream.ts

@@ -669,7 +669,7 @@ export class PersistentWriteStream extends PersistentStream<
    * PersistentWriteStream manages propagating this value from responses to the
    * next request.
    */
-  private lastStreamToken: ByteString = ByteString.EMPTY_BYTE_STRING;
+  lastStreamToken: ByteString = ByteString.EMPTY_BYTE_STRING;
 
   /**
    * Tracks whether or not a handshake has been successfully exchanged and

packages/firestore/src/remote/remote_store.ts

@@ -18,7 +18,7 @@
 import { SnapshotVersion } from '../core/snapshot_version';
 import { Transaction } from '../core/transaction';
 import { OnlineState, TargetId } from '../core/types';
-import { LocalStore } from '../local/local_store';
+import { ignoreIfPrimaryLeaseLoss, LocalStore } from '../local/local_store';
 import { TargetData, TargetPurpose } from '../local/target_data';
 import { MutationResult } from '../model/mutation';
 import {
@@ -43,7 +43,7 @@ import {
   PersistentWriteStream
 } from './persistent_stream';
 import { RemoteSyncer } from './remote_syncer';
-import { isPermanentWriteError } from './rpc_error';
+import { isPermanentError, isPermanentWriteError } from './rpc_error';
 import {
   DocumentWatchChange,
   ExistenceFilterChange,
@@ -199,6 +199,8 @@ export class RemoteStore implements TargetMetadataProvider {
 
   private async enableNetworkInternal(): Promise<void> {
     if (this.canUseNetwork()) {
+      this.writeStream.lastStreamToken = await this.localStore.getLastStreamToken();
+
       if (this.shouldStartWatchStream()) {
         this.startWatchStream();
       } else {
@@ -656,11 +658,16 @@ export class RemoteStore implements TargetMetadataProvider {
   }
 
   private onWriteHandshakeComplete(): Promise<void> {
-    // Send the write pipeline now that the stream is established.
-    for (const batch of this.writePipeline) {
-      this.writeStream.writeMutations(batch.mutations);
-    }
-    return Promise.resolve();
+    // Record the stream token.
+    return this.localStore
+      .setLastStreamToken(this.writeStream.lastStreamToken)
+      .then(() => {
+        // Send the write pipeline now that the stream is established.
+        for (const batch of this.writePipeline) {
+          this.writeStream.writeMutations(batch.mutations);
+        }
+      })
+      .catch(ignoreIfPrimaryLeaseLoss);
   }
 
   private async onMutationResult(
@@ -674,7 +681,12 @@ export class RemoteStore implements TargetMetadataProvider {
       'Got result for empty write pipeline'
     );
     const batch = this.writePipeline.shift()!;
-    const success = MutationBatchResult.from(batch, commitVersion, results);
+    const success = MutationBatchResult.from(
+      batch,
+      commitVersion,
+      results,
+      this.writeStream.lastStreamToken
+    );
     try {
       await this.syncEngine.applySuccessfulWrite(success);
     } catch (e) {
@@ -698,19 +710,45 @@ export class RemoteStore implements TargetMetadataProvider {
         'Write stream was stopped gracefully while still needed.'
       );
     }
-
     // An error that occurs after the write handshake completes is an indication
     // that the write operation itself failed.
     if (error && this.writeStream.handshakeComplete) {
       // This error affects the actual write.
       await this.handleWriteError(error!);
+    } else if (error) {
+      // If there was an error before the handshake has finished, it's
+      // possible that the server is unable to process the stream token
+      // we're sending. (Perhaps it's too old?)
+      await this.handleHandshakeError(error!);
     }
 
     // The write stream might have been started by refilling the write
     // pipeline for failed writes
     if (this.shouldStartWriteStream()) {
       this.startWriteStream();
     }
+    // No pending writes, nothing to do
+  }
+
+  private async handleHandshakeError(error: FirestoreError): Promise<void> {
+    // Reset the token if it's a permanent error, signaling the write stream is
+    // no longer valid. Note that the handshake does not count as a write: see
+    // comments on isPermanentWriteError for details.
+    if (isPermanentError(error.code)) {
+      logDebug(
+        LOG_TAG,
+        'RemoteStore error before completed handshake; resetting stream token: ',
+        this.writeStream.lastStreamToken
+      );
+      this.writeStream.lastStreamToken = ByteString.EMPTY_BYTE_STRING;
+
+      return this.localStore
+        .setLastStreamToken(ByteString.EMPTY_BYTE_STRING)
+        .catch(ignoreIfPrimaryLeaseLoss);
+    } else {
+      // Some other error, don't reset stream token. Our stream logic will
+      // just retry with exponential backoff.
+    }
   }
 
   private async handleWriteError(error: FirestoreError): Promise<void> {