diff --git a/lib/OnyxUtils.ts b/lib/OnyxUtils.ts index 669eda86..951b5008 100644 --- a/lib/OnyxUtils.ts +++ b/lib/OnyxUtils.ts @@ -582,6 +582,96 @@ function tryGetCachedValue(key: TKey): OnyxValue return val; } +/** + * Marks items that existed in previousCollection but not in preservedCollection as null. + * This ensures subscribers are properly notified about item removals. + * @param preservedCollection - The collection to mark removed items in (mutated in place) + * @param previousCollection - The previous collection state to compare against + */ +function markRemovedItemsAsNull(preservedCollection: OnyxInputKeyValueMapping, previousCollection: OnyxCollection) { + if (!previousCollection) { + return preservedCollection; + } + + const mutablePreservedCollection = {...preservedCollection}; + Object.keys(previousCollection).forEach((key) => { + if (key in preservedCollection) { + return; + } + + mutablePreservedCollection[key] = null; + }); + + return mutablePreservedCollection; +} + +/** + * Utility function to preserve object references for unchanged items in collection operations. + * Compares new values with cached values using deep equality and preserves references when data is identical. + * @param keyValuePairs - Array of key-value pairs to process + * @param previousCollection - Optional previous collection state. If provided, removed items will be included as null + * @returns The preserved collection with unchanged references maintained and removed items marked as null + */ +function preserveCollectionReferences(keyValuePairs: StorageKeyValuePair[], previousCollection?: OnyxCollection): OnyxInputKeyValueMapping { + const preservedCollection: OnyxInputKeyValueMapping = {}; + + keyValuePairs.forEach(([key, value]) => { + const cachedValue = cache.get(key, false); + + // If no cached value exists, we need to add the new value (skip expensive deep equality check) + // Use deep equality check to preserve references for unchanged items + if (cachedValue !== undefined && deepEqual(value, cachedValue)) { + // Keep the existing reference + preservedCollection[key] = cachedValue; + } else { + // Update cache only for changed items + cache.set(key, value); + preservedCollection[key] = value; + } + }); + + if (previousCollection) { + return markRemovedItemsAsNull(preservedCollection, previousCollection); + } + + return preservedCollection; +} + +/** + * Utility function for merge operations that preserves references after cache merge has been performed. + * Compares merged values with original cached values and preserves references when data is unchanged. + * @param collection - Collection of merged data + * @param originalCachedValues - Original cached values before merge + * @param previousCollection - Optional previous collection state. If provided, removed items will be included as null + * @returns The preserved collection with unchanged references maintained and removed items marked as null + */ +function preserveCollectionReferencesAfterMerge( + collection: Record>, + originalCachedValues: Record>, + previousCollection?: Record>, +): Record> { + const preservedCollection: Record> = {}; + + Object.keys(collection).forEach((key) => { + const newMergedValue = cache.get(key, false); + const originalValue = originalCachedValues[key]; + + // Use deep equality check to preserve references for unchanged items + if (originalValue !== undefined && deepEqual(newMergedValue, originalValue)) { + // Keep the existing reference for subscribers + preservedCollection[key] = originalValue; + } else { + preservedCollection[key] = newMergedValue; + } + }); + + if (previousCollection) { + return markRemovedItemsAsNull(preservedCollection, previousCollection); + } + + return preservedCollection; +} + function getCachedCollection(collectionKey: TKey, collectionMemberKeys?: string[]): NonNullable> { // Use optimized collection data retrieval when cache is populated const collectionData = cache.getCollectionData(collectionKey); @@ -1469,9 +1559,10 @@ function setCollectionWithRetry({collectionKey, const keyValuePairs = OnyxUtils.prepareKeyValuePairsForStorage(mutableCollection, true, undefined, true); const previousCollection = OnyxUtils.getCachedCollection(collectionKey); - keyValuePairs.forEach(([key, value]) => cache.set(key, value)); + // Preserve references for unchanged items and include removed items as null in setCollection + const preservedCollection = preserveCollectionReferences(keyValuePairs, previousCollection); - const updatePromise = OnyxUtils.scheduleNotifyCollectionSubscribers(collectionKey, mutableCollection, previousCollection); + const updatePromise = OnyxUtils.scheduleNotifyCollectionSubscribers(collectionKey, preservedCollection, previousCollection); return Storage.multiSet(keyValuePairs) .catch((error) => OnyxUtils.retryOperation(error, setCollectionWithRetry, {collectionKey, collection}, retryAttempt)) @@ -1532,6 +1623,9 @@ function mergeCollectionWithPatches( return getAllKeys() .then((persistedKeys) => { + // Capture keys that will be removed (before calling remove()) + const keysToRemove = resultCollectionKeys.filter((key) => resultCollection[key] === null && persistedKeys.has(key)); + // Split to keys that exist in storage and keys that don't const keys = resultCollectionKeys.filter((key) => { if (resultCollection[key] === null) { @@ -1543,6 +1637,8 @@ function mergeCollectionWithPatches( const existingKeys = keys.filter((key) => persistedKeys.has(key)); + // Get previous values for both existing keys and keys that will be removed + const allAffectedKeys = [...existingKeys, ...keysToRemove]; const cachedCollectionForExistingKeys = getCachedCollection(collectionKey, existingKeys); const existingKeyCollection = existingKeys.reduce((obj: OnyxInputKeyValueMapping, key) => { @@ -1579,7 +1675,8 @@ function mergeCollectionWithPatches( // We need to get the previously existing values so we can compare the new ones // against them, to avoid unnecessary subscriber updates. - const previousCollectionPromise = Promise.all(existingKeys.map((key) => get(key).then((value) => [key, value]))).then(Object.fromEntries); + // Include keys that will be removed so subscribers are notified about removals + const previousCollectionPromise = Promise.all(allAffectedKeys.map((key) => get(key).then((value) => [key, value]))).then(Object.fromEntries); // New keys will be added via multiSet while existing keys will be updated using multiMerge // This is because setting a key that doesn't exist yet with multiMerge will throw errors @@ -1597,8 +1694,19 @@ function mergeCollectionWithPatches( // Prefill cache if necessary by calling get() on any existing keys and then merge original data to cache // and update all subscribers const promiseUpdate = previousCollectionPromise.then((previousCollection) => { + // Capture the original cached values before merging + const originalCachedValues: Record> = {}; + Object.keys(finalMergedCollection).forEach((key) => { + originalCachedValues[key] = cache.get(key, false); + }); + + // Then merge all the data into cache as normal + cache.merge(finalMergedCollection); - return scheduleNotifyCollectionSubscribers(collectionKey, finalMergedCollection, previousCollection); + + // Finally, preserve references for items that didn't actually change and include removed items as null + const preservedCollection = preserveCollectionReferencesAfterMerge(finalMergedCollection, originalCachedValues, previousCollection); + return scheduleNotifyCollectionSubscribers(collectionKey, preservedCollection, previousCollection); }); return Promise.all(promises) @@ -1662,9 +1770,10 @@ function partialSetCollection({collectionKey, co const previousCollection = getCachedCollection(collectionKey, existingKeys); const keyValuePairs = prepareKeyValuePairsForStorage(mutableCollection, true, undefined, true); - keyValuePairs.forEach(([key, value]) => cache.set(key, value)); + // Preserve references for unchanged items and include removed items as null in partialSetCollection + const preservedCollection = preserveCollectionReferences(keyValuePairs, previousCollection); - const updatePromise = scheduleNotifyCollectionSubscribers(collectionKey, mutableCollection, previousCollection); + const updatePromise = scheduleNotifyCollectionSubscribers(collectionKey, preservedCollection, previousCollection); return Storage.multiSet(keyValuePairs) .catch((error) => retryOperation(error, partialSetCollection, {collectionKey, collection}, retryAttempt)) diff --git a/tests/unit/onyxUtilsTest.ts b/tests/unit/onyxUtilsTest.ts index a390ed79..2744f6e5 100644 --- a/tests/unit/onyxUtilsTest.ts +++ b/tests/unit/onyxUtilsTest.ts @@ -228,6 +228,157 @@ describe('OnyxUtils', () => { await Onyx.disconnect(connection); }); + + it('should notify subscribers when removing collection items with null values', async () => { + let collectionResult: OnyxCollection; + let callbackCount = 0; + const routeA = `${ONYXKEYS.COLLECTION.ROUTES}A`; + const routeB = `${ONYXKEYS.COLLECTION.ROUTES}B`; + const routeC = `${ONYXKEYS.COLLECTION.ROUTES}C`; + + // Test with waitForCollectionCallback: true + const connection = Onyx.connect({ + key: ONYXKEYS.COLLECTION.ROUTES, + initWithStoredValues: false, + callback: (value) => { + collectionResult = value; + callbackCount++; + }, + waitForCollectionCallback: true, + }); + + // Set initial collection state + await Onyx.setCollection(ONYXKEYS.COLLECTION.ROUTES, { + [routeA]: {name: 'Route A'}, + [routeB]: {name: 'Route B'}, + [routeC]: {name: 'Route C'}, + } as GenericCollection); + + callbackCount = 0; // Reset counter after initial set + + // Remove items by setting them to null + await OnyxUtils.partialSetCollection({ + collectionKey: ONYXKEYS.COLLECTION.ROUTES, + collection: { + [routeA]: null, + [routeB]: null, + } as GenericCollection, + }); + + // Should be notified about the removal + expect(callbackCount).toBeGreaterThan(0); + expect(collectionResult).toEqual({ + [routeC]: {name: 'Route C'}, + }); + + await Onyx.disconnect(connection); + }); + + it('should notify individual key subscribers when removing collection items with null values', async () => { + let routeAValue: unknown; + let routeBValue: unknown; + let routeCValue: unknown; + const routeA = `${ONYXKEYS.COLLECTION.ROUTES}A`; + const routeB = `${ONYXKEYS.COLLECTION.ROUTES}B`; + const routeC = `${ONYXKEYS.COLLECTION.ROUTES}C`; + + // Test with waitForCollectionCallback: false (individual callbacks) + const connectionA = Onyx.connect({ + key: routeA, + initWithStoredValues: false, + callback: (value) => { + routeAValue = value; + }, + }); + + const connectionB = Onyx.connect({ + key: routeB, + initWithStoredValues: false, + callback: (value) => { + routeBValue = value; + }, + }); + + const connectionC = Onyx.connect({ + key: routeC, + initWithStoredValues: false, + callback: (value) => { + routeCValue = value; + }, + }); + + // Set initial collection state + await Onyx.setCollection(ONYXKEYS.COLLECTION.ROUTES, { + [routeA]: {name: 'Route A'}, + [routeB]: {name: 'Route B'}, + [routeC]: {name: 'Route C'}, + } as GenericCollection); + + // Remove items by setting them to null + await OnyxUtils.partialSetCollection({ + collectionKey: ONYXKEYS.COLLECTION.ROUTES, + collection: { + [routeA]: null, + [routeB]: null, + } as GenericCollection, + }); + + // Individual subscribers should be notified about removals + expect(routeAValue).toBeUndefined(); + expect(routeBValue).toBeUndefined(); + expect(routeCValue).toEqual({name: 'Route C'}); + + await Onyx.disconnect(connectionA); + await Onyx.disconnect(connectionB); + await Onyx.disconnect(connectionC); + }); + + it('should notify collection subscribers without waitForCollectionCallback when removing items', async () => { + const callbackResults: Array<{value: unknown; key: string}> = []; + const routeA = `${ONYXKEYS.COLLECTION.ROUTES}A`; + const routeB = `${ONYXKEYS.COLLECTION.ROUTES}B`; + const routeC = `${ONYXKEYS.COLLECTION.ROUTES}C`; + + // Test without waitForCollectionCallback (gets called for each item) + const connection = Onyx.connect({ + key: ONYXKEYS.COLLECTION.ROUTES, + initWithStoredValues: false, + callback: (value, key) => { + callbackResults.push({value, key}); + }, + waitForCollectionCallback: false, + }); + + // Set initial collection state + await Onyx.setCollection(ONYXKEYS.COLLECTION.ROUTES, { + [routeA]: {name: 'Route A'}, + [routeB]: {name: 'Route B'}, + [routeC]: {name: 'Route C'}, + } as GenericCollection); + + callbackResults.length = 0; // Clear after initial set + + // Remove items by setting them to null + await OnyxUtils.partialSetCollection({ + collectionKey: ONYXKEYS.COLLECTION.ROUTES, + collection: { + [routeA]: null, + [routeB]: null, + } as GenericCollection, + }); + + // Should be notified about removals (value will be undefined for removed keys) + const removalCallbacks = callbackResults.filter((result) => result.value === undefined); + + // We expect at least routeA or routeB to have undefined callback + // Both should ideally be notified, but the exact behavior depends on iteration order + expect(removalCallbacks.length).toBeGreaterThanOrEqual(1); + const hasRouteARemoval = removalCallbacks.some((result) => result.key === routeA); + const hasRouteBRemoval = removalCallbacks.some((result) => result.key === routeB); + expect(hasRouteARemoval || hasRouteBRemoval).toBe(true); + + await Onyx.disconnect(connection); + }); }); describe('keysChanged', () => {