IndexeddbIndexManager.updateEntries() (#5999)

schmidt-sebastian · web-flow · commit 05a6aff7ece5 · 2022-02-15T10:02:23.000-07:00
diff --git a/packages/firestore/src/index/index_byte_encoder.ts b/packages/firestore/src/index/index_byte_encoder.ts
@@ -14,6 +14,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+import { IndexKind } from '../model/field_index';
 import { ByteString } from '../util/byte_string';
 
 import { DirectionalIndexByteEncoder } from './directional_index_byte_encoder';
@@ -69,8 +70,8 @@ export class IndexByteEncoder {
     this.orderedCode.seed(encodedBytes);
   }
 
-  forKind(kind: 'asc' | 'desc'): DirectionalIndexByteEncoder {
-    return kind === 'asc' ? this.ascending : this.descending;
+  forKind(kind: IndexKind): DirectionalIndexByteEncoder {
+    return kind === IndexKind.ASCENDING ? this.ascending : this.descending;
   }
 
   encodedBytes(): Uint8Array {
diff --git a/packages/firestore/src/index/index_entry.ts b/packages/firestore/src/index/index_entry.ts
@@ -0,0 +1,60 @@
+/**
+ * @license
+ * Copyright 2022 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 { DocumentKey } from '../model/document_key';
+
+/** Represents an index entry saved by the SDK in persisted storage. */
+export class IndexEntry {
+  constructor(
+    readonly indexId: number,
+    readonly documentKey: DocumentKey,
+    readonly arrayValue: Uint8Array,
+    readonly directionalValue: Uint8Array
+  ) {}
+}
+
+export function indexEntryComparator(
+  left: IndexEntry,
+  right: IndexEntry
+): number {
+  let cmp = left.indexId - right.indexId;
+  if (cmp !== 0) {
+    return cmp;
+  }
+
+  cmp = DocumentKey.comparator(left.documentKey, right.documentKey);
+  if (cmp !== 0) {
+    return cmp;
+  }
+
+  cmp = compareByteArrays(left.arrayValue, right.arrayValue);
+  if (cmp !== 0) {
+    return cmp;
+  }
+
+  return compareByteArrays(left.directionalValue, right.directionalValue);
+}
+
+function compareByteArrays(left: Uint8Array, right: Uint8Array): number {
+  for (let i = 0; i < left.length && i < right.length; ++i) {
+    const compare = left[i] - right[i];
+    if (compare !== 0) {
+      return compare;
+    }
+  }
+  return left.length - right.length;
+}
diff --git a/packages/firestore/src/local/indexeddb_index_manager.ts b/packages/firestore/src/local/indexeddb_index_manager.ts
@@ -17,15 +17,30 @@
 
 import { User } from '../auth/user';
 import { Target } from '../core/target';
+import { FirestoreIndexValueWriter } from '../index/firestore_index_value_writer';
+import { IndexByteEncoder } from '../index/index_byte_encoder';
+import { IndexEntry, indexEntryComparator } from '../index/index_entry';
 import {
   documentKeySet,
   DocumentKeySet,
   DocumentMap
 } from '../model/collections';
-import { FieldIndex, IndexOffset } from '../model/field_index';
+import { Document } from '../model/document';
+import { DocumentKey } from '../model/document_key';
+import {
+  FieldIndex,
+  fieldIndexGetArraySegment,
+  fieldIndexGetDirectionalSegments,
+  IndexKind,
+  IndexOffset
+} from '../model/field_index';
 import { ResourcePath } from '../model/path';
+import { isArray } from '../model/values';
+import { Value as ProtoValue } from '../protos/firestore_proto_api';
 import { debugAssert } from '../util/assert';
+import { logDebug } from '../util/log';
 import { immediateSuccessor } from '../util/misc';
+import { diffSortedSets, SortedSet } from '../util/sorted_set';
 
 import {
   decodeResourcePath,
@@ -53,6 +68,8 @@ import { PersistencePromise } from './persistence_promise';
 import { PersistenceTransaction } from './persistence_transaction';
 import { SimpleDbStore } from './simple_db';
 
+const LOG_TAG = 'IndexedDbIndexManager';
+
 /**
  * A persisted implementation of IndexManager.
  *
@@ -194,6 +211,40 @@ export class IndexedDbIndexManager implements IndexManager {
     return PersistencePromise.resolve<FieldIndex | null>(null);
   }
 
+  /**
+   * Returns the byte encoded form of the directional values in the field index.
+   * Returns `null` if the document does not have all fields specified in the
+   * index.
+   */
+  private encodeDirectionalElements(
+    fieldIndex: FieldIndex,
+    document: Document
+  ): Uint8Array | null {
+    const encoder = new IndexByteEncoder();
+    for (const segment of fieldIndexGetDirectionalSegments(fieldIndex)) {
+      const field = document.data.field(segment.fieldPath);
+      if (field == null) {
+        return null;
+      }
+      const directionalEncoder = encoder.forKind(segment.kind);
+      FirestoreIndexValueWriter.INSTANCE.writeIndexValue(
+        field,
+        directionalEncoder
+      );
+    }
+    return encoder.encodedBytes();
+  }
+
+  /** Encodes a single value to the ascending index format. */
+  private encodeSingleElement(value: ProtoValue): Uint8Array {
+    const encoder = new IndexByteEncoder();
+    FirestoreIndexValueWriter.INSTANCE.writeIndexValue(
+      value,
+      encoder.forKind(IndexKind.ASCENDING)
+    );
+    return encoder.encodedBytes();
+  }
+
   getFieldIndexes(
     transaction: PersistenceTransaction,
     collectionGroup?: string
@@ -269,8 +320,181 @@ export class IndexedDbIndexManager implements IndexManager {
     transaction: PersistenceTransaction,
     documents: DocumentMap
   ): PersistencePromise<void> {
-    // TODO(indexing): Implement
-    return PersistencePromise.resolve();
+    // Porting Note: `getFieldIndexes()` on Web does not cache index lookups as
+    // it could be used across different IndexedDB transactions. As any cached
+    // data might be invalidated by other multi-tab clients, we can only trust
+    // data within a single IndexedDB transaction. We therefore add a cache
+    // here.
+    const memoizedIndexes = new Map<string, FieldIndex[]>();
+    return PersistencePromise.forEach(documents, (key, doc) => {
+      const memoizedCollectionIndexes = memoizedIndexes.get(
+        key.collectionGroup
+      );
+      const fieldIndexes = memoizedCollectionIndexes
+        ? PersistencePromise.resolve(memoizedCollectionIndexes)
+        : this.getFieldIndexes(transaction, key.collectionGroup);
+
+      return fieldIndexes.next(fieldIndexes => {
+        memoizedIndexes.set(key.collectionGroup, fieldIndexes);
+        return PersistencePromise.forEach(
+          fieldIndexes,
+          (fieldIndex: FieldIndex) => {
+            return this.getExistingIndexEntries(
+              transaction,
+              key,
+              fieldIndex
+            ).next(existingEntries => {
+              const newEntries = this.computeIndexEntries(doc, fieldIndex);
+              if (!existingEntries.isEqual(newEntries)) {
+                return this.updateEntries(
+                  transaction,
+                  doc,
+                  existingEntries,
+                  newEntries
+                );
+              }
+              return PersistencePromise.resolve();
+            });
+          }
+        );
+      });
+    });
+  }
+
+  private addIndexEntry(
+    transaction: PersistenceTransaction,
+    document: Document,
+    indexEntry: IndexEntry
+  ): PersistencePromise<void> {
+    const indexEntries = indexEntriesStore(transaction);
+    return indexEntries.put(
+      new DbIndexEntry(
+        indexEntry.indexId,
+        this.uid,
+        indexEntry.arrayValue,
+        indexEntry.directionalValue,
+        encodeResourcePath(document.key.path)
+      )
+    );
+  }
+
+  private deleteIndexEntry(
+    transaction: PersistenceTransaction,
+    document: Document,
+    indexEntry: IndexEntry
+  ): PersistencePromise<void> {
+    const indexEntries = indexEntriesStore(transaction);
+    return indexEntries.delete([
+      indexEntry.indexId,
+      this.uid,
+      indexEntry.arrayValue,
+      indexEntry.directionalValue,
+      encodeResourcePath(document.key.path)
+    ]);
+  }
+
+  private getExistingIndexEntries(
+    transaction: PersistenceTransaction,
+    documentKey: DocumentKey,
+    fieldIndex: FieldIndex
+  ): PersistencePromise<SortedSet<IndexEntry>> {
+    const indexEntries = indexEntriesStore(transaction);
+    let results = new SortedSet<IndexEntry>(indexEntryComparator);
+    return indexEntries
+      .iterate(
+        {
+          index: DbIndexEntry.documentKeyIndex,
+          range: IDBKeyRange.only([
+            fieldIndex.indexId,
+            this.uid,
+            encodeResourcePath(documentKey.path)
+          ])
+        },
+        (_, entry) => {
+          results = results.add(
+            new IndexEntry(
+              fieldIndex.indexId,
+              documentKey,
+              entry.arrayValue,
+              entry.directionalValue
+            )
+          );
+        }
+      )
+      .next(() => results);
+  }
+
+  /** Creates the index entries for the given document. */
+  private computeIndexEntries(
+    document: Document,
+    fieldIndex: FieldIndex
+  ): SortedSet<IndexEntry> {
+    let results = new SortedSet<IndexEntry>(indexEntryComparator);
+
+    const directionalValue = this.encodeDirectionalElements(
+      fieldIndex,
+      document
+    );
+    if (directionalValue == null) {
+      return results;
+    }
+
+    const arraySegment = fieldIndexGetArraySegment(fieldIndex);
+    if (arraySegment != null) {
+      const value = document.data.field(arraySegment.fieldPath);
+      if (isArray(value)) {
+        for (const arrayValue of value.arrayValue!.values || []) {
+          results = results.add(
+            new IndexEntry(
+              fieldIndex.indexId,
+              document.key,
+              this.encodeSingleElement(arrayValue),
+              directionalValue
+            )
+          );
+        }
+      }
+    } else {
+      results = results.add(
+        new IndexEntry(
+          fieldIndex.indexId,
+          document.key,
+          new Uint8Array(),
+          directionalValue
+        )
+      );
+    }
+
+    return results;
+  }
+
+  /**
+   * Updates the index entries for the provided document by deleting entries
+   * that are no longer referenced in `newEntries` and adding all newly added
+   * entries.
+   */
+  private updateEntries(
+    transaction: PersistenceTransaction,
+    document: Document,
+    existingEntries: SortedSet<IndexEntry>,
+    newEntries: SortedSet<IndexEntry>
+  ): PersistencePromise<void> {
+    logDebug(LOG_TAG, "Updating index entries for document '%s'", document.key);
+
+    const promises: Array<PersistencePromise<void>> = [];
+    diffSortedSets(
+      existingEntries,
+      newEntries,
+      indexEntryComparator,
+      /* onAdd= */ entry => {
+        promises.push(this.addIndexEntry(transaction, document, entry));
+      },
+      /* onRemove= */ entry => {
+        promises.push(this.deleteIndexEntry(transaction, document, entry));
+      }
+    );
+
+    return PersistencePromise.waitFor(promises);
   }
 
   private getNextSequenceNumber(
diff --git a/packages/firestore/src/local/indexeddb_schema.ts b/packages/firestore/src/local/indexeddb_schema.ts
@@ -808,6 +808,10 @@ export class DbIndexEntry {
     'documentKey'
   ];
 
+  static documentKeyIndex = 'documentKeyIndex';
+
+  static documentKeyIndexPath = ['indexId', 'uid', 'documentKey'];
+
   constructor(
     /** The index id for this entry. */
     public indexId: number,
diff --git a/packages/firestore/src/local/indexeddb_schema_converter.ts b/packages/firestore/src/local/indexeddb_schema_converter.ts
@@ -545,9 +545,14 @@ function createFieldIndex(db: IDBDatabase): void {
     { unique: false }
   );
 
-  db.createObjectStore(DbIndexEntry.store, {
+  const indexEntryStore = db.createObjectStore(DbIndexEntry.store, {
     keyPath: DbIndexEntry.keyPath
   });
+  indexEntryStore.createIndex(
+    DbIndexEntry.documentKeyIndex,
+    DbIndexEntry.documentKeyIndexPath,
+    { unique: false }
+  );
 }
 
 function createDocumentOverlayStore(db: IDBDatabase): void {
diff --git a/packages/firestore/src/model/document_key.ts b/packages/firestore/src/model/document_key.ts
@@ -43,6 +43,14 @@ export class DocumentKey {
     return new DocumentKey(ResourcePath.emptyPath());
   }
 
+  get collectionGroup(): string {
+    debugAssert(
+      !this.path.isEmpty(),
+      'Cannot get collection group for empty key'
+    );
+    return this.path.popLast().lastSegment();
+  }
+
   /** Returns true if the document is in the specified collectionId. */
   hasCollectionId(collectionId: string): boolean {
     return (
diff --git a/packages/firestore/src/util/sorted_set.ts b/packages/firestore/src/util/sorted_set.ts
diff --git a/packages/firestore/test/unit/local/index_manager.test.ts b/packages/firestore/test/unit/local/index_manager.test.ts
diff --git a/packages/firestore/test/unit/util/sorted_set.test.ts b/packages/firestore/test/unit/util/sorted_set.test.ts
