Keep values written while Onyx.clear is in progress

lib/Onyx.ts

@@ -428,7 +428,16 @@ function mergeCollection<TKey extends CollectionKeyBase>(collectionKey: TKey, co
  * @param keysToPreserve is a list of ONYXKEYS that should not be cleared with the rest of the data
  */
 function clear(keysToPreserve: OnyxKey[] = []): Promise<void> {
-    return PendingWrites.trackPendingGlobalWrite(
+    // Writes can land while storage clears (e.g. a new sign-in's session), so record them to avoid clobbering them.
+    // Matching on values instead would confuse them with reads refilling the cache from storage.
+    const keysWrittenDuringClear = new Set<OnyxKey>();
+    const removeWriteListener = PendingWrites.addWriteListener((keys) => {
+        for (const key of keys) {
+            keysWrittenDuringClear.add(key);
+        }
+    });
+
+    const clearPromise = PendingWrites.trackPendingGlobalWrite(
         OnyxUtils.afterInit(() => {
             const pendingMergeEntries = OnyxUtils.getPendingMergeEntries(keysToPreserve);
             const defaultKeyStates = OnyxUtils.getDefaultKeyStates();
@@ -449,6 +458,17 @@ function clear(keysToPreserve: OnyxKey[] = []): Promise<void> {
 
                     const allKeys = new Set([...cachedKeys, ...initialKeys]);
 
+                    // A write issued before the reset below is overwritten by it, so only a cache value that differs from
+                    // what clear set counts as a write that landed after the reset.
+                    const valuesSetByClear: KeyValueMapping = {};
+                    const wasWrittenDuringClear = (key: OnyxKey): boolean => {
+                        if (!keysWrittenDuringClear.has(key) || !(key in valuesSetByClear)) {
+                            return false;
+                        }
+                        const currentValue = cache.get(key);
+                        return currentValue !== undefined && currentValue !== valuesSetByClear[key];
+                    };
+
                     OnyxUtils.cancelPendingMerges(pendingMergeEntries);
 
                     // The only keys that should not be cleared are:
@@ -468,6 +488,7 @@ function clear(keysToPreserve: OnyxKey[] = []): Promise<void> {
                         if (!isKeyToPreserve) {
                             const oldValue = cache.get(key);
                             const newValue = defaultKeyStates[key] ?? null;
+                            valuesSetByClear[key] = newValue;
                             if (newValue !== oldValue) {
                                 cache.set(key, newValue);
 
@@ -507,17 +528,35 @@ function clear(keysToPreserve: OnyxKey[] = []): Promise<void> {
                     // Remove only the items that we want cleared from storage, and reset others to default
                     for (const key of keysToBeClearedFromStorage) cache.drop(key);
                     return Storage.removeItems(keysToBeClearedFromStorage)
-                        .then(() => Storage.multiSet(defaultKeyValuePairs))
+                        .then(() => {
+                            // Save written keys again in case `removeItems` deleted them after they were stored.
+                            const writtenKeyValuePairs = Object.keys(valuesSetByClear)
+                                .filter((key) => !OnyxKeys.isRamOnlyKey(key) && wasWrittenDuringClear(key))
+                                .map((key): [OnyxKey, OnyxValue<OnyxKey>] => [key, cache.get(key)]);
+                            return Storage.multiSet([...defaultKeyValuePairs.filter(([key]) => !wasWrittenDuringClear(key)), ...writtenKeyValuePairs]);
+                        })
                         .then(() => {
                             DevTools.clearState(keysToPreserve);
 
                             onyxSubscriptionManager.resetCacheMissReports();
 
-                            // Notify the subscribers for each key/value group so they can receive the new values
+                            // Notify the subscribers for each key/value group so they can receive the new values.
+                            // Keys written during the clear are skipped, their subscribers already have the newer value.
                             for (const [key, value] of Object.entries(keyValuesToResetIndividually)) {
+                                if (wasWrittenDuringClear(key)) {
+                                    continue;
+                                }
                                 OnyxUtils.notifyKey(key, value);
                             }
                             for (const [key, value] of Object.entries(keyValuesToResetAsCollection)) {
+                                const writtenMemberKeys = Object.keys(value.newValues).filter(wasWrittenDuringClear);
+                                for (const memberKey of writtenMemberKeys) {
+                                    delete value.newValues[memberKey];
+                                    delete value.oldValues[memberKey];
+                                }
+                                if (Object.keys(value.newValues).length === 0) {
+                                    continue;
+                                }
                                 OnyxUtils.notifyCollection(key, value.newValues, value.oldValues);
                             }
                         });
@@ -527,6 +566,8 @@ function clear(keysToPreserve: OnyxKey[] = []): Promise<void> {
             return cache.captureTask(TASK.CLEAR, promise) as Promise<void>;
         }),
     );
+    clearPromise.then(removeWriteListener, removeWriteListener);
+    return clearPromise;
 }
 
 /**

lib/PendingWrites.ts

@@ -8,6 +8,9 @@ const pendingWritesByKey = new Map<OnyxKey, Set<Promise<unknown>>>();
 // In-flight writes that affect every key (Onyx.clear). Any initial notification waits for these.
 const pendingGlobalWrites = new Set<Promise<unknown>>();
 
+// Told about every write as it is issued, so Onyx.clear can tell writes apart from reads refilling the cache.
+const writeListeners = new Set<(keys: OnyxKey[]) => void>();
+
 /**
  * Registers an in-flight write under each key it can change, so `scheduleInitialSubscriberNotification`
  * waits only for the writes relevant to a connecting key. Returns the same promise so callers can wrap a
@@ -16,6 +19,9 @@ const pendingGlobalWrites = new Set<Promise<unknown>>();
 function trackPendingWrite<T>(keys: OnyxKey | OnyxKey[], pendingWrite: Promise<T>): Promise<T> {
     // Drop nullish keys (e.g. a keyless `clear` item) so they never reach `getPendingWritesForKey`'s scan.
     const keyList = (Array.isArray(keys) ? keys : [keys]).filter((key) => typeof key === 'string');
+    for (const listener of writeListeners) {
+        listener(keyList);
+    }
     for (const key of keyList) {
         let pendingWritesForKey = pendingWritesByKey.get(key);
         if (!pendingWritesForKey) {
@@ -50,6 +56,14 @@ function trackPendingGlobalWrite<T>(pendingGlobalWrite: Promise<T>): Promise<T>
     return pendingGlobalWrite;
 }
 
+/**
+ * Calls `listener` with the keys of every write issued from now on, until the returned function is called.
+ */
+function addWriteListener(listener: (keys: OnyxKey[]) => void): () => void {
+    writeListeners.add(listener);
+    return () => writeListeners.delete(listener);
+}
+
 /**
  * In-flight writes that can change the value delivered to a subscriber of `key`: writes to the key
  * itself, writes to any member when `key` is a collection root, and global writes (clear).
@@ -96,11 +110,13 @@ function scheduleInitialSubscriberNotification(key: OnyxKey, notify: () => void)
 function clearPendingWrites(): void {
     pendingWritesByKey.clear();
     pendingGlobalWrites.clear();
+    writeListeners.clear();
 }
 
 export default {
     trackPendingWrite,
     trackPendingGlobalWrite,
+    addWriteListener,
     getPendingWritesForKey,
     scheduleInitialSubscriberNotification,
     clearPendingWrites,

tests/unit/onyxClearConcurrentWriteTest.ts

@@ -0,0 +1,144 @@
+import waitForPromisesToResolve from '../utils/waitForPromisesToResolve';
+import StorageMock from '../../lib/storage';
+import Onyx from '../../lib/Onyx';
+import type OnyxCache from '../../lib/OnyxCache';
+import type GenericCollection from '../utils/GenericCollection';
+import type {Connection} from '../../lib/Onyx';
+
+const ONYX_KEYS = {
+    DEFAULT_KEY: 'defaultKey',
+    REGULAR_KEY: 'regularKey',
+    COLLECTION: {
+        TEST: 'test_',
+    },
+};
+const OLD_VALUE = 'old';
+const NEW_VALUE = 'new';
+const DEFAULT_VALUE = 'default';
+
+describe('Write data while Onyx.clear is in progress', () => {
+    let connections: Connection[] = [];
+    let defaultKeyValue: unknown;
+    let regularKeyValue: unknown;
+    let collectionValue: GenericCollection | undefined;
+    let cache: typeof OnyxCache;
+    let finishRemoveItems: () => void;
+
+    beforeEach(() => {
+        cache = require('../../lib/OnyxCache').default;
+        Onyx.init({
+            keys: ONYX_KEYS,
+            initialKeyStates: {
+                [ONYX_KEYS.DEFAULT_KEY]: DEFAULT_VALUE,
+            },
+        });
+        connections = [
+            Onyx.connect({key: ONYX_KEYS.DEFAULT_KEY, callback: (value) => (defaultKeyValue = value)}),
+            Onyx.connect({key: ONYX_KEYS.REGULAR_KEY, callback: (value) => (regularKeyValue = value)}),
+            Onyx.connect({key: ONYX_KEYS.COLLECTION.TEST, callback: (value) => (collectionValue = value as GenericCollection)}),
+        ];
+        return waitForPromisesToResolve();
+    });
+
+    afterEach(() => {
+        for (const connection of connections) {
+            connection.unsubscribe();
+        }
+        return Onyx.clear();
+    });
+
+    // Holds Onyx.clear() in its storage step, so the test can write while the clear is still running.
+    function holdNextRemoveItems() {
+        const removeItemsMock = StorageMock.removeItems as jest.Mock;
+        const originalImplementation = removeItemsMock.getMockImplementation();
+        removeItemsMock.mockImplementationOnce(
+            (keys: string[]) =>
+                new Promise((resolve) => {
+                    finishRemoveItems = () => resolve(originalImplementation?.(keys));
+                }),
+        );
+    }
+
+    it('should keep values written to keys after the clear reset them', async () => {
+        // Given keys that hold values from a previous session
+        await Onyx.multiSet({
+            [ONYX_KEYS.DEFAULT_KEY]: OLD_VALUE,
+            [ONYX_KEYS.REGULAR_KEY]: OLD_VALUE,
+            [`${ONYX_KEYS.COLLECTION.TEST}1`]: OLD_VALUE,
+        });
+
+        // When a clear starts and new values are written before its storage step finishes, like a new sign-in
+        // saving its session while the previous sign-out is still wiping storage
+        holdNextRemoveItems();
+        const clearPromise = Onyx.clear();
+        await waitForPromisesToResolve();
+        Onyx.set(ONYX_KEYS.DEFAULT_KEY, NEW_VALUE);
+        Onyx.merge(ONYX_KEYS.REGULAR_KEY, NEW_VALUE);
+        Onyx.set(`${ONYX_KEYS.COLLECTION.TEST}1`, NEW_VALUE);
+        await waitForPromisesToResolve();
+        finishRemoveItems();
+        await clearPromise;
+        await waitForPromisesToResolve();
+
+        // Then the clear does not overwrite the new values in subscribers, the cache or storage
+        expect(defaultKeyValue).toBe(NEW_VALUE);
+        expect(regularKeyValue).toBe(NEW_VALUE);
+        expect(collectionValue).toEqual({[`${ONYX_KEYS.COLLECTION.TEST}1`]: NEW_VALUE});
+        expect(cache.get(ONYX_KEYS.DEFAULT_KEY)).toBe(NEW_VALUE);
+        expect(cache.get(ONYX_KEYS.REGULAR_KEY)).toBe(NEW_VALUE);
+        expect(cache.get(`${ONYX_KEYS.COLLECTION.TEST}1`)).toBe(NEW_VALUE);
+        await expect(StorageMock.getItem(ONYX_KEYS.DEFAULT_KEY)).resolves.toBe(NEW_VALUE);
+        await expect(StorageMock.getItem(ONYX_KEYS.REGULAR_KEY)).resolves.toBe(NEW_VALUE);
+        await expect(StorageMock.getItem(`${ONYX_KEYS.COLLECTION.TEST}1`)).resolves.toBe(NEW_VALUE);
+    });
+
+    it('should keep a value written during the clear even when it equals the value from before the clear', async () => {
+        // Given keys that hold values from a previous session
+        await Onyx.multiSet({
+            [ONYX_KEYS.DEFAULT_KEY]: OLD_VALUE,
+            [ONYX_KEYS.REGULAR_KEY]: OLD_VALUE,
+        });
+
+        // When the same values are written again while the clear is running, like a re-auth restoring unchanged session data
+        holdNextRemoveItems();
+        const clearPromise = Onyx.clear();
+        await waitForPromisesToResolve();
+        Onyx.set(ONYX_KEYS.DEFAULT_KEY, OLD_VALUE);
+        Onyx.set(ONYX_KEYS.REGULAR_KEY, OLD_VALUE);
+        await waitForPromisesToResolve();
+        finishRemoveItems();
+        await clearPromise;
+        await waitForPromisesToResolve();
+
+        // Then the written values are kept, because they are new writes and not leftovers from before the clear
+        expect(defaultKeyValue).toBe(OLD_VALUE);
+        expect(regularKeyValue).toBe(OLD_VALUE);
+        expect(cache.get(ONYX_KEYS.DEFAULT_KEY)).toBe(OLD_VALUE);
+        await expect(StorageMock.getItem(ONYX_KEYS.DEFAULT_KEY)).resolves.toBe(OLD_VALUE);
+        await expect(StorageMock.getItem(ONYX_KEYS.REGULAR_KEY)).resolves.toBe(OLD_VALUE);
+    });
+
+    it('should still reset keys that were not written during the clear', async () => {
+        // Given keys that hold values from a previous session
+        await Onyx.multiSet({
+            [ONYX_KEYS.DEFAULT_KEY]: OLD_VALUE,
+            [ONYX_KEYS.REGULAR_KEY]: OLD_VALUE,
+        });
+
+        // When a clear runs and only one key gets a new value while it is in progress
+        holdNextRemoveItems();
+        const clearPromise = Onyx.clear();
+        await waitForPromisesToResolve();
+        Onyx.set(ONYX_KEYS.REGULAR_KEY, NEW_VALUE);
+        await waitForPromisesToResolve();
+        finishRemoveItems();
+        await clearPromise;
+        await waitForPromisesToResolve();
+
+        // Then the written key keeps its new value and the other key is reset to its default
+        expect(regularKeyValue).toBe(NEW_VALUE);
+        expect(defaultKeyValue).toBe(DEFAULT_VALUE);
+        expect(cache.get(ONYX_KEYS.DEFAULT_KEY)).toBe(DEFAULT_VALUE);
+        await expect(StorageMock.getItem(ONYX_KEYS.DEFAULT_KEY)).resolves.toBe(DEFAULT_VALUE);
+    });
+});