Improve useOnyx performance

lib/useOnyx.ts

@@ -1,15 +1,15 @@
-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 OnyxUtils from './OnyxUtils';
 import type {CollectionKeyBase, KeyValueMapping, NonNull, OnyxCollection, OnyxEntry, OnyxKey, Selector} from './types';
 import useLiveRef from './useLiveRef';
 import usePrevious from './usePrevious';
-import Onyx from './Onyx';
 
 /**
  * Represents a Onyx value that can be either a single entry or a collection of entries, depending on the `TKey` provided.
- * It's a variation of `OnyxValue` type that is read-only and excludes the `null` type.
+ * It's a variation of `OnyxValue` type that excludes the `null` type.
  */
 type UseOnyxValue<TKey extends OnyxKey> = string extends TKey
     ? unknown
@@ -84,7 +84,10 @@ function useOnyx<TKey extends OnyxKey, TReturnValue = UseOnyxValue<TKey>>(key: T
 
     // Stores the previous cached value as it's necessary to compare with the new value in `getSnapshot()`.
     // We initialize it to `undefined` to simulate that we don't have any value from cache yet.
-    const cachedValueRef = useRef<CachedValue<TKey, TReturnValue> | undefined>(undefined);
+    const previousValueRef = useRef<CachedValue<TKey, TReturnValue> | undefined>(undefined);
+
+    // Stores the newest cached value in order to compare with the previous one and optimize `getSnapshot()` execution.
+    const newValueRef = useRef<CachedValue<TKey, TReturnValue> | undefined>(undefined);
 
     // 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.
@@ -100,6 +103,9 @@ function useOnyx<TKey extends OnyxKey, TReturnValue = UseOnyxValue<TKey>>(key: T
     // 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) {
@@ -125,12 +131,20 @@ function useOnyx<TKey extends OnyxKey, TReturnValue = UseOnyxValue<TKey>>(key: T
     }, [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 if 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);
+        // 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 if 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;
+        }
 
         // Since the fetch status can be different given the use cases below, we define the variable right away.
         let newFetchStatus: FetchStatus | undefined;
@@ -139,24 +153,35 @@ function useOnyx<TKey extends OnyxKey, TReturnValue = UseOnyxValue<TKey>>(key: T
         // 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 (if it's `undefined`) 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 && newValue === undefined && options?.initialValue !== undefined) {
-            newValue = options?.initialValue as CachedValue<TKey, TReturnValue>;
+        if (isFirstConnectionRef.current && newValueRef.current === undefined && options?.initialValue !== undefined) {
+            newValueRef.current = options?.initialValue 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 = false;
+        if (OnyxUtils.isCollectionKey(key) && selectorRef.current) {
+            areValuesEqual = deepEqual(previousValueRef.current, newValueRef.current);
+        } else {
+            areValuesEqual = shallowEqual(previousValueRef.current, 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 (!deepEqual(cachedValueRef.current, newValue)) {
-            cachedValueRef.current = newValue;
+        if (!areValuesEqual) {
+            previousValueRef.current = newValueRef.current;
 
             // If the new value is `null` we default it to `undefined` to ensure the consumer get a consistent result from the hook.
-            resultRef.current = [(cachedValueRef.current ?? undefined) as CachedValue<TKey, TReturnValue>, {status: newFetchStatus ?? 'loaded'}];
+            resultRef.current = [(previousValueRef.current ?? undefined) as CachedValue<TKey, TReturnValue>, {status: newFetchStatus ?? 'loaded'}];
         }
 
         return resultRef.current;
@@ -167,9 +192,14 @@ function useOnyx<TKey extends OnyxKey, TReturnValue = UseOnyxValue<TKey>>(key: T
             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,
@@ -212,4 +242,4 @@ function useOnyx<TKey extends OnyxKey, TReturnValue = UseOnyxValue<TKey>>(key: T
 
 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 tha 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}),
                 }),
             );