Fix problems related to Onyx.clear and useOnyx

lib/Onyx.ts

@@ -2,7 +2,7 @@
 import _ from 'underscore';
 import lodashPick from 'lodash/pick';
 import * as Logger from './Logger';
-import cache from './OnyxCache';
+import cache, {TASK} from './OnyxCache';
 import * as PerformanceUtils from './PerformanceUtils';
 import Storage from './storage';
 import utils from './utils';
@@ -449,7 +449,7 @@ function clear(keysToPreserve: OnyxKey[] = []): Promise<void> {
     const defaultKeyStates = OnyxUtils.getDefaultKeyStates();
     const initialKeys = Object.keys(defaultKeyStates);
 
-    return OnyxUtils.getAllKeys()
+    const promise = OnyxUtils.getAllKeys()
         .then((cachedKeys) => {
             cache.clearNullishStorageKeys();
 
@@ -529,13 +529,16 @@ function clear(keysToPreserve: OnyxKey[] = []): Promise<void> {
             // Remove only the items that we want cleared from storage, and reset others to default
             keysToBeClearedFromStorage.forEach((key) => cache.drop(key));
             return Storage.removeItems(keysToBeClearedFromStorage)
+                .then(() => connectionManager.refreshSessionID())
                 .then(() => Storage.multiSet(defaultKeyValuePairs))
                 .then(() => {
                     DevTools.clearState(keysToPreserve);
                     return Promise.all(updatePromises);
                 });
         })
         .then(() => undefined);
+
+    return cache.captureTask(TASK.CLEAR, promise) as Promise<void>;
 }
 
 function updateSnapshots(data: OnyxUpdate[]) {

lib/OnyxCache.ts

@@ -1,8 +1,18 @@
 import {deepEqual} from 'fast-equals';
 import bindAll from 'lodash/bindAll';
+import type {ValueOf} from 'type-fest';
 import utils from './utils';
 import type {OnyxKey, OnyxValue} from './types';
 
+// Task constants
+const TASK = {
+    GET: 'get',
+    GET_ALL_KEYS: 'getAllKeys',
+    CLEAR: 'clear',
+} as const;
+
+type CacheTask = ValueOf<typeof TASK> | `${ValueOf<typeof TASK>}:${string}`;
+
 /**
  * In memory cache providing data by reference
  * Encapsulates Onyx cache related functionality
@@ -172,7 +182,7 @@ class OnyxCache {
      * Check whether the given task is already running
      * @param taskName - unique name given for the task
      */
-    hasPendingTask(taskName: string): boolean {
+    hasPendingTask(taskName: CacheTask): boolean {
         return this.pendingPromises.get(taskName) !== undefined;
     }
 
@@ -182,7 +192,7 @@ class OnyxCache {
      * provided from this function
      * @param taskName - unique name given for the task
      */
-    getTaskPromise(taskName: string): Promise<OnyxValue<OnyxKey> | OnyxKey[]> | undefined {
+    getTaskPromise(taskName: CacheTask): Promise<OnyxValue<OnyxKey> | OnyxKey[]> | undefined {
         return this.pendingPromises.get(taskName);
     }
 
@@ -191,7 +201,7 @@ class OnyxCache {
      * hook up to the promise if it's still pending
      * @param taskName - unique name for the task
      */
-    captureTask(taskName: string, promise: Promise<OnyxValue<OnyxKey>>): Promise<OnyxValue<OnyxKey>> {
+    captureTask(taskName: CacheTask, promise: Promise<OnyxValue<OnyxKey>>): Promise<OnyxValue<OnyxKey>> {
         const returnPromise = promise.finally(() => {
             this.pendingPromises.delete(taskName);
         });
@@ -242,3 +252,5 @@ class OnyxCache {
 const instance = new OnyxCache();
 
 export default instance;
+export {TASK};
+export type {CacheTask};

lib/OnyxConnectionManager.ts

@@ -73,12 +73,28 @@ class OnyxConnectionManager {
      */
     private lastCallbackID: number;
 
+    /**
+     * Stores the last generated session ID for the connection manager. The current session ID
+     * is appended to the connection IDs and it's used to create new different connections for the same key
+     * when `refreshSessionID()` is called.
+     *
+     * When calling `Onyx.clear()` after a logout operation some connections might remain active as they
+     * aren't tied to the React's lifecycle e.g. `Onyx.connect()` usage, causing infinite loading state issues to new `useOnyx()` subscribers
+     * that are connecting to the same key as we didn't populate the cache again because we are still reusing such connections.
+     *
+     * To elimitate this problem, the session ID must be refreshed during the `Onyx.clear()` call (by using `refreshSessionID()`)
+     * in order to create fresh connections when new subscribers connect to the same keys again, allowing them
+     * to use the cache system correctly and avoid the mentioned issues in `useOnyx()`.
+     */
+    private sessionID: string;
+
     constructor() {
         this.connectionsMap = new Map();
         this.lastCallbackID = 0;
+        this.sessionID = Str.guid();
 
         // Binds all public methods to prevent problems with `this`.
-        bindAll(this, 'generateConnectionID', 'fireCallbacks', 'connect', 'disconnect', 'disconnectAll', 'addToEvictionBlockList', 'removeFromEvictionBlockList');
+        bindAll(this, 'generateConnectionID', 'fireCallbacks', 'connect', 'disconnect', 'disconnectAll', 'refreshSessionID', 'addToEvictionBlockList', 'removeFromEvictionBlockList');
     }
 
     /**
@@ -89,7 +105,10 @@ class OnyxConnectionManager {
      */
     private generateConnectionID<TKey extends OnyxKey>(connectOptions: ConnectOptions<TKey>): string {
         const {key, initWithStoredValues, reuseConnection, waitForCollectionCallback} = connectOptions;
-        let suffix = '';
+
+        // The current session ID is appended to the connection ID so we can have different connections
+        // after an `Onyx.clear()` operation.
+        let suffix = `,sessionID=${this.sessionID}`;
 
         // We will generate a unique ID in any of the following situations:
         // - `reuseConnection` is `false`. That means the subscriber explicitly wants the connection to not be reused.
@@ -229,6 +248,13 @@ class OnyxConnectionManager {
         this.connectionsMap.clear();
     }
 
+    /**
+     * Refreshes the connection manager's session ID.
+     */
+    refreshSessionID(): void {
+        this.sessionID = Str.guid();
+    }
+
     /**
      * Adds the connection to the eviction block list. Connections added to this list can never be evicted.
      * */

lib/OnyxUtils.ts

@@ -6,7 +6,7 @@ import type {ValueOf} from 'type-fest';
 import DevTools from './DevTools';
 import * as Logger from './Logger';
 import type Onyx from './Onyx';
-import cache from './OnyxCache';
+import cache, {TASK} from './OnyxCache';
 import * as PerformanceUtils from './PerformanceUtils';
 import * as Str from './Str';
 import unstable_batchedUpdates from './batch';
@@ -244,7 +244,7 @@ function get<TKey extends OnyxKey, TValue extends OnyxValue<TKey>>(key: TKey): P
         return Promise.resolve(cache.get(key) as TValue);
     }
 
-    const taskName = `get:${key}`;
+    const taskName = `${TASK.GET}:${key}` as const;
 
     // When a value retrieving task for this key is still running hook to it
     if (cache.hasPendingTask(taskName)) {
@@ -296,7 +296,7 @@ function multiGet<TKey extends OnyxKey>(keys: CollectionKeyBase[]): Promise<Map<
             return;
         }
 
-        const pendingKey = `get:${key}`;
+        const pendingKey = `${TASK.GET}:${key}` as const;
         if (cache.hasPendingTask(pendingKey)) {
             pendingTasks.push(cache.getTaskPromise(pendingKey) as Promise<OnyxValue<TKey>>);
             pendingKeys.push(key);
@@ -378,11 +378,9 @@ function getAllKeys(): Promise<Set<OnyxKey>> {
         return Promise.resolve(cachedKeys);
     }
 
-    const taskName = 'getAllKeys';
-
     // When a value retrieving task for all keys is still running hook to it
-    if (cache.hasPendingTask(taskName)) {
-        return cache.getTaskPromise(taskName) as Promise<Set<OnyxKey>>;
+    if (cache.hasPendingTask(TASK.GET_ALL_KEYS)) {
+        return cache.getTaskPromise(TASK.GET_ALL_KEYS) as Promise<Set<OnyxKey>>;
     }
 
     // Otherwise retrieve the keys from storage and capture a promise to aid concurrent usages
@@ -393,7 +391,7 @@ function getAllKeys(): Promise<Set<OnyxKey>> {
         return cache.getAllKeys();
     });
 
-    return cache.captureTask(taskName, promise) as Promise<Set<OnyxKey>>;
+    return cache.captureTask(TASK.GET_ALL_KEYS, promise) as Promise<Set<OnyxKey>>;
 }
 
 /**

lib/useOnyx.ts

@@ -1,6 +1,6 @@
 import {deepEqual, shallowEqual} from 'fast-equals';
 import {useCallback, useEffect, useRef, useSyncExternalStore} from 'react';
-import OnyxCache from './OnyxCache';
+import OnyxCache, {TASK} from './OnyxCache';
 import type {Connection} from './OnyxConnectionManager';
 import connectionManager from './OnyxConnectionManager';
 import OnyxUtils from './OnyxUtils';
@@ -23,6 +23,12 @@ type BaseUseOnyxOptions = {
      * If set to `true`, data will be retrieved from cache during the first render even if there is a pending merge for the key.
      */
     allowStaleData?: boolean;
+
+    /**
+     * If set to `false`, the connection won't be reused between other subscribers that are listening to the same Onyx key
+     * with the same connect configurations.
+     */
+    reuseConnection?: boolean;
 };
 
 type UseOnyxInitialValueOption<TInitialValue> = {
@@ -230,11 +236,14 @@ function useOnyx<TKey extends OnyxKey, TReturnValue = OnyxValue<TKey>>(key: TKey
             areValuesEqual = shallowEqual(previousValueRef.current ?? undefined, newValueRef.current);
         }
 
-        // If the previously cached value is different from the new value, we update both cached value
-        // and the result to be returned by the hook.
-        // If the cache was set for the first time, we also update the cached value and the result.
-        const isCacheSetFirstTime = previousValueRef.current === null && hasCacheForKey;
-        if (isCacheSetFirstTime || !areValuesEqual) {
+        // We updated the cached value and the result in the following conditions:
+        // We will update the cached value and the result in any of the following situations:
+        // - The previously cached value is different from the new value.
+        // - The previously cached value is `null` (not set from cache yet) and we have cache for this key
+        //   OR we have a pending `Onyx.clear()` task (if `Onyx.clear()` is running cache might not be available anymore
+        //   so we update the cached value/result right away in order to prevent infinite loading state issues).
+        const shouldUpdateResult = !areValuesEqual || (previousValueRef.current === null && (hasCacheForKey || OnyxCache.hasPendingTask(TASK.CLEAR)));
+        if (shouldUpdateResult) {
             previousValueRef.current = newValueRef.current;
 
             // If the new value is `null` we default it to `undefined` to ensure the consumer gets a consistent result from the hook.
@@ -261,6 +270,7 @@ function useOnyx<TKey extends OnyxKey, TReturnValue = OnyxValue<TKey>>(key: TKey
                 },
                 initWithStoredValues: options?.initWithStoredValues,
                 waitForCollectionCallback: OnyxUtils.isCollectionKey(key) as true,
+                reuseConnection: options?.reuseConnection,
             });
 
             checkEvictableKey();
@@ -274,7 +284,7 @@ function useOnyx<TKey extends OnyxKey, TReturnValue = OnyxValue<TKey>>(key: TKey
                 isFirstConnectionRef.current = false;
             };
         },
-        [key, options?.initWithStoredValues, checkEvictableKey],
+        [key, options?.initWithStoredValues, options?.reuseConnection, checkEvictableKey],
     );
 
     useEffect(() => {

tests/unit/OnyxConnectionManagerTest.ts

@@ -10,12 +10,14 @@ import type {WithOnyxInstance} from '../../lib/withOnyx/types';
 import type GenericCollection from '../utils/GenericCollection';
 import waitForPromisesToResolve from '../utils/waitForPromisesToResolve';
 
-// We need access to `connectionsMap` and `generateConnectionID` during the tests but the properties are private,
+// We need access to some internal properties of `connectionManager` during the tests but they are private,
 // so this workaround allows us to have access to them.
 // eslint-disable-next-line dot-notation
 const connectionsMap = connectionManager['connectionsMap'];
 // eslint-disable-next-line dot-notation
 const generateConnectionID = connectionManager['generateConnectionID'];
+// eslint-disable-next-line dot-notation
+const getSessionID = () => connectionManager['sessionID'];
 
 function generateEmptyWithOnyxInstance() {
     return new (class {
@@ -51,31 +53,39 @@ describe('OnyxConnectionManager', () => {
     describe('generateConnectionID', () => {
         it('should generate a stable connection ID', async () => {
             const connectionID = generateConnectionID({key: ONYXKEYS.TEST_KEY});
-            expect(connectionID).toEqual(`onyxKey=${ONYXKEYS.TEST_KEY},initWithStoredValues=true,waitForCollectionCallback=false`);
+            expect(connectionID).toEqual(`onyxKey=${ONYXKEYS.TEST_KEY},initWithStoredValues=true,waitForCollectionCallback=false,sessionID=${getSessionID()}`);
         });
 
         it("should generate a stable connection ID regardless of the order which the option's properties were passed", async () => {
             const connectionID = generateConnectionID({key: ONYXKEYS.TEST_KEY, waitForCollectionCallback: true, initWithStoredValues: true});
-            expect(connectionID).toEqual(`onyxKey=${ONYXKEYS.TEST_KEY},initWithStoredValues=true,waitForCollectionCallback=true`);
+            expect(connectionID).toEqual(`onyxKey=${ONYXKEYS.TEST_KEY},initWithStoredValues=true,waitForCollectionCallback=true,sessionID=${getSessionID()}`);
         });
 
         it('should generate unique connection IDs if certain options are passed', async () => {
             const connectionID1 = generateConnectionID({key: ONYXKEYS.TEST_KEY, reuseConnection: false});
             const connectionID2 = generateConnectionID({key: ONYXKEYS.TEST_KEY, reuseConnection: false});
-            expect(connectionID1.startsWith(`onyxKey=${ONYXKEYS.TEST_KEY},initWithStoredValues=true,waitForCollectionCallback=false,uniqueID=`)).toBeTruthy();
-            expect(connectionID2.startsWith(`onyxKey=${ONYXKEYS.TEST_KEY},initWithStoredValues=true,waitForCollectionCallback=false,uniqueID=`)).toBeTruthy();
+            expect(connectionID1.startsWith(`onyxKey=${ONYXKEYS.TEST_KEY},initWithStoredValues=true,waitForCollectionCallback=false,sessionID=${getSessionID()},uniqueID=`)).toBeTruthy();
+            expect(connectionID2.startsWith(`onyxKey=${ONYXKEYS.TEST_KEY},initWithStoredValues=true,waitForCollectionCallback=false,sessionID=${getSessionID()},uniqueID=`)).toBeTruthy();
             expect(connectionID1).not.toEqual(connectionID2);
 
             const connectionID3 = generateConnectionID({key: ONYXKEYS.TEST_KEY, initWithStoredValues: false});
-            expect(connectionID3.startsWith(`onyxKey=${ONYXKEYS.TEST_KEY},initWithStoredValues=false,waitForCollectionCallback=false,uniqueID=`)).toBeTruthy();
+            expect(connectionID3.startsWith(`onyxKey=${ONYXKEYS.TEST_KEY},initWithStoredValues=false,waitForCollectionCallback=false,sessionID=${getSessionID()},uniqueID=`)).toBeTruthy();
 
             const connectionID4 = generateConnectionID({
                 key: ONYXKEYS.TEST_KEY,
                 displayName: 'Component1',
                 statePropertyName: 'prop1',
                 withOnyxInstance: generateEmptyWithOnyxInstance(),
             } as WithOnyxConnectOptions<OnyxKey>);
-            expect(connectionID4.startsWith(`onyxKey=${ONYXKEYS.TEST_KEY},initWithStoredValues=true,waitForCollectionCallback=false,uniqueID=`)).toBeTruthy();
+            expect(connectionID4.startsWith(`onyxKey=${ONYXKEYS.TEST_KEY},initWithStoredValues=true,waitForCollectionCallback=false,sessionID=${getSessionID()},uniqueID=`)).toBeTruthy();
+        });
+
+        it('should generate an unique connection ID if the session ID is changed', async () => {
+            const connectionID1 = generateConnectionID({key: ONYXKEYS.TEST_KEY});
+            connectionManager.refreshSessionID();
+            const connectionID2 = generateConnectionID({key: ONYXKEYS.TEST_KEY});
+
+            expect(connectionID1).not.toEqual(connectionID2);
         });
     });
 
@@ -320,6 +330,56 @@ describe('OnyxConnectionManager', () => {
             }).not.toThrow();
         });
 
+        it('should create a separate connection for the same key after a Onyx.clear() call', async () => {
+            await StorageMock.setItem(ONYXKEYS.TEST_KEY, 'test');
+
+            const callback1 = jest.fn();
+            connectionManager.connect({key: ONYXKEYS.TEST_KEY, callback: callback1});
+            expect(connectionsMap.size).toEqual(1);
+
+            await act(async () => waitForPromisesToResolve());
+
+            expect(callback1).toHaveBeenCalledTimes(1);
+            expect(callback1).toHaveBeenCalledWith('test', ONYXKEYS.TEST_KEY);
+            callback1.mockReset();
+
+            await act(async () => Onyx.clear());
+
+            expect(callback1).toHaveBeenCalledTimes(1);
+            expect(callback1).toHaveBeenCalledWith(undefined, ONYXKEYS.TEST_KEY);
+            callback1.mockReset();
+
+            const callback2 = jest.fn();
+            connectionManager.connect({key: ONYXKEYS.TEST_KEY, callback: callback2});
+
+            const callback3 = jest.fn();
+            connectionManager.connect({key: ONYXKEYS.TEST_KEY, callback: callback3});
+
+            // We expect to have two connections for ONYXKEYS.TEST_KEY, one for the first subscription before Onyx.clear(),
+            // and the other for the two subscriptions with the same key after Onyx.clear().
+            expect(connectionsMap.size).toEqual(2);
+
+            await act(async () => waitForPromisesToResolve());
+
+            expect(callback2).toHaveBeenCalledTimes(1);
+            expect(callback2).toHaveBeenCalledWith(undefined, undefined);
+            expect(callback3).toHaveBeenCalledTimes(1);
+            expect(callback3).toHaveBeenCalledWith(undefined, undefined);
+            callback1.mockReset();
+            callback2.mockReset();
+            callback3.mockReset();
+
+            Onyx.merge(ONYXKEYS.TEST_KEY, 'test2');
+            await act(async () => waitForPromisesToResolve());
+
+            expect(callback1).toHaveBeenCalledTimes(1);
+            expect(callback1).toHaveBeenCalledWith('test2', ONYXKEYS.TEST_KEY);
+            expect(callback2).toHaveBeenCalledTimes(1);
+            expect(callback2).toHaveBeenCalledWith('test2', ONYXKEYS.TEST_KEY);
+            expect(callback3).toHaveBeenCalledTimes(1);
+            expect(callback3).toHaveBeenCalledWith('test2', ONYXKEYS.TEST_KEY);
+        });
+
         describe('withOnyx', () => {
             it('should connect to a key two times with withOnyx and create two connections instead of one', async () => {
                 await StorageMock.setItem(ONYXKEYS.TEST_KEY, 'test');
@@ -406,6 +466,25 @@ describe('OnyxConnectionManager', () => {
         });
     });
 
+    describe('refreshSessionID', () => {
+        it('should create a separate connection for the same key if the session ID changes', async () => {
+            await StorageMock.setItem(ONYXKEYS.TEST_KEY, 'test');
+            await StorageMock.setItem(ONYXKEYS.TEST_KEY_2, 'test2');
+
+            const connection1 = connectionManager.connect({key: ONYXKEYS.TEST_KEY, callback: jest.fn()});
+
+            expect(connectionsMap.size).toEqual(1);
+
+            connectionManager.refreshSessionID();
+
+            const connection2 = connectionManager.connect({key: ONYXKEYS.TEST_KEY, callback: jest.fn()});
+
+            expect(connectionsMap.size).toEqual(2);
+            expect(connectionsMap.has(connection1.id)).toBeTruthy();
+            expect(connectionsMap.has(connection2.id)).toBeTruthy();
+        });
+    });
+
     describe('addToEvictionBlockList / removeFromEvictionBlockList', () => {
         it('should add and remove connections from the eviction block list correctly', async () => {
             const evictionBlocklist = OnyxUtils.getEvictionBlocklist();