Improve useOnyx performance

lib/useOnyx.ts

@@ -1,12 +1,12 @@
-import {deepEqual} from 'fast-equals';
+import {deepEqual, shallowEqual} from 'fast-equals';
 import {useCallback, useEffect, useRef, useSyncExternalStore} from 'react';
 import type {IsEqual} from 'type-fest';
+import Onyx from './Onyx';
+import OnyxCache from './OnyxCache';
 import OnyxUtils from './OnyxUtils';
 import type {CollectionKeyBase, OnyxCollection, OnyxKey, OnyxValue, Selector} from './types';
 import useLiveRef from './useLiveRef';
 import usePrevious from './usePrevious';
-import Onyx from './Onyx';
-import cache from './OnyxCache';
 
 type BaseUseOnyxOptions = {
     /**
@@ -75,7 +75,10 @@ function useOnyx<TKey extends OnyxKey, TReturnValue = OnyxValue<TKey>>(key: TKey
 
     // Stores the previous cached value as it's necessary to compare with the new value in `getSnapshot()`.
     // We initialize it to `null` to simulate that we don't have any value from cache yet.
-    const cachedValueRef = useRef<CachedValue<TKey, TReturnValue> | undefined | null>(null);
+    const previousValueRef = useRef<CachedValue<TKey, TReturnValue> | undefined | null>(null);
+
+    // Stores the newest cached value in order to compare with the previous one and optimize `getSnapshot()` execution.
+    const newValueRef = useRef<CachedValue<TKey, TReturnValue> | undefined | null>(null);
 
     // Stores the previously result returned by the hook, containing the data from cache and the fetch status.
     // We initialize it to `undefined` and `loading` fetch status to simulate the initial result when the hook is loading from the cache.
@@ -91,6 +94,9 @@ function useOnyx<TKey extends OnyxKey, TReturnValue = OnyxValue<TKey>>(key: TKey
     // in `getSnapshot()` to be satisfied several times.
     const isFirstConnectionRef = useRef(true);
 
+    // Indicates if we should get the newest cached value from Onyx during `getSnapshot()` execution.
+    const shouldGetCachedValueRef = useRef(true);
+
     useEffect(() => {
         // These conditions will ensure we can only handle dynamic collection member keys from the same collection.
         if (previousKey === key) {
@@ -116,13 +122,22 @@ function useOnyx<TKey extends OnyxKey, TReturnValue = OnyxValue<TKey>>(key: TKey
     }, [previousKey, key]);
 
     const getSnapshot = useCallback(() => {
-        // We get the value from the cache, supplying a selector too in case it's defined.
-        // If `newValue` is `undefined` it means that the cache doesn't have a value for that key yet.
-        // If `newValue` is `null` or any other value it means that the cache does have a value for that key.
-        // This difference between `undefined` and other values is crucial and it's used to address the following
-        // conditions and use cases.
-        let newValue = getCachedValue<TKey, TReturnValue>(key, selectorRef.current);
-        const hasCacheForKey = cache.hasCacheForKey(key);
+        // We get the value from cache while the first connection to Onyx is being made,
+        // so we can return any cached value right away. After the connection is made, we only
+        // update `newValueRef` when `Onyx.connect()` callback is fired.
+        if (isFirstConnectionRef.current || shouldGetCachedValueRef.current) {
+            // If `newValueRef.current` is `undefined` it means that the cache doesn't have a value for that key yet.
+            // If `newValueRef.current` is `null` or any other value it means that the cache does have a value for that key.
+            // This difference between `undefined` and other values is crucial and it's used to address the following
+            // conditions and use cases.
+            newValueRef.current = getCachedValue<TKey, TReturnValue>(key, selectorRef.current);
+
+            // We set this flag to `false` again since we don't want to get the newest cached value every time `getSnapshot()` is executed,
+            // and only when `Onyx.connect()` callback is fired.
+            shouldGetCachedValueRef.current = false;
+        }
+
+        const hasCacheForKey = OnyxCache.hasCacheForKey(key);
 
         // Since the fetch status can be different given the use cases below, we define the variable right away.
         let newFetchStatus: FetchStatus | undefined;
@@ -131,24 +146,37 @@ function useOnyx<TKey extends OnyxKey, TReturnValue = OnyxValue<TKey>>(key: TKey
         // and fetch status to `loading` to simulate that it is still being loaded until we have the most updated data.
         // If `allowStaleData` is `true` this logic will be ignored and cached value will be used, even if it's stale data.
         if (isFirstConnectionRef.current && OnyxUtils.hasPendingMergeForKey(key) && !options?.allowStaleData) {
-            newValue = undefined;
+            newValueRef.current = undefined;
             newFetchStatus = 'loading';
         }
 
         // If data is not present in cache and `initialValue` is set during the first connection,
         // we set the new value to `initialValue` and fetch status to `loaded` since we already have some data to return to the consumer.
         if (isFirstConnectionRef.current && !hasCacheForKey && options?.initialValue !== undefined) {
-            newValue = (options?.initialValue ?? undefined) as CachedValue<TKey, TReturnValue>;
+            newValueRef.current = (options?.initialValue ?? undefined) as CachedValue<TKey, TReturnValue>;
             newFetchStatus = 'loaded';
         }
 
+        // We do a deep equality check if we are subscribed to a collection key and `selector` is defined,
+        // since each `OnyxUtils.tryGetCachedValue()` call will generate a plain new collection object with new records as well,
+        // all of them created using the `selector` function.
+        // For the other cases we will only deal with object reference checks, so just a shallow equality check is enough.
+        let areValuesEqual: boolean;
+        if (OnyxUtils.isCollectionKey(key) && selectorRef.current) {
+            areValuesEqual = deepEqual(previousValueRef.current ?? undefined, newValueRef.current);
+        } else {
+            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 = cachedValueRef.current === null && hasCacheForKey;
-        if (isCacheSetFirstTime || !deepEqual(cachedValueRef.current ?? undefined, newValue)) {
-            cachedValueRef.current = newValue;
-            resultRef.current = [cachedValueRef.current as CachedValue<TKey, TReturnValue>, {status: newFetchStatus ?? 'loaded'}];
+        const isCacheSetFirstTime = previousValueRef.current === null && hasCacheForKey;
+        if (isCacheSetFirstTime || !areValuesEqual) {
+            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.
+            resultRef.current = [previousValueRef.current as CachedValue<TKey, TReturnValue>, {status: newFetchStatus ?? 'loaded'}];
         }
 
         return resultRef.current;
@@ -159,9 +187,14 @@ function useOnyx<TKey extends OnyxKey, TReturnValue = OnyxValue<TKey>>(key: TKey
             connectionIDRef.current = Onyx.connect<CollectionKeyBase>({
                 key,
                 callback: () => {
-                    // We don't need to update the Onyx cache again here, when `callback` is called the cache is already
-                    // expected to be updated, so we just signal that the store changed and `getSnapshot()` can be called again.
+                    // Signals that the first connection was made, so some logics in `getSnapshot()`
+                    // won't be executed anymore.
                     isFirstConnectionRef.current = false;
+
+                    // Signals that we want to get the newest cached value again in `getSnapshot()`.
+                    shouldGetCachedValueRef.current = true;
+
+                    // Finally, we signal that the store changed, making `getSnapshot()` be called again.
                     onStoreChange();
                 },
                 initWithStoredValues: options?.initWithStoredValues,
@@ -204,4 +237,4 @@ function useOnyx<TKey extends OnyxKey, TReturnValue = OnyxValue<TKey>>(key: TKey
 
 export default useOnyx;
 
-export type {UseOnyxResult, ResultMetadata, FetchStatus};
+export type {FetchStatus, ResultMetadata, UseOnyxResult};

tests/unit/useOnyxTest.ts

@@ -195,7 +195,27 @@ describe('useOnyx', () => {
             expect(result.current[1].status).toEqual('loaded');
         });
 
-        it('should not change selected data if a property outside the selector was changed', async () => {
+        it('should not change selected data if a property outside that data was changed', async () => {
+            Onyx.set(ONYXKEYS.TEST_KEY, {id: 'test_id', name: 'test_name'});
+
+            const {result} = renderHook(() =>
+                useOnyx(ONYXKEYS.TEST_KEY, {
+                    // @ts-expect-error bypass
+                    selector: (entry: OnyxEntry<{id: string; name: string}>) => ({id: entry?.id}),
+                }),
+            );
+
+            await act(async () => waitForPromisesToResolve());
+
+            const oldResult = result.current;
+
+            await act(async () => Onyx.merge(ONYXKEYS.TEST_KEY, {name: 'test_name_changed'}));
+
+            // must be the same reference
+            expect(oldResult).toBe(result.current);
+        });
+
+        it('should not change selected collection data if a property outside that data was changed', async () => {
             // @ts-expect-error bypass
             Onyx.mergeCollection(ONYXKEYS.COLLECTION.TEST_KEY, {
                 [`${ONYXKEYS.COLLECTION.TEST_KEY}entry1`]: {id: 'entry1_id', name: 'entry1_name'},
@@ -206,7 +226,7 @@ describe('useOnyx', () => {
             const {result} = renderHook(() =>
                 useOnyx(ONYXKEYS.COLLECTION.TEST_KEY, {
                     // @ts-expect-error bypass
-                    selector: (entry: OnyxEntry<{id: string; name: string}>) => entry?.id,
+                    selector: (entry: OnyxEntry<{id: string; name: string}>) => ({id: entry?.id}),
                 }),
             );