[Structural Sharing] PR 5: Batched multiSetWithRetry notifications

lib/OnyxUtils.ts

@@ -1356,15 +1356,49 @@ function multiSetWithRetry(data: OnyxMultiSetInput, retryAttempt?: number): Prom
 
     const keyValuePairsToSet = OnyxUtils.prepareKeyValuePairsForStorage(newData, true);
 
+    // Group collection members by their parent collection key so each collection can be notified
+    // via a single batched keysChanged() call instead of one keyChanged() per member. For each
+    // collection, `partial` holds the new values being set and `previous` holds the cached values
+    // from before the set, which keysChanged() uses to skip subscribers whose value didn't change.
+    const collectionBatches = new Map<string, {partial: Record<string, OnyxValue<OnyxKey>>; previous: Record<string, OnyxValue<OnyxKey>>}>();
+    const nonCollectionPairs: Array<[string, OnyxValue<OnyxKey>]> = [];
+
     for (const [key, value] of keyValuePairsToSet) {
         // When we use multiSet to set a key we want to clear the current delta changes from Onyx.merge that were queued
         // before the value was set. If Onyx.merge is currently reading the old value from storage, it will then not apply the changes.
         if (OnyxUtils.hasPendingMergeForKey(key)) {
             delete OnyxUtils.getMergeQueue()[key];
         }
 
-        // Update cache and optimistically inform subscribers
-        cache.set(key, value);
+        const collectionKey = OnyxKeys.getCollectionKey(key);
+        if (collectionKey && OnyxKeys.isCollectionMemberKey(collectionKey, key)) {
+            // Capture the previous cached value BEFORE calling cache.set() so keysChanged()
+            // can diff old vs new per-member.
+            const previousValue = cache.get(key);
+            cache.set(key, value);
+
+            let batch = collectionBatches.get(collectionKey);
+            if (!batch) {
+                batch = {partial: {}, previous: {}};
+                collectionBatches.set(collectionKey, batch);
+            }
+            batch.partial[key] = value;
+            batch.previous[key] = previousValue;
+        } else {
+            // Non-collection keys are notified individually — no batching.
+            cache.set(key, value);
+            nonCollectionPairs.push([key, value]);
+        }
+    }
+
+    // One keysChanged() per collection — fires each collection-level subscriber once and lets
+    // keysChanged() internally decide which individual member subscribers need notification.
+    for (const [collectionKey, batch] of collectionBatches) {
+        keysChanged(collectionKey as CollectionKeyBase, batch.partial, batch.previous);
+    }
+
+    // Non-collection keys go through the regular per-key notification path.
+    for (const [key, value] of nonCollectionPairs) {
         keyChanged(key, value);
     }
 

tests/unit/onyxUtilsTest.ts

@@ -266,6 +266,175 @@ describe('OnyxUtils', () => {
         });
     });
 
