Batch removeItem/removeItems through WriteBuffer as tombstones

sosek108 · sosek108 · commit de20c96f2060 · 2026-05-18T13:45:37.000+02:00
diff --git a/lib/storage/NativeFlushWorker.ts b/lib/storage/NativeFlushWorker.ts
@@ -67,14 +67,17 @@ function createNativeFlushWorker(bufferStore: BufferStore) {
             return;
         }
 
-        // Separate SET and MERGE entries for the provider
+        // Separate SET, MERGE, and REMOVE entries for the provider
         const setPairs: StorageKeyValuePair[] = [];
         const mergePairs: StorageKeyValuePair[] = [];
+        const removeKeys: string[] = [];
 
         for (const [, entry] of entries) {
             const typedEntry = entry as BufferEntry;
             if (typedEntry.entryType === 'set') {
                 setPairs.push([typedEntry.key, typedEntry.value]);
+            } else if (typedEntry.entryType === 'remove') {
+                removeKeys.push(typedEntry.key);
             } else {
                 mergePairs.push([typedEntry.key, typedEntry.value, typedEntry.replaceNullPatches]);
             }
@@ -88,6 +91,9 @@ function createNativeFlushWorker(bufferStore: BufferStore) {
             if (mergePairs.length > 0) {
                 nativeSQLiteProvider.multiMerge(mergePairs);
             }
+            if (removeKeys.length > 0) {
+                nativeSQLiteProvider.removeItems(removeKeys);
+            }
         } catch (error) {
             console.error('[Onyx] NativeFlushWorker: flush error:', error);
         }
diff --git a/lib/storage/WriteBuffer.ts b/lib/storage/WriteBuffer.ts
@@ -27,7 +27,7 @@ import utils from '../utils';
 import type {StorageKeyValuePair} from './providers/types';
 import type BufferStore from './BufferStore/types';
 
-type EntryType = 'set' | 'merge';
+type EntryType = 'set' | 'merge' | 'remove';
 
 type BufferEntry = {
     key: OnyxKey;
@@ -36,10 +36,11 @@ type BufferEntry = {
     replaceNullPatches?: FastMergeReplaceNullPatch[];
 };
 
-/** Flush handlers for the two entry types. */
+/** Flush handlers for the three entry types. */
 type FlushHandlers = {
     multiSet: (pairs: StorageKeyValuePair[]) => Promise<void>;
     multiMerge: (pairs: StorageKeyValuePair[]) => Promise<void>;
+    multiRemove: (keys: OnyxKey[]) => Promise<void>;
 };
 
 /**
@@ -73,7 +74,9 @@ const FLUSH_TIMEOUT_MS = 200;
 /** Default web flush scheduler using requestIdleCallback with setTimeout fallback. */
 function defaultScheduleFlush(doFlush: () => void): number | null {
     if (typeof requestIdleCallback === 'function') {
-        return requestIdleCallback(doFlush, {timeout: FLUSH_TIMEOUT_MS}) as unknown as number;
+        return requestIdleCallback(doFlush, {
+            timeout: FLUSH_TIMEOUT_MS,
+        }) as unknown as number;
     }
     return setTimeout(doFlush, FLUSH_TIMEOUT_MS) as unknown as number;
 }
@@ -144,20 +147,38 @@ class WriteBuffer {
      * - No existing entry: create a MERGE entry with just the patch
      * - Existing SET entry: apply patch to the full value in-memory, keep as SET
      * - Existing MERGE entry: merge patches together, keep as MERGE
+     * - Existing REMOVE (tombstone): the underlying value will be null after flush,
+     *   so merging into null is just the patch itself -- replace tombstone with SET
      */
     merge(key: OnyxKey, patch: OnyxValue<OnyxKey>, replaceNullPatches?: FastMergeReplaceNullPatch[]): void {
         const existing = this.store.get(key);
 
         if (!existing) {
             // No pending write -- stage as a MERGE entry (just the patch)
-            this.store.set(key, {key, value: patch, entryType: 'merge', replaceNullPatches});
+            this.store.set(key, {
+                key,
+                value: patch,
+                entryType: 'merge',
+                replaceNullPatches,
+            });
+        } else if (existing.entryType === 'remove') {
+            // Pending REMOVE -- merging into null is just the patch itself, replace with SET
+            this.store.set(key, {
+                key,
+                value: patch,
+                entryType: 'set',
+            });
         } else if (existing.entryType === 'set') {
             // Pending SET -- apply patch to the full value, stay as SET
             const {result: merged} = utils.fastMerge(existing.value as Record<string, unknown>, patch as Record<string, unknown>, {
                 shouldRemoveNestedNulls: true,
                 objectRemovalMode: 'replace',
             });
-            this.store.set(key, {key, value: merged as OnyxValue<OnyxKey>, entryType: 'set'});
+            this.store.set(key, {
+                key,
+                value: merged as OnyxValue<OnyxKey>,
+                entryType: 'set',
+            });
         } else {
             // Pending MERGE -- merge patches together, stay as MERGE
             const {result: mergedPatch} = utils.fastMerge(existing.value as Record<string, unknown>, patch as Record<string, unknown>, {
@@ -176,20 +197,32 @@ class WriteBuffer {
     }
 
     /**
-     * Remove a key from the write buffer. This is used when a key is being
-     * removed from storage entirely (not just set to null).
+     * Stage a tombstone for a key (REMOVE entry). A pending SET or MERGE is
+     * discarded; the tombstone wins and is flushed via multiRemove. Scheduling
+     * a flush mirrors set/merge so deletes are batched off the worker hot path.
      */
     remove(key: OnyxKey): void {
-        this.store.delete(key);
+        this.store.set(key, {
+            key,
+            value: null as OnyxValue<OnyxKey>,
+            entryType: 'remove',
+        });
+        this.scheduleFlush();
     }
 
     /**
-     * Remove multiple keys from the write buffer.
+     * Stage tombstones for multiple keys (all REMOVE entries). A single flush
+     * is scheduled regardless of how many keys are removed.
      */
     removeMany(keys: OnyxKey[]): void {
         for (const key of keys) {
-            this.store.delete(key);
+            this.store.set(key, {
+                key,
+                value: null as OnyxValue<OnyxKey>,
+                entryType: 'remove',
+            });
         }
+        this.scheduleFlush();
     }
 
     /**
@@ -217,12 +250,21 @@ class WriteBuffer {
     }
 
     /**
-     * Check whether the key has any pending entry (SET or MERGE).
+     * Check whether the key has any pending entry (SET, MERGE, or REMOVE).
      */
     hasAny(key: OnyxKey): boolean {
         return this.store.has(key);
     }
 
+    /**
+     * Inspect the pending entry type for a key without touching it. Used by the
+     * storage facade to disambiguate reads: SET serves from buffer, REMOVE
+     * returns null immediately, MERGE forces a flush before reading the provider.
+     */
+    peekEntryType(key: OnyxKey): EntryType | undefined {
+        return this.store.get(key)?.entryType;
+    }
+
     /**
      * Returns the number of pending entries (both SET and MERGE).
      */
@@ -277,11 +319,16 @@ class WriteBuffer {
             const setSnapshot = new Map<OnyxKey, BufferEntry>();
             const mergePairs: StorageKeyValuePair[] = [];
             const mergeKeys: OnyxKey[] = [];
+            const removeKeys: OnyxKey[] = [];
+            const removeSnapshot = new Map<OnyxKey, BufferEntry>();
 
             for (const [key, entry] of this.store.entries()) {
                 if (entry.entryType === 'set') {
                     setPairs.push([entry.key, entry.value]);
                     setSnapshot.set(key, entry);
+                } else if (entry.entryType === 'remove') {
+                    removeKeys.push(entry.key);
+                    removeSnapshot.set(key, entry);
                 } else {
                     mergePairs.push([entry.key, entry.value, entry.replaceNullPatches]);
                     mergeKeys.push(key);
@@ -294,14 +341,17 @@ class WriteBuffer {
                 this.store.delete(key);
             }
 
-            // Flush both types concurrently
+            // Flush all three types concurrently
             const promises: Array<Promise<void>> = [];
             if (setPairs.length > 0) {
                 promises.push(this.handlers.multiSet(setPairs));
             }
             if (mergePairs.length > 0) {
                 promises.push(this.handlers.multiMerge(mergePairs));
             }
+            if (removeKeys.length > 0) {
+                promises.push(this.handlers.multiRemove(removeKeys));
+            }
 
             await Promise.all(promises);
 
@@ -312,6 +362,14 @@ class WriteBuffer {
                     this.store.delete(key);
                 }
             }
+
+            // Same reference-identity check for REMOVE tombstones: a set/merge
+            // arriving during flush replaces the entry and must not be erased.
+            for (const [key, flushedEntry] of removeSnapshot) {
+                if (this.store.get(key) === flushedEntry) {
+                    this.store.delete(key);
+                }
+            }
         } finally {
             this.isFlushing = false;
         }
diff --git a/lib/storage/index.ts b/lib/storage/index.ts
@@ -70,6 +70,7 @@ const writeBuffer = new WriteBuffer({
     handlers: {
         multiSet: (pairs) => provider.multiSet(pairs),
         multiMerge: (pairs) => provider.multiMerge(pairs),
+        multiRemove: (keys) => provider.removeItems(keys),
     },
     store: createBufferStore(),
 });
@@ -104,18 +105,23 @@ const storage: Storage = {
     /**
      * Get the value of a given key or return `null` if it's not available.
      *
-     * - If the key has a pending SET entry, it is returned directly from memory.
-     * - If the key has a pending MERGE entry (a patch delta, not a full value),
-     *   the write buffer is flushed first so the provider has the correct merged
-     *   value on disk, then the read proceeds normally.
-     * - Otherwise, the read goes straight to the provider.
+     * - SET entry: served directly from memory.
+     * - REMOVE entry (tombstone): returns null immediately. The on-disk value
+     *   will be gone once the buffer flushes; callers see the post-remove view.
+     * - MERGE entry: the patch is incomplete, so the buffer is flushed first,
+     *   then the merged value is read back from the provider.
+     * - No pending entry: read straight from the provider.
      */
     getItem: <TKey extends OnyxKey>(key: TKey) =>
         tryOrDegradePerformance<OnyxValue<TKey>>(() => {
-            if (writeBuffer.has(key)) {
+            const entryType = writeBuffer.peekEntryType(key);
+            if (entryType === 'set') {
                 return Promise.resolve(writeBuffer.get(key) as OnyxValue<TKey>);
             }
-            if (writeBuffer.hasAny(key)) {
+            if (entryType === 'remove') {
+                return Promise.resolve(null as unknown as OnyxValue<TKey>);
+            }
+            if (entryType === 'merge') {
                 return writeBuffer.flushNow().then(() => provider.getItem(key));
             }
             return provider.getItem(key);
@@ -131,23 +137,30 @@ const storage: Storage = {
     multiGet: (keys) =>
         tryOrDegradePerformance(() => {
             const bufferedKeys: string[] = [];
+            const tombstoneKeys: string[] = [];
             const cleanKeys: string[] = [];
             let hasPendingMerges = false;
 
             for (const key of keys) {
-                if (writeBuffer.has(key)) {
+                const entryType = writeBuffer.peekEntryType(key);
+                if (entryType === 'set') {
                     bufferedKeys.push(key);
+                } else if (entryType === 'remove') {
+                    tombstoneKeys.push(key);
                 } else {
                     cleanKeys.push(key);
-                    if (writeBuffer.hasAny(key)) {
+                    if (entryType === 'merge') {
                         hasPendingMerges = true;
                     }
                 }
             }
 
-            // If all keys have SET entries, skip the provider call entirely
+            const tombstoneResults = tombstoneKeys.map((key) => [key, null] as [string, unknown]);
+
+            // If no key needs the provider, skip the worker round-trip entirely
             if (cleanKeys.length === 0) {
-                return Promise.resolve(bufferedKeys.map((key) => [key, writeBuffer.get(key)]));
+                const bufferedResults = bufferedKeys.map((key) => [key, writeBuffer.get(key)] as [string, unknown]);
+                return Promise.resolve([...bufferedResults, ...tombstoneResults]);
             }
 
             // If any clean keys have pending MERGE entries, flush first
@@ -156,7 +169,7 @@ const storage: Storage = {
             return readyPromise.then(() =>
                 provider.multiGet(cleanKeys).then((providerResults) => {
                     const bufferedResults = bufferedKeys.map((key) => [key, writeBuffer.get(key)] as [string, unknown]);
-                    return [...providerResults, ...bufferedResults];
+                    return [...providerResults, ...bufferedResults, ...tombstoneResults];
                 }),
             );
         }),
@@ -211,23 +224,24 @@ const storage: Storage = {
         }),
 
     /**
-     * Removes given key and its value.
-     * Also removes the key from the write buffer if it has a pending write.
+     * Removes given key and its value. Staged as a REMOVE tombstone in the
+     * write buffer; the actual provider delete happens in the next coalesced
+     * flush via multiRemove, so deletes no longer block reads in the worker.
      */
     removeItem: (key) =>
         tryOrDegradePerformance(() => {
             writeBuffer.remove(key);
-            return provider.removeItem(key);
+            return Promise.resolve();
         }),
 
     /**
-     * Remove given keys and their values.
-     * Also removes the keys from the write buffer if they have pending writes.
+     * Remove given keys and their values. Staged as REMOVE tombstones in the
+     * write buffer; flushed together as a single multiRemove on next flush.
      */
     removeItems: (keys) =>
         tryOrDegradePerformance(() => {
             writeBuffer.removeMany(keys);
-            return provider.removeItems(keys);
+            return Promise.resolve();
         }),
 
     /**
