Implement SessionToken Persistence (#8391)

tom-andersen · web-flow · commit 6bb2e8931182 · 2024-07-30T11:32:22.000-04:00
* Implement SessionToken Persistence

* Add migration tests.

* Fix migration test
diff --git a/packages/firestore/src/local/globals_cache.ts b/packages/firestore/src/local/globals_cache.ts
@@ -0,0 +1,48 @@
+/**
+ * @license
+ * Copyright 2024 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import { ByteString } from '../util/byte_string';
+
+import { PersistencePromise } from './persistence_promise';
+import { PersistenceTransaction } from './persistence_transaction';
+
+/**
+ * General purpose cache for global values.
+ *
+ * Global state that cuts across components should be saved here. Following are contained herein:
+ *
+ * `sessionToken` tracks server interaction across Listen and Write streams. This facilitates cache
+ * synchronization and invalidation.
+ */
+export interface GlobalsCache {
+  /**
+   * Gets session token.
+   */
+  getSessionToken(
+    transaction: PersistenceTransaction
+  ): PersistencePromise<ByteString>;
+
+  /**
+   * Sets session token.
+   *
+   * @param sessionToken - The new session token.
+   */
+  setSessionToken(
+    transaction: PersistenceTransaction,
+    sessionToken: ByteString
+  ): PersistencePromise<void>;
+}
diff --git a/packages/firestore/src/local/indexeddb_globals_cache.ts b/packages/firestore/src/local/indexeddb_globals_cache.ts
@@ -0,0 +1,54 @@
+/**
+ * @license
+ * Copyright 2024 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+import { ByteString } from '../util/byte_string';
+
+import { GlobalsCache } from './globals_cache';
+import { DbGlobals } from './indexeddb_schema';
+import { DbGlobalsStore, DbGlobalsKey } from './indexeddb_sentinels';
+import { getStore } from './indexeddb_transaction';
+import { PersistencePromise } from './persistence_promise';
+import { PersistenceTransaction } from './persistence_transaction';
+import { SimpleDbStore } from './simple_db';
+
+export class IndexedDbGlobalsCache implements GlobalsCache {
+  private globalsStore(
+    txn: PersistenceTransaction
+  ): SimpleDbStore<DbGlobalsKey, DbGlobals> {
+    return getStore<DbGlobalsKey, DbGlobals>(txn, DbGlobalsStore);
+  }
+
+  getSessionToken(txn: PersistenceTransaction): PersistencePromise<ByteString> {
+    const globals = this.globalsStore(txn);
+    return globals.get('sessionToken').next(global => {
+      const value = global?.value;
+      return value
+        ? ByteString.fromUint8Array(value)
+        : ByteString.EMPTY_BYTE_STRING;
+    });
+  }
+
+  setSessionToken(
+    txn: PersistenceTransaction,
+    sessionToken: ByteString
+  ): PersistencePromise<void> {
+    const globals = this.globalsStore(txn);
+    return globals.put({
+      name: 'sessionToken',
+      value: sessionToken.toUint8Array()
+    });
+  }
+}
diff --git a/packages/firestore/src/local/indexeddb_persistence.ts b/packages/firestore/src/local/indexeddb_persistence.ts
@@ -29,9 +29,11 @@ import { DocumentLike, WindowLike } from '../util/types';
 
 import { BundleCache } from './bundle_cache';
 import { DocumentOverlayCache } from './document_overlay_cache';
+import { GlobalsCache } from './globals_cache';
 import { IndexManager } from './index_manager';
 import { IndexedDbBundleCache } from './indexeddb_bundle_cache';
 import { IndexedDbDocumentOverlayCache } from './indexeddb_document_overlay_cache';
+import { IndexedDbGlobalsCache } from './indexeddb_globals_cache';
 import { IndexedDbIndexManager } from './indexeddb_index_manager';
 import { IndexedDbLruDelegateImpl } from './indexeddb_lru_delegate_impl';
 import { IndexedDbMutationQueue } from './indexeddb_mutation_queue';
@@ -188,6 +190,7 @@ export class IndexedDbPersistence implements Persistence {
   /** A listener to notify on primary state changes. */
   private primaryStateListener: PrimaryStateListener = _ => Promise.resolve();
 
+  private readonly globalsCache: IndexedDbGlobalsCache;
   private readonly targetCache: IndexedDbTargetCache;
   private readonly remoteDocumentCache: IndexedDbRemoteDocumentCache;
   private readonly bundleCache: IndexedDbBundleCache;
@@ -232,6 +235,7 @@ export class IndexedDbPersistence implements Persistence {
       this.schemaVersion,
       new SchemaConverter(this.serializer)
     );
+    this.globalsCache = new IndexedDbGlobalsCache();
     this.targetCache = new IndexedDbTargetCache(
       this.referenceDelegate,
       this.serializer
@@ -708,6 +712,14 @@ export class IndexedDbPersistence implements Persistence {
     return this._started;
   }
 
+  getGlobalsCache(): GlobalsCache {
+    debugAssert(
+      this.started,
+      'Cannot initialize GlobalsCache before persistence is started.'
+    );
+    return this.globalsCache;
+  }
+
   getMutationQueue(
     user: User,
     indexManager: IndexManager
diff --git a/packages/firestore/src/local/indexeddb_schema.ts b/packages/firestore/src/local/indexeddb_schema.ts
@@ -54,7 +54,7 @@ import { DbTimestampKey } from './indexeddb_sentinels';
  * 16. Parse timestamp strings before creating index entries.
  */
 
-export const SCHEMA_VERSION = 16;
+export const SCHEMA_VERSION = 17;
 
 /**
  * Wrapper class to store timestamps (seconds and nanos) in IndexedDb objects.
@@ -536,3 +536,13 @@ export interface DbDocumentOverlay {
   /** The overlay mutation. */
   overlayMutation: ProtoWrite;
 }
+
+/**
+ * An object containing global name/value pair.
+ */
+export interface DbGlobals {
+  /** Name is a globally unique identifier for a value. */
+  name: string;
+  /** Value is a general purpose storage for global data. */
+  value: Uint8Array;
+}
diff --git a/packages/firestore/src/local/indexeddb_schema_converter.ts b/packages/firestore/src/local/indexeddb_schema_converter.ts
@@ -68,6 +68,8 @@ import {
   DbDocumentOverlayCollectionPathOverlayIndexPath,
   DbDocumentOverlayKeyPath,
   DbDocumentOverlayStore,
+  DbGlobalsKeyPath,
+  DbGlobalsStore,
   DbIndexConfigurationCollectionGroupIndex,
   DbIndexConfigurationCollectionGroupIndexPath,
   DbIndexConfigurationKeyPath,
@@ -269,6 +271,12 @@ export class SchemaConverter implements SimpleDbSchemaConverter {
         });
     }
 
+    if (fromVersion < 17 && toVersion >= 17) {
+      p = p.next(() => {
+        createGlobalsStore(db);
+      });
+    }
+
     return p;
   }
 