+    describe('multiSetWithRetry', () => {
+        it('should fire collection-level callback only once per collection even with multiple members', async () => {
+            const collectionCallback = jest.fn();
+            const connection = Onyx.connect({
+                key: ONYXKEYS.COLLECTION.TEST_KEY,
+                callback: collectionCallback,
+                waitForCollectionCallback: true,
+                initWithStoredValues: false,
+            });
+
+            await waitForPromisesToResolve();
+            collectionCallback.mockClear();
+
+            // multiSet with 3 members of the same collection
+            await Onyx.multiSet({
+                [`${ONYXKEYS.COLLECTION.TEST_KEY}1`]: {id: 1},
+                [`${ONYXKEYS.COLLECTION.TEST_KEY}2`]: {id: 2},
+                [`${ONYXKEYS.COLLECTION.TEST_KEY}3`]: {id: 3},
+            });
+
+            // Should be called only ONCE with the batched collection (not 3 times)
+            expect(collectionCallback).toHaveBeenCalledTimes(1);
+            const [collection] = collectionCallback.mock.calls[0];
+            expect(collection[`${ONYXKEYS.COLLECTION.TEST_KEY}1`]).toEqual({id: 1});
+            expect(collection[`${ONYXKEYS.COLLECTION.TEST_KEY}2`]).toEqual({id: 2});
+            expect(collection[`${ONYXKEYS.COLLECTION.TEST_KEY}3`]).toEqual({id: 3});
+
+            Onyx.disconnect(connection);
+        });
+
+        it('should fire individual member-key subscribers once per key', async () => {
+            const spy1 = jest.fn();
+            const spy2 = jest.fn();
+            const spy3 = jest.fn();
+
+            const conn1 = Onyx.connect({key: `${ONYXKEYS.COLLECTION.TEST_KEY}1`, callback: spy1, initWithStoredValues: false});
+            const conn2 = Onyx.connect({key: `${ONYXKEYS.COLLECTION.TEST_KEY}2`, callback: spy2, initWithStoredValues: false});
+            const conn3 = Onyx.connect({key: `${ONYXKEYS.COLLECTION.TEST_KEY}3`, callback: spy3, initWithStoredValues: false});
+            await waitForPromisesToResolve();
+            spy1.mockClear();
+            spy2.mockClear();
+            spy3.mockClear();
+
+            await Onyx.multiSet({
+                [`${ONYXKEYS.COLLECTION.TEST_KEY}1`]: {id: 1},
+                [`${ONYXKEYS.COLLECTION.TEST_KEY}2`]: {id: 2},
+                [`${ONYXKEYS.COLLECTION.TEST_KEY}3`]: {id: 3},
+            });
+
+            expect(spy1).toHaveBeenCalledTimes(1);
+            expect(spy1).toHaveBeenCalledWith({id: 1}, `${ONYXKEYS.COLLECTION.TEST_KEY}1`);
+            expect(spy2).toHaveBeenCalledTimes(1);
+            expect(spy2).toHaveBeenCalledWith({id: 2}, `${ONYXKEYS.COLLECTION.TEST_KEY}2`);
+            expect(spy3).toHaveBeenCalledTimes(1);
+            expect(spy3).toHaveBeenCalledWith({id: 3}, `${ONYXKEYS.COLLECTION.TEST_KEY}3`);
+
+            Onyx.disconnect(conn1);
+            Onyx.disconnect(conn2);
+            Onyx.disconnect(conn3);
+        });
+
+        it('should notify non-collection keys individually alongside batched collection updates', async () => {
+            const collectionCallback = jest.fn();
+            const singleKeyCallback = jest.fn();
+
+            const connCollection = Onyx.connect({
+                key: ONYXKEYS.COLLECTION.TEST_KEY,
+                callback: collectionCallback,
+                waitForCollectionCallback: true,
+                initWithStoredValues: false,
+            });
+            const connSingle = Onyx.connect({
+                key: ONYXKEYS.TEST_KEY,
+                callback: singleKeyCallback,
+                initWithStoredValues: false,
+            });
+            await waitForPromisesToResolve();
+            collectionCallback.mockClear();
+            singleKeyCallback.mockClear();
+
+            // Mix of collection members and a non-collection key
+            await Onyx.multiSet({
+                [`${ONYXKEYS.COLLECTION.TEST_KEY}1`]: {id: 1},
+                [`${ONYXKEYS.COLLECTION.TEST_KEY}2`]: {id: 2},
+                [ONYXKEYS.TEST_KEY]: 'standalone',
+            });
+
+            // Collection callback fires once (batched)
+            expect(collectionCallback).toHaveBeenCalledTimes(1);
+            // Non-collection key callback fires once
+            expect(singleKeyCallback).toHaveBeenCalledTimes(1);
+            expect(singleKeyCallback).toHaveBeenCalledWith('standalone', ONYXKEYS.TEST_KEY);
+
+            Onyx.disconnect(connCollection);
+            Onyx.disconnect(connSingle);
+        });
+
+        it('should batch notifications per-collection when members span multiple collections', async () => {
+            const testCallback = jest.fn();
+            const routesCallback = jest.fn();
+
+            const connTest = Onyx.connect({
+                key: ONYXKEYS.COLLECTION.TEST_KEY,
+                callback: testCallback,
+                waitForCollectionCallback: true,
+                initWithStoredValues: false,
+            });
+            const connRoutes = Onyx.connect({
+                key: ONYXKEYS.COLLECTION.ROUTES,
+                callback: routesCallback,
+                waitForCollectionCallback: true,
+                initWithStoredValues: false,
+            });
+            await waitForPromisesToResolve();
+            testCallback.mockClear();
+            routesCallback.mockClear();
+
+            // multiSet with members of two different collections
+            await Onyx.multiSet({
+                [`${ONYXKEYS.COLLECTION.TEST_KEY}1`]: {id: 1},
+                [`${ONYXKEYS.COLLECTION.TEST_KEY}2`]: {id: 2},
+                [`${ONYXKEYS.COLLECTION.ROUTES}A`]: {name: 'A'},
+                [`${ONYXKEYS.COLLECTION.ROUTES}B`]: {name: 'B'},
+            });
+
+            // Each collection callback fires once
+            expect(testCallback).toHaveBeenCalledTimes(1);
+            expect(routesCallback).toHaveBeenCalledTimes(1);
+
+            Onyx.disconnect(connTest);
+            Onyx.disconnect(connRoutes);
+        });
+
+        it('should pass previous values to keysChanged so unchanged members skip notification', async () => {
+            // Set initial data
+            const initial1 = {id: 1, name: 'A'};
+            const initial2 = {id: 2, name: 'B'};
+            await Onyx.multiSet({
+                [`${ONYXKEYS.COLLECTION.TEST_KEY}1`]: initial1,
+                [`${ONYXKEYS.COLLECTION.TEST_KEY}2`]: initial2,
+            });
+
+            const spy1 = jest.fn();
+            const spy2 = jest.fn();
+            const conn1 = Onyx.connect({key: `${ONYXKEYS.COLLECTION.TEST_KEY}1`, callback: spy1, initWithStoredValues: false});
+            const conn2 = Onyx.connect({key: `${ONYXKEYS.COLLECTION.TEST_KEY}2`, callback: spy2, initWithStoredValues: false});
+            await waitForPromisesToResolve();
+            spy1.mockClear();
+            spy2.mockClear();
+
+            // multiSet: change key 1, keep key 2 with same content (but new reference)
+            await Onyx.multiSet({
+                [`${ONYXKEYS.COLLECTION.TEST_KEY}1`]: {id: 1, name: 'A-updated'},
+                [`${ONYXKEYS.COLLECTION.TEST_KEY}2`]: initial2,
+            });
+
+            // Key 1 subscriber fires (value changed)
+            expect(spy1).toHaveBeenCalledTimes(1);
+            expect(spy1).toHaveBeenCalledWith({id: 1, name: 'A-updated'}, `${ONYXKEYS.COLLECTION.TEST_KEY}1`);
+
+            // Key 2 keeps the same reference (passed as-is in multiSet) — subscriber should not fire
+            // because keysChanged sees the same reference as previousCollection[key]
+            expect(spy2).not.toHaveBeenCalled();
+
+            Onyx.disconnect(conn1);
+            Onyx.disconnect(conn2);
+        });
+    });
+
     describe('keysChanged', () => {
         beforeEach(() => {
             Onyx.clear();