diff --git a/.changeset/fix-index-stale-bucket.md b/.changeset/fix-index-stale-bucket.md new file mode 100644 index 0000000000..a633378695 --- /dev/null +++ b/.changeset/fix-index-stale-bucket.md @@ -0,0 +1,9 @@ +--- +'@tanstack/db': patch +--- + +fix(db): prevent indexes from leaving a key in two buckets at once + +`BasicIndex` and `BTreeIndex` now track each key's currently-indexed value internally and use that on `remove`/`update`, instead of trusting the caller-supplied `oldItem`. If `oldItem` evaluated to a value that didn't match the bucket the key actually lived in, the previous implementation silently no-oped on the bucket map while the subsequent `add` placed the key in a new bucket, leaving the key in two buckets simultaneously. Calling `add` for a key that's already indexed also now drops the previous bucket first, so a missed `remove` can't double-bucket the key. + +This was the underlying cause of stale auto-index buckets after an optimistic→synced update on the indexed column (e.g. a kanban card snapping back to its old column). diff --git a/packages/db/src/indexes/basic-index.ts b/packages/db/src/indexes/basic-index.ts index b80f7fb439..55139ea431 100644 --- a/packages/db/src/indexes/basic-index.ts +++ b/packages/db/src/indexes/basic-index.ts @@ -54,6 +54,11 @@ export class BasicIndex< private sortedValues: Array = [] // Set of all indexed PKs private indexedKeys = new Set() + // Reverse map: key -> currently indexed value. Lets `remove`/`update` + // find the bucket a key actually lives in instead of trusting the + // caller-supplied old item, which can be stale across optimistic→synced + // lifecycle races. + private keyToValue = new Map() // Comparator function private compareFn: (a: any, b: any) => number = defaultComparator @@ -88,6 +93,13 @@ export class BasicIndex< const normalizedValue = normalizeValue(indexedValue) + // If the key is already indexed, drop it from its current bucket first so + // an `add` after a stale or missed `remove` can't leave the key in two + // buckets at once. + if (this.indexedKeys.has(key)) { + this.removeFromBucket(key, this.keyToValue.get(key)) + } + if (this.valueMap.has(normalizedValue)) { // Value already exists, just add the key to the set this.valueMap.get(normalizedValue)!.add(key) @@ -105,40 +117,49 @@ export class BasicIndex< } this.indexedKeys.add(key) + this.keyToValue.set(key, normalizedValue) this.updateTimestamp() } + private removeFromBucket(key: TKey, normalizedValue: any): void { + if (normalizedValue === undefined && !this.valueMap.has(undefined)) { + // Key wasn't tracked yet; nothing to remove from valueMap. + return + } + const keySet = this.valueMap.get(normalizedValue) + if (!keySet) return + keySet.delete(key) + if (keySet.size === 0) { + this.valueMap.delete(normalizedValue) + deleteInSortedArray(this.sortedValues, normalizedValue, this.compareFn) + } + } + /** * Removes a value from the index */ remove(key: TKey, item: any): void { - let indexedValue: any - try { - indexedValue = this.evaluateIndexExpression(item) - } catch (error) { - console.warn( - `Failed to evaluate index expression for key ${key} during removal:`, - error, - ) - this.indexedKeys.delete(key) - this.updateTimestamp() - return - } - - const normalizedValue = normalizeValue(indexedValue) - - if (this.valueMap.has(normalizedValue)) { - const keySet = this.valueMap.get(normalizedValue)! - keySet.delete(key) - - if (keySet.size === 0) { - // No more keys for this value, remove from map and sorted array - this.valueMap.delete(normalizedValue) - deleteInSortedArray(this.sortedValues, normalizedValue, this.compareFn) + // Use the tracked bucket for this key. We intentionally do NOT trust + // the caller-supplied `item`: across optimistic→synced lifecycle races + // the previousValue passed in can disagree with the value we actually + // indexed, which would silently leave the key in a stale bucket. + if (this.keyToValue.has(key)) { + this.removeFromBucket(key, this.keyToValue.get(key)) + } else { + // Fall back to evaluating the supplied item (covers the case where + // the key was never indexed but caller still asks to remove it). + try { + const indexedValue = this.evaluateIndexExpression(item) + this.removeFromBucket(key, normalizeValue(indexedValue)) + } catch (error) { + console.warn( + `Failed to evaluate index expression for key ${key} during removal:`, + error, + ) } } - this.indexedKeys.delete(key) + this.keyToValue.delete(key) this.updateTimestamp() } @@ -168,8 +189,10 @@ export class BasicIndex< { cause: error }, ) } - entriesArray.push({ key, value: normalizeValue(indexedValue) }) + const normalizedValue = normalizeValue(indexedValue) + entriesArray.push({ key, value: normalizedValue }) this.indexedKeys.add(key) + this.keyToValue.set(key, normalizedValue) } // Group by value @@ -194,6 +217,7 @@ export class BasicIndex< this.valueMap.clear() this.sortedValues = [] this.indexedKeys.clear() + this.keyToValue.clear() this.updateTimestamp() } diff --git a/packages/db/src/indexes/btree-index.ts b/packages/db/src/indexes/btree-index.ts index 17608950a7..2afe00097a 100644 --- a/packages/db/src/indexes/btree-index.ts +++ b/packages/db/src/indexes/btree-index.ts @@ -50,6 +50,11 @@ export class BTreeIndex< private orderedEntries: BTree // we don't associate values with the keys of the B+ tree (the keys are indexed values) private valueMap = new Map>() // instead we store a mapping of indexed values to a set of PKs private indexedKeys = new Set() + // Reverse map: key -> currently indexed value. Lets `remove`/`update` + // find the bucket a key actually lives in instead of trusting the + // caller-supplied old item, which can be stale across optimistic→synced + // lifecycle races. + private keyToValue = new Map() private compareFn: (a: any, b: any) => number = defaultComparator constructor( @@ -93,6 +98,13 @@ export class BTreeIndex< // Normalize the value for Map key usage const normalizedValue = normalizeForBTree(indexedValue) + // If the key is already indexed, drop it from its current bucket first so + // an `add` after a stale or missed `remove` can't leave the key in two + // buckets at once. + if (this.indexedKeys.has(key)) { + this.removeFromBucket(key, this.keyToValue.get(key)) + } + // Check if this value already exists if (this.valueMap.has(normalizedValue)) { // Add to existing set @@ -105,41 +117,45 @@ export class BTreeIndex< } this.indexedKeys.add(key) + this.keyToValue.set(key, normalizedValue) this.updateTimestamp() } + private removeFromBucket(key: TKey, normalizedValue: any): void { + const keySet = this.valueMap.get(normalizedValue) + if (!keySet) return + keySet.delete(key) + if (keySet.size === 0) { + this.valueMap.delete(normalizedValue) + this.orderedEntries.delete(normalizedValue) + } + } + /** * Removes a value from the index */ remove(key: TKey, item: any): void { - let indexedValue: any - try { - indexedValue = this.evaluateIndexExpression(item) - } catch (error) { - console.warn( - `Failed to evaluate index expression for key ${key} during removal:`, - error, - ) - return - } - - // Normalize the value for Map key usage - const normalizedValue = normalizeForBTree(indexedValue) - - if (this.valueMap.has(normalizedValue)) { - const keySet = this.valueMap.get(normalizedValue)! - keySet.delete(key) - - // If set is now empty, remove the entry entirely - if (keySet.size === 0) { - this.valueMap.delete(normalizedValue) - - // Remove from ordered entries - this.orderedEntries.delete(normalizedValue) + // Use the tracked bucket for this key. We intentionally do NOT trust + // the caller-supplied `item`: across optimistic→synced lifecycle races + // the previousValue passed in can disagree with the value we actually + // indexed, which would silently leave the key in a stale bucket. + if (this.keyToValue.has(key)) { + this.removeFromBucket(key, this.keyToValue.get(key)) + } else { + // Fall back to evaluating the supplied item (covers the case where + // the key was never indexed but caller still asks to remove it). + try { + const indexedValue = this.evaluateIndexExpression(item) + this.removeFromBucket(key, normalizeForBTree(indexedValue)) + } catch (error) { + console.warn( + `Failed to evaluate index expression for key ${key} during removal:`, + error, + ) } } - this.indexedKeys.delete(key) + this.keyToValue.delete(key) this.updateTimestamp() } @@ -169,6 +185,7 @@ export class BTreeIndex< this.orderedEntries.clear() this.valueMap.clear() this.indexedKeys.clear() + this.keyToValue.clear() this.updateTimestamp() } diff --git a/packages/db/tests/index-key-bucket-tracking.test.ts b/packages/db/tests/index-key-bucket-tracking.test.ts new file mode 100644 index 0000000000..b8f7d83cc4 --- /dev/null +++ b/packages/db/tests/index-key-bucket-tracking.test.ts @@ -0,0 +1,127 @@ +import { describe, expect, it } from 'vitest' +import { createCollection } from '../src/collection/index.js' +import { createTransaction } from '../src/transactions' +import { BasicIndex } from '../src/indexes/basic-index' +import { BTreeIndex } from '../src/indexes/btree-index' +import { eq } from '../src/query/builder/functions' +import { PropRef } from '../src/query/ir' +import type { BaseIndex, IndexConstructor } from '../src/indexes/base-index' +import type { ChangeMessageOrDeleteKeyMessage } from '../src/types' + +interface Project { + id: string + name: string + current_stage_id: string + updated_at: string +} + +const row = ( + id: string, + stage: string, + t: string = `2023-01-01T00:00:00Z`, +): Project => ({ + id, + name: `Project-${id}`, + current_stage_id: stage, + updated_at: t, +}) + +describe.each([ + [`BasicIndex`, BasicIndex], + [`BTreeIndex`, BTreeIndex], +] as Array<[string, IndexConstructor]>)( + `%s: a key lives in at most one bucket at a time`, + (_name, IndexType) => { + it(`update moves the key out of its actual bucket, regardless of the oldItem passed`, () => { + const index = new IndexType(1, new PropRef([`current_stage_id`])) + index.add(`P`, row(`P`, `A`)) + expect(index.equalityLookup(`A`)).toEqual(new Set([`P`])) + + // The oldItem disagrees with the value the index recorded for P. + index.update(`P`, row(`P`, `B`), row(`P`, `C`)) + + expect(index.equalityLookup(`A`)).toEqual(new Set()) + expect(index.equalityLookup(`B`)).toEqual(new Set()) + expect(index.equalityLookup(`C`)).toEqual(new Set([`P`])) + }) + + it(`add called twice without an intervening remove keeps the key in the latest bucket only`, () => { + const index = new IndexType(1, new PropRef([`current_stage_id`])) + index.add(`P`, row(`P`, `A`)) + index.add(`P`, row(`P`, `B`)) + + expect(index.equalityLookup(`A`)).toEqual(new Set()) + expect(index.equalityLookup(`B`)).toEqual(new Set([`P`])) + }) + + it(`an optimistic→synced update on the indexed column ends with the row in the new bucket`, async () => { + let resolveMutation: () => void + let begin!: () => void + let write!: (m: ChangeMessageOrDeleteKeyMessage) => void + let commit!: () => void + let markReady!: () => void + + const collection = createCollection({ + id: `projects`, + getKey: (item) => item.id, + autoIndex: `eager`, + defaultIndexType: IndexType, + startSync: true, + sync: { + sync: (ctx) => { + begin = ctx.begin + write = ctx.write + commit = ctx.commit + markReady = ctx.markReady + begin() + write({ type: `insert`, value: row(`P`, `A`) }) + commit() + markReady() + }, + }, + }) + await collection.stateWhenReady() + + collection.subscribeChanges(() => {}, { + whereExpression: eq(new PropRef([`current_stage_id`]), `A`), + includeInitialState: true, + }) + const idx = Array.from( + collection.indexes.values(), + )[0] as BaseIndex + + // Offline-style optimistic mutation (non-direct ambient transaction). + const tx = createTransaction({ + autoCommit: false, + mutationFn: async () => { + await new Promise((resolve) => { + resolveMutation = resolve + }) + }, + }) + tx.mutate(() => { + collection.update(`P`, (draft) => { + draft.current_stage_id = `B` + }) + }) + + const commitPromise = tx.commit() + + begin() + write({ + type: `update`, + value: row(`P`, `B`, `2023-01-02T00:00:00Z`), + }) + commit() + + resolveMutation!() + await commitPromise + await Promise.resolve() + await Promise.resolve() + + expect(collection.get(`P`)?.current_stage_id).toBe(`B`) + expect(idx.equalityLookup(`A`)).toEqual(new Set()) + expect(idx.equalityLookup(`B`)).toEqual(new Set([`P`])) + }) + }, +)