@@ -748,6 +756,12 @@ function createDocumentOverlayStore(db: IDBDatabase): void {
   );
 }
 
+function createGlobalsStore(db: IDBDatabase): void {
+  db.createObjectStore(DbGlobalsStore, {
+    keyPath: DbGlobalsKeyPath
+  });
+}
+
 function extractKey(remoteDoc: DbRemoteDocumentLegacy): DocumentKey {
   if (remoteDoc.document) {
     return new DocumentKey(
diff --git a/packages/firestore/src/local/indexeddb_sentinels.ts b/packages/firestore/src/local/indexeddb_sentinels.ts
@@ -220,6 +220,7 @@ export const DbTargetDocumentDocumentTargetsKeyPath = ['path', 'targetId'];
  * The type to represent the single allowed key for the DbTargetGlobal store.
  */
 export type DbTargetGlobalKey = typeof DbTargetGlobalKey;
+
 /**
  * The key string used for the single object that exists in the
  * DbTargetGlobal store.
@@ -371,6 +372,14 @@ export const DbDocumentOverlayCollectionGroupOverlayIndexPath = [
   'largestBatchId'
 ];
 
+/** Name of the IndexedDb object store. */
+export const DbGlobalsStore = 'globals';
+
+export const DbGlobalsKeyPath = 'name';
+
+/** Names of global values */
+export type DbGlobalsKey = 'sessionToken';
+
 // Visible for testing
 export const V1_STORES = [
   DbMutationQueueStore,
@@ -415,6 +424,7 @@ export const V15_STORES = [
   DbIndexEntryStore
 ];
 export const V16_STORES = V15_STORES;
+export const V17_STORES = [...V15_STORES, DbGlobalsStore];
 
 /**
  * The list of all default IndexedDB stores used throughout the SDK. This is
@@ -425,7 +435,9 @@ export const ALL_STORES = V12_STORES;
 
 /** Returns the object stores for the provided schema. */
 export function getObjectStores(schemaVersion: number): string[] {
-  if (schemaVersion === 16) {
+  if (schemaVersion === 17) {
+    return V17_STORES;
+  } else if (schemaVersion === 16) {
     return V16_STORES;
   } else if (schemaVersion === 15) {
     return V15_STORES;
diff --git a/packages/firestore/src/local/memory_globals_cache.ts b/packages/firestore/src/local/memory_globals_cache.ts
@@ -0,0 +1,39 @@
+/**
+ * @license
+ * Copyright 2024 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+import { ByteString } from '../util/byte_string';
+
+import { GlobalsCache } from './globals_cache';
+import { PersistencePromise } from './persistence_promise';
+import { PersistenceTransaction } from './persistence_transaction';
+
+export class MemoryGlobalsCache implements GlobalsCache {
+  private sessionToken: ByteString = ByteString.EMPTY_BYTE_STRING;
+
+  getSessionToken(
+    transaction: PersistenceTransaction
+  ): PersistencePromise<ByteString> {
+    return PersistencePromise.resolve(this.sessionToken);
+  }
+
+  setSessionToken(
+    transaction: PersistenceTransaction,
+    sessionToken: ByteString
+  ): PersistencePromise<void> {
+    this.sessionToken = sessionToken;
+    return PersistencePromise.resolve();
+  }
+}
diff --git a/packages/firestore/src/local/memory_persistence.ts b/packages/firestore/src/local/memory_persistence.ts
@@ -29,6 +29,7 @@ import { ObjectMap } from '../util/obj_map';
 
 import { DocumentOverlayCache } from './document_overlay_cache';
 import { encodeResourcePath } from './encoded_resource_path';
+import { GlobalsCache } from './globals_cache';
 import { IndexManager } from './index_manager';
 import { LocalSerializer } from './local_serializer';
 import {
@@ -40,6 +41,7 @@ import {
 import { newLruGarbageCollector } from './lru_garbage_collector_impl';
 import { MemoryBundleCache } from './memory_bundle_cache';
 import { MemoryDocumentOverlayCache } from './memory_document_overlay_cache';
+import { MemoryGlobalsCache } from './memory_globals_cache';
 import { MemoryIndexManager } from './memory_index_manager';
 import { MemoryMutationQueue } from './memory_mutation_queue';
 import {
@@ -71,6 +73,7 @@ export class MemoryPersistence implements Persistence {
    * persisting values.
    */
   private readonly indexManager: MemoryIndexManager;
+  private readonly globalsCache: MemoryGlobalsCache;
   private mutationQueues: { [user: string]: MemoryMutationQueue } = {};
   private overlays: { [user: string]: MemoryDocumentOverlayCache } = {};
   private readonly remoteDocumentCache: MemoryRemoteDocumentCache;
@@ -94,6 +97,7 @@ export class MemoryPersistence implements Persistence {
     serializer: JsonProtoSerializer
   ) {
     this._started = true;
+    this.globalsCache = new MemoryGlobalsCache();
     this.referenceDelegate = referenceDelegateFactory(this);
     this.targetCache = new MemoryTargetCache(this);
     const sizer = (doc: Document): number =>
@@ -150,6 +154,10 @@ export class MemoryPersistence implements Persistence {
     return queue;
   }
 
+  getGlobalsCache(): GlobalsCache {
+    return this.globalsCache;
+  }
+
   getTargetCache(): MemoryTargetCache {
     return this.targetCache;
   }
diff --git a/packages/firestore/src/local/persistence.ts b/packages/firestore/src/local/persistence.ts
@@ -21,6 +21,7 @@ import { DocumentKey } from '../model/document_key';
 
 import { BundleCache } from './bundle_cache';
 import { DocumentOverlayCache } from './document_overlay_cache';
+import { GlobalsCache } from './globals_cache';
 import { IndexManager } from './index_manager';
 import { MutationQueue } from './mutation_queue';
 import { PersistencePromise } from './persistence_promise';
@@ -167,6 +168,11 @@ export interface Persistence {
    */
   setNetworkEnabled(networkEnabled: boolean): void;
 
+  /**
+   * Returns GlobalCache representing a general purpose cache for global values.
+   */
+  getGlobalsCache(): GlobalsCache;
+
   /**
    * Returns a MutationQueue representing the persisted mutations for the
    * given user.
diff --git a/packages/firestore/test/unit/local/globals_cache.test.ts b/packages/firestore/test/unit/local/globals_cache.test.ts
diff --git a/packages/firestore/test/unit/local/indexeddb_persistence.test.ts b/packages/firestore/test/unit/local/indexeddb_persistence.test.ts
diff --git a/packages/firestore/test/unit/local/test_globals_cache.ts b/packages/firestore/test/unit/local/test_globals_cache.ts
