Merge pull request #774 from callstack-internal/feature/structural-sharing-cache-pr-4

[Structural Sharing] PR 4: Reference equality in notification and subscription paths

Author: puneetlath

API-INTERNAL.md

@@ -293,23 +293,6 @@ If the requested key is a collection, it will return an object with all the coll
 When a collection of keys change, search for any callbacks matching the collection key and trigger those callbacks
 
 **Kind**: global function  
-
-* [keysChanged()](#keysChanged)
-    * [~isSubscribedToCollectionKey](#keysChanged..isSubscribedToCollectionKey)
-    * [~isSubscribedToCollectionMemberKey](#keysChanged..isSubscribedToCollectionMemberKey)
-
-<a name="keysChanged..isSubscribedToCollectionKey"></a>
-
-### keysChanged~isSubscribedToCollectionKey
-e.g. Onyx.connect({key: ONYXKEYS.COLLECTION.REPORT, callback: ...});
-
-**Kind**: inner constant of [<code>keysChanged</code>](#keysChanged)  
-<a name="keysChanged..isSubscribedToCollectionMemberKey"></a>
-
-### keysChanged~isSubscribedToCollectionMemberKey
-e.g. Onyx.connect({key: `${ONYXKEYS.COLLECTION.REPORT}{reportID}`, callback: ...});
-
-**Kind**: inner constant of [<code>keysChanged</code>](#keysChanged)  
 <a name="keyChanged"></a>
 
 ## keyChanged()

lib/OnyxUtils.ts

@@ -1,4 +1,4 @@
-import {deepEqual, shallowEqual} from 'fast-equals';
+import {shallowEqual} from 'fast-equals';
 import type {ValueOf} from 'type-fest';
 import _ from 'underscore';
 import DevTools from './DevTools';
@@ -510,8 +510,8 @@ function getCachedCollection<TKey extends CollectionKeyBase>(collectionKey: TKey
             return filteredCollection;
         }
 
-        // Return a copy to avoid mutations affecting the cache
-        return {...collectionData};
+        // Snapshot is frozen — safe to return by reference
+        return collectionData;
     }
 
     // Fallback to original implementation if collection data not available
@@ -546,81 +546,67 @@ function keysChanged<TKey extends CollectionKeyBase>(
     partialCollection: OnyxCollection<KeyValueMapping[TKey]>,
     partialPreviousCollection: OnyxCollection<KeyValueMapping[TKey]> | undefined,
 ): void {
-    // We prepare the "cached collection" which is the entire collection + the new partial data that
-    // was merged in via mergeCollection().
     const cachedCollection = getCachedCollection(collectionKey);
-
     const previousCollection = partialPreviousCollection ?? {};
-
-    // We are iterating over all subscribers similar to keyChanged(). However, we are looking for subscribers who are subscribing to either a collection key or
-    // individual collection key member for the collection that is being updated. It is important to note that the collection parameter cane be a PARTIAL collection
-    // and does not represent all of the combined keys and values for a collection key. It is just the "new" data that was merged in via mergeCollection().
-    const stateMappingKeys = Object.keys(callbackToStateMapping);
-
-    for (const stateMappingKey of stateMappingKeys) {
-        const subscriber = callbackToStateMapping[stateMappingKey];
-        if (!subscriber) {
-            continue;
+    const changedMemberKeys = Object.keys(partialCollection ?? {});
+
+    // Use indexed lookup instead of scanning all subscribers.
+    // We need subscribers for: (1) the collection key itself, and (2) individual changed member keys.
+    const collectionSubscriberIDs = onyxKeyToSubscriptionIDs.get(collectionKey) ?? [];
+    const memberSubscriberIDs: number[] = [];
+    for (const memberKey of changedMemberKeys) {
+        const ids = onyxKeyToSubscriptionIDs.get(memberKey);
+        if (ids) {
+            for (const id of ids) {
+                memberSubscriberIDs.push(id);
+            }
         }
+    }
 
-        // Skip iteration if we do not have a collection key or a collection member key on this subscriber
-        if (!Str.startsWith(subscriber.key, collectionKey)) {
+    // Notify collection-level subscribers
+    for (const subID of collectionSubscriberIDs) {
+        const subscriber = callbackToStateMapping[subID];
+        if (!subscriber || typeof subscriber.callback !== 'function') {
             continue;
         }
 
-        /**
-         * e.g. Onyx.connect({key: ONYXKEYS.COLLECTION.REPORT, callback: ...});
-         */
-        const isSubscribedToCollectionKey = subscriber.key === collectionKey;
-
-        /**
-         * e.g. Onyx.connect({key: `${ONYXKEYS.COLLECTION.REPORT}{reportID}`, callback: ...});
-         */
-        const isSubscribedToCollectionMemberKey = OnyxKeys.isCollectionMemberKey(collectionKey, subscriber.key);
-
-        // Regular Onyx.connect() subscriber found.
-        if (typeof subscriber.callback === 'function') {
-            try {
-                // If they are subscribed to the collection key and using waitForCollectionCallback then we'll
-                // send the whole cached collection.
-                if (isSubscribedToCollectionKey) {
-                    lastConnectionCallbackData.set(subscriber.subscriptionID, {value: cachedCollection, matchedKey: subscriber.key});
-
-                    if (subscriber.waitForCollectionCallback) {
-                        subscriber.callback(cachedCollection, subscriber.key, partialCollection);
-                        continue;
-                    }
+        try {
+            lastConnectionCallbackData.set(subscriber.subscriptionID, {value: cachedCollection, matchedKey: subscriber.key});
 
-                    // If they are not using waitForCollectionCallback then we notify the subscriber with
-                    // the new merged data but only for any keys in the partial collection.
-                    const dataKeys = Object.keys(partialCollection ?? {});
-                    for (const dataKey of dataKeys) {
-                        if (deepEqual(cachedCollection[dataKey], previousCollection[dataKey])) {
-                            continue;
-                        }
+            if (subscriber.waitForCollectionCallback) {
+                subscriber.callback(cachedCollection, subscriber.key, partialCollection);
+                continue;
+            }
 
-                        subscriber.callback(cachedCollection[dataKey], dataKey);
-                    }
+            // Not using waitForCollectionCallback — notify per changed key
+            for (const dataKey of changedMemberKeys) {
+                if (cachedCollection[dataKey] === previousCollection[dataKey]) {
                     continue;
                 }
+                subscriber.callback(cachedCollection[dataKey], dataKey);
+            }
+        } catch (error) {
+            Logger.logAlert(`[OnyxUtils.keysChanged] Subscriber callback threw an error for key '${collectionKey}': ${error}`);
+        }
+    }
 
-                // And if the subscriber is specifically only tracking a particular collection member key then we will
-                // notify them with the cached data for that key only.
-                if (isSubscribedToCollectionMemberKey) {
-                    if (deepEqual(cachedCollection[subscriber.key], previousCollection[subscriber.key])) {
-                        continue;
-                    }
+    // Notify member-level subscribers (e.g. subscribed to `report_123`)
+    for (const subID of memberSubscriberIDs) {
+        const subscriber = callbackToStateMapping[subID];
+        if (!subscriber || typeof subscriber.callback !== 'function') {
+            continue;
+        }
 
-                    const subscriberCallback = subscriber.callback as DefaultConnectCallback<TKey>;
-                    subscriberCallback(cachedCollection[subscriber.key], subscriber.key as TKey);
-                    lastConnectionCallbackData.set(subscriber.subscriptionID, {value: cachedCollection[subscriber.key], matchedKey: subscriber.key});
-                    continue;
-                }
+        if (cachedCollection[subscriber.key] === previousCollection[subscriber.key]) {
+            continue;
+        }
 
-                continue;
-            } catch (error) {
-                Logger.logAlert(`[OnyxUtils.keysChanged] Subscriber callback threw an error for key '${collectionKey}': ${error}`);
-            }
+        try {
+            const subscriberCallback = subscriber.callback as DefaultConnectCallback<TKey>;
+            subscriberCallback(cachedCollection[subscriber.key], subscriber.key as TKey);
+            lastConnectionCallbackData.set(subscriber.subscriptionID, {value: cachedCollection[subscriber.key], matchedKey: subscriber.key});
+        } catch (error) {
+            Logger.logAlert(`[OnyxUtils.keysChanged] Subscriber callback threw an error for key '${collectionKey}': ${error}`);
         }
     }
 }
@@ -660,6 +646,9 @@ function keyChanged<TKey extends OnyxKey>(
         }
     }
 
+    // Cache the collection snapshot per dispatch so all subscribers to the same collection
+    // see a consistent view, even if an earlier subscriber's callback synchronously writes
+    // to the same collection.
     const cachedCollections: Record<string, ReturnType<typeof getCachedCollection>> = {};
 
     for (const stateMappingKey of stateMappingKeys) {
@@ -682,14 +671,13 @@ function keyChanged<TKey extends OnyxKey>(
                     if (isProcessingCollectionUpdate) {
                         continue;
                     }
+                    // Cache once per dispatch to ensure all subscribers see a consistent snapshot
+                    // even if a previous callback synchronously wrote to the same collection.
                     let cachedCollection = cachedCollections[subscriber.key];
-
                     if (!cachedCollection) {
                         cachedCollection = getCachedCollection(subscriber.key);
                         cachedCollections[subscriber.key] = cachedCollection;
                     }
-
-                    cachedCollection[key] = value;
                     lastConnectionCallbackData.set(subscriber.subscriptionID, {value: cachedCollection, matchedKey: subscriber.key});
                     subscriber.callback(cachedCollection, subscriber.key, {[key]: value});
                     continue;

lib/useOnyx.ts

@@ -70,17 +70,17 @@ function useOnyx<TKey extends OnyxKey, TReturnValue = OnyxValue<TKey>>(
             // Recompute if input changed, dependencies changed, or first time
             const dependenciesChanged = !shallowEqual(lastDependencies, currentDependencies);
             if (!hasComputed || lastInput !== input || dependenciesChanged) {
-                // Only proceed if we have a valid selector
-                if (selector) {
-                    const newOutput = selector(input);
-
-                    // Deep equality mode: only update if output actually changed
-                    if (!hasComputed || !deepEqual(lastOutput, newOutput) || dependenciesChanged) {
-                        lastInput = input;
-                        lastOutput = newOutput;
-                        lastDependencies = [...currentDependencies];
-                        hasComputed = true;
-                    }
+                const newOutput = selector(input);
+
+                // Always track the current input to avoid re-running the selector
+                // when the same input is seen again (even if the output didn't change).
+                lastInput = input;
+
+                // Only update the output reference if it actually changed
+                if (!hasComputed || !deepEqual(lastOutput, newOutput) || dependenciesChanged) {
+                    lastOutput = newOutput;
+                    lastDependencies = [...currentDependencies];
+                    hasComputed = true;
                 }
             }
 
@@ -218,18 +218,11 @@ function useOnyx<TKey extends OnyxKey, TReturnValue = OnyxValue<TKey>>(
             newFetchStatus = 'loading';
         }
 
-        // Optimized equality checking:
-        // - Memoized selectors already handle deep equality internally, so we can use fast reference equality
-        // - Non-selector cases use shallow equality for object reference checks
-        // - Normalize null to undefined to ensure consistent comparison (both represent "no value")
-        let areValuesEqual: boolean;
-        if (memoizedSelector) {
-            const normalizedPrevious = previousValueRef.current ?? undefined;
-            const normalizedNew = newValueRef.current ?? undefined;
-            areValuesEqual = normalizedPrevious === normalizedNew;
-        } else {
-            areValuesEqual = shallowEqual(previousValueRef.current ?? undefined, newValueRef.current);
-        }
+        // shallowEqual checks === first (O(1) for frozen snapshots and stable selector references),
+        // then falls back to comparing top-level properties for individual keys that may have
+        // new references with equivalent content.
+        // Normalize null to undefined to ensure consistent comparison (both represent "no value").
+        const areValuesEqual = shallowEqual(previousValueRef.current ?? undefined, newValueRef.current ?? undefined);
 
         // We update 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:

tests/unit/onyxUtilsTest.ts

@@ -349,6 +349,133 @@ describe('OnyxUtils', () => {
 
             await Onyx.disconnect(connection);
         });
+
+        it('should notify collection-level subscribers with waitForCollectionCallback', async () => {
+            const entryKey = `${ONYXKEYS.COLLECTION.TEST_KEY}789`;
+            const entryData = {value: 'data'};
+
+            const collectionCallback = jest.fn();
+            const connection = Onyx.connect({
+                key: ONYXKEYS.COLLECTION.TEST_KEY,
+                callback: collectionCallback,
+                waitForCollectionCallback: true,
+                initWithStoredValues: false,
+            });
+
+            await Onyx.set(entryKey, entryData);
+            collectionCallback.mockClear();
+
+            // Trigger keysChanged directly with a partial collection
+            OnyxUtils.keysChanged(ONYXKEYS.COLLECTION.TEST_KEY, {[entryKey]: entryData}, {});
+
+            expect(collectionCallback).toHaveBeenCalledTimes(1);
+            // Collection subscriber receives the full cached collection, subscriber.key, and partial
+            const [receivedCollection, receivedKey, receivedPartial] = collectionCallback.mock.calls[0];
+            expect(receivedKey).toBe(ONYXKEYS.COLLECTION.TEST_KEY);
+            expect(receivedCollection[entryKey]).toEqual(entryData);
+            expect(receivedPartial).toEqual({[entryKey]: entryData});
+
+            Onyx.disconnect(connection);
+        });
+
+        it('should skip notification when member value has same reference in previous and current collection', async () => {
+            const entryKey = `${ONYXKEYS.COLLECTION.TEST_KEY}same`;
+            const sameValue = {value: 'unchanged'};
+
+            await Onyx.set(entryKey, sameValue);
+
+            const callbackSpy = jest.fn();
+            const connection = Onyx.connect({
+                key: entryKey,
+                callback: callbackSpy,
+                initWithStoredValues: false,
+            });
+            await waitForPromisesToResolve();
+            callbackSpy.mockClear();
+
+            // Simulate keysChanged where the previous and current value are the SAME reference
+            // (which happens with frozen snapshots when nothing changed). === should skip notification.
+            OnyxUtils.keysChanged(ONYXKEYS.COLLECTION.TEST_KEY, {[entryKey]: sameValue}, {[entryKey]: sameValue});
+
+            expect(callbackSpy).not.toHaveBeenCalled();
+
+            Onyx.disconnect(connection);
+        });
+
+        it('should notify member subscribers only for changed keys in a batched update', async () => {
+            const keyA = `${ONYXKEYS.COLLECTION.TEST_KEY}A`;
+            const keyB = `${ONYXKEYS.COLLECTION.TEST_KEY}B`;
+            const keyC = `${ONYXKEYS.COLLECTION.TEST_KEY}C`;
+
+            const dataA = {value: 'A'};
+            const dataB = {value: 'B'};
+            const dataC = {value: 'C'};
+
+            await Onyx.multiSet({[keyA]: dataA, [keyB]: dataB, [keyC]: dataC});
+
+            const spyA = jest.fn();
+            const spyB = jest.fn();
+            const spyC = jest.fn();
+            const connA = Onyx.connect({key: keyA, callback: spyA, initWithStoredValues: false});
+            const connB = Onyx.connect({key: keyB, callback: spyB, initWithStoredValues: false});
+            const connC = Onyx.connect({key: keyC, callback: spyC, initWithStoredValues: false});
+            await waitForPromisesToResolve();
+            spyA.mockClear();
+            spyB.mockClear();
+            spyC.mockClear();
+
+            // Update cache so keysChanged reads the new values via getCachedCollection
+            const newA = {value: 'A-updated'};
+            const newC = {value: 'C-updated'};
+            OnyxCache.set(keyA, newA);
+            OnyxCache.set(keyC, newC);
+            // keyB stays the same reference
+
+            OnyxUtils.keysChanged(ONYXKEYS.COLLECTION.TEST_KEY, {[keyA]: newA, [keyB]: dataB, [keyC]: newC}, {[keyA]: dataA, [keyB]: dataB, [keyC]: dataC});
+
+            expect(spyA).toHaveBeenCalledTimes(1);
+            expect(spyB).not.toHaveBeenCalled();
+            expect(spyC).toHaveBeenCalledTimes(1);
+
+            Onyx.disconnect(connA);
+            Onyx.disconnect(connB);
+            Onyx.disconnect(connC);
+        });
+
+        it('should catch errors thrown by subscriber callbacks and continue notifying others', async () => {
+            const entryKey = `${ONYXKEYS.COLLECTION.TEST_KEY}errorTest`;
+            const entryData = {value: 'data'};
+
+            await Onyx.set(entryKey, entryData);
+
+            const failingCallback = jest.fn(() => {
+                throw new Error('subscriber failure');
+            });
+            const workingCallback = jest.fn();
+
+            const connFailing = Onyx.connect({key: entryKey, callback: failingCallback, initWithStoredValues: false});
+            const connWorking = Onyx.connect({key: entryKey, callback: workingCallback, initWithStoredValues: false});
+            await waitForPromisesToResolve();
+            failingCallback.mockClear();
+            workingCallback.mockClear();
+
+            // Spy on Logger to verify the error is logged
+            const logSpy = jest.spyOn(Logger, 'logAlert').mockImplementation(() => undefined);
+
+            const newData = {value: 'new'};
+            // Update the cache so keysChanged sees the new value as different from previous
+            OnyxCache.set(entryKey, newData);
+            OnyxUtils.keysChanged(ONYXKEYS.COLLECTION.TEST_KEY, {[entryKey]: newData}, {[entryKey]: entryData});
+
+            // Both callbacks should have been attempted; error should be logged
+            expect(failingCallback).toHaveBeenCalled();
+            expect(workingCallback).toHaveBeenCalled();
+            expect(logSpy).toHaveBeenCalled();
+
+            logSpy.mockRestore();
+            Onyx.disconnect(connFailing);
+            Onyx.disconnect(connWorking);
+        });
     });
 
     describe('mergeChanges', () => {