Merge pull request Expensify#799 from callstack-internal/feature/onyx-store-pr-1

Remove waitForCollectionCallback from Onyx

Author: mountiny

API.md

@@ -80,7 +80,6 @@ This method will be deprecated soon. Please use `Onyx.connectWithoutView()` inst
 | connectOptions | The options object that will define the behavior of the connection. |
 | connectOptions.key | The Onyx key to subscribe to. |
 | connectOptions.callback | A function that will be called when the Onyx data we are subscribed changes. |
-| connectOptions.waitForCollectionCallback | If set to `true`, it will return the entire collection to the callback as a single object. |
 | connectOptions.selector | This will be used to subscribe to a subset of an Onyx key's data. **Only used inside `useOnyx()` hook.**        Using this setting on `useOnyx()` can have very positive performance benefits because the component will only re-render        when the subset of data changes. Otherwise, any change of data on any property would normally        cause the component to re-render (and that can be expensive from a performance standpoint). |
 
 **Example**  
@@ -103,7 +102,6 @@ Connects to an Onyx key given the options passed and listens to its changes.
 | connectOptions | The options object that will define the behavior of the connection. |
 | connectOptions.key | The Onyx key to subscribe to. |
 | connectOptions.callback | A function that will be called when the Onyx data we are subscribed changes. |
-| connectOptions.waitForCollectionCallback | If set to `true`, it will return the entire collection to the callback as a single object. |
 | connectOptions.selector | This will be used to subscribe to a subset of an Onyx key's data. **Only used inside `useOnyx()` hook.**        Using this setting on `useOnyx()` can have very positive performance benefits because the component will only re-render        when the subset of data changes. Otherwise, any change of data on any property would normally        cause the component to re-render (and that can be expensive from a performance standpoint). |
 
 **Example**  

README.md

@@ -257,20 +257,10 @@ export default MyComponent;
 This will add a prop to the component called `allReports` which is an object of collection member key/values. Changes to the individual member keys will modify the entire object and new props will be passed with each individual key update. The prop doesn't update on the initial rendering of the component until the entire collection has been read out of Onyx.
 
 ```js
-Onyx.connect({key: ONYXKEYS.COLLECTION.REPORT}, callback: (memberValue, memberKey) => {...});
+Onyx.connect({key: ONYXKEYS.COLLECTION.REPORT}, callback: (allReports, collectionKey, sourceValue) => {...});
 ```
 
-This will fire the callback once per member key depending on how many collection member keys are currently stored. Changes to those keys after the initial callbacks fire will occur when each individual key is updated.
-
-```js
-Onyx.connect({
-    key: ONYXKEYS.COLLECTION.REPORT,
-    waitForCollectionCallback: true,
-    callback: (allReports, collectionKey, sourceValue) => {...},
-});
-```
-
-This final option forces `Onyx.connect()` to behave more like `useOnyx()` and only update the callback once with the entire collection initially and later with an updated version of the collection when individual keys update. The `sourceValue` parameter contains only the specific keys and values that triggered the current update, which can be useful for optimizing your code to only process what changed. This parameter is not available when `waitForCollectionCallback` is false or the key is not a collection key.
+This will fire the callback once with the entire collection initially and later with an updated version of the collection when individual keys update. The `sourceValue` parameter contains only the specific keys and values that triggered the current update, which can be useful for optimizing your code to only process what changed.
 
 ### Performance Considerations When Using Collections
 

lib/Onyx.ts

@@ -94,7 +94,6 @@ function init({
  * @param connectOptions The options object that will define the behavior of the connection.
  * @param connectOptions.key The Onyx key to subscribe to.
  * @param connectOptions.callback A function that will be called when the Onyx data we are subscribed changes.
- * @param connectOptions.waitForCollectionCallback If set to `true`, it will return the entire collection to the callback as a single object.
  * @param connectOptions.selector This will be used to subscribe to a subset of an Onyx key's data. **Only used inside `useOnyx()` hook.**
  *        Using this setting on `useOnyx()` can have very positive performance benefits because the component will only re-render
  *        when the subset of data changes. Otherwise, any change of data on any property would normally
@@ -119,7 +118,6 @@ function connect<TKey extends OnyxKey>(connectOptions: ConnectOptions<TKey>): Co
  * @param connectOptions The options object that will define the behavior of the connection.
  * @param connectOptions.key The Onyx key to subscribe to.
  * @param connectOptions.callback A function that will be called when the Onyx data we are subscribed changes.
- * @param connectOptions.waitForCollectionCallback If set to `true`, it will return the entire collection to the callback as a single object.
  * @param connectOptions.selector This will be used to subscribe to a subset of an Onyx key's data. **Only used inside `useOnyx()` hook.**
  *        Using this setting on `useOnyx()` can have very positive performance benefits because the component will only re-render
  *        when the subset of data changes. Otherwise, any change of data on any property would normally

lib/OnyxConnectionManager.ts

@@ -4,7 +4,7 @@ import type {ConnectOptions} from './Onyx';
 import OnyxUtils from './OnyxUtils';
 import OnyxKeys from './OnyxKeys';
 import * as Str from './Str';
-import type {CollectionConnectCallback, DefaultConnectCallback, DefaultConnectOptions, OnyxKey, OnyxValue} from './types';
+import type {CollectionConnectCallback, DefaultConnectCallback, OnyxKey, OnyxValue} from './types';
 import onyxSnapshotCache from './OnyxSnapshotCache';
 
 type ConnectCallback = DefaultConnectCallback<OnyxKey> | CollectionConnectCallback<OnyxKey>;
@@ -48,11 +48,6 @@ type ConnectionMetadata = {
      * The value that triggered the last update
      */
     sourceValue?: OnyxValue<OnyxKey>;
-
-    /**
-     * Whether the subscriber is waiting for the collection callback to be fired.
-     */
-    waitForCollectionCallback?: boolean;
 };
 
 /**
@@ -115,21 +110,20 @@ class OnyxConnectionManager {
      * according to their purpose and effect they produce in the Onyx connection.
      */
     private generateConnectionID<TKey extends OnyxKey>(connectOptions: ConnectOptions<TKey>): string {
-        const {key, reuseConnection, waitForCollectionCallback} = connectOptions;
+        const {key, reuseConnection} = connectOptions;
 
         // The current session ID is appended to the connection ID so we can have different connections
         // after an `Onyx.clear()` operation.
         let suffix = `,sessionID=${this.sessionID}`;
 
-        // We will generate a unique ID in any of the following situations:
-        // - `reuseConnection` is `false`. That means the subscriber explicitly wants the connection to not be reused.
-        // - `key` is a collection key AND `waitForCollectionCallback` is `undefined/false`. This combination needs a new connection at every subscription
-        //   in order to send all the collection entries, so the connection can't be reused.
-        if (reuseConnection === false || (OnyxKeys.isCollectionKey(key) && (waitForCollectionCallback === undefined || waitForCollectionCallback === false))) {
+        // We will generate a unique ID when `reuseConnection` is `false`, which means the subscriber
+        // explicitly wants the connection to not be reused. Collection-root subscriptions are now always
+        // snapshot mode, so they can be reused like any other connection.
+        if (reuseConnection === false) {
             suffix += `,uniqueID=${Str.guid()}`;
         }
 
-        return `onyxKey=${key},waitForCollectionCallback=${waitForCollectionCallback ?? false}${suffix}`;
+        return `onyxKey=${key}${suffix}`;
     }
 
     /**
@@ -142,10 +136,18 @@ class OnyxConnectionManager {
         }
 
         for (const callback of connection.callbacks.values()) {
-            if (connection.waitForCollectionCallback) {
-                (callback as CollectionConnectCallback<OnyxKey>)(connection.cachedCallbackValue as Record<string, unknown>, connection.cachedCallbackKey as OnyxKey, connection.sourceValue);
-            } else {
-                (callback as DefaultConnectCallback<OnyxKey>)(connection.cachedCallbackValue, connection.cachedCallbackKey as OnyxKey);
+            try {
+                if (OnyxKeys.isCollectionKey(connection.onyxKey)) {
+                    (callback as CollectionConnectCallback<OnyxKey>)(
+                        connection.cachedCallbackValue as Record<string, unknown>,
+                        connection.cachedCallbackKey as OnyxKey,
+                        connection.sourceValue,
+                    );
+                } else {
+                    (callback as DefaultConnectCallback<OnyxKey>)(connection.cachedCallbackValue, connection.cachedCallbackKey as OnyxKey);
+                }
+            } catch (error) {
+                Logger.logAlert(`[ConnectionManager] Subscriber callback threw an error for key '${connection.onyxKey}': ${error}`);
             }
         }
     }
@@ -180,15 +182,14 @@ class OnyxConnectionManager {
 
             subscriptionID = OnyxUtils.subscribeToKey({
                 ...connectOptions,
-                callback: callback as DefaultConnectCallback<TKey>,
-            });
+                callback,
+            } as ConnectOptions<TKey>);
 
             connectionMetadata = {
                 subscriptionID,
                 onyxKey: connectOptions.key,
                 isConnectionMade: false,
                 callbacks: new Map(),
-                waitForCollectionCallback: connectOptions.waitForCollectionCallback,
             };
 
             this.connectionsMap.set(connectionID, connectionMetadata);
@@ -204,7 +205,7 @@ class OnyxConnectionManager {
             // Defer the callback execution to the next tick of the event loop.
             // This ensures that the current execution flow completes and the result connection object is available when the callback fires.
             Promise.resolve().then(() => {
-                (connectOptions as DefaultConnectOptions<OnyxKey>).callback?.(connectionMetadata.cachedCallbackValue, connectionMetadata.cachedCallbackKey as OnyxKey);
+                (connectOptions.callback as DefaultConnectCallback<OnyxKey> | undefined)?.(connectionMetadata.cachedCallbackValue, connectionMetadata.cachedCallbackKey as OnyxKey);
             });
         }
 

lib/OnyxUtils.ts

@@ -10,11 +10,11 @@ import StorageCircuitBreaker from './StorageCircuitBreaker';
 import Storage from './storage';
 import {StorageErrorClass} from './storage/errors';
 import type {
+    CollectionConnectCallback,
     CollectionKeyBase,
     ConnectOptions,
     DeepRecord,
     DefaultConnectCallback,
-    DefaultConnectOptions,
     KeyValueMapping,
     CallbackToStateMapping,
     MultiMergeReplaceNullPatches,
@@ -580,31 +580,8 @@ function keysChanged<TKey extends CollectionKeyBase>(
         }
 
         try {
-            lastConnectionCallbackData.set(subscriber.subscriptionID, {
-                value: cachedCollection,
-                matchedKey: subscriber.key,
-            });
-
-            if (subscriber.waitForCollectionCallback) {
-                subscriber.callback(cachedCollection, subscriber.key, partialCollection);
-                continue;
-            }
-
-            // Not using waitForCollectionCallback — notify per changed key.
-            // Re-check the subscription on each iteration because the callback may
-            // synchronously disconnect itself (removing it from callbackToStateMapping),
-            // in which case we must stop firing further callbacks for this subscriber.
-            for (const dataKey of changedMemberKeys) {
-                const currentSubscriber = callbackToStateMapping[subID];
-                if (!currentSubscriber || typeof currentSubscriber.callback !== 'function') {
-                    break;
-                }
-                if (cachedCollection[dataKey] === previousCollection[dataKey]) {
-                    continue;
-                }
-                const currentSubscriberCallback = currentSubscriber.callback as DefaultConnectCallback<TKey>;
-                currentSubscriberCallback(cachedCollection[dataKey], dataKey as TKey);
-            }
+            lastConnectionCallbackData.set(subscriber.subscriptionID, {value: cachedCollection, matchedKey: subscriber.key});
+            subscriber.callback(cachedCollection, subscriber.key, partialCollection);
         } catch (error) {
             Logger.logAlert(`[OnyxUtils.keysChanged] Subscriber callback threw an error for key '${collectionKey}': ${error}`);
         }
@@ -685,9 +662,9 @@ function keyChanged<TKey extends OnyxKey>(
                     continue;
                 }
 
-                if (OnyxKeys.isCollectionKey(subscriber.key) && subscriber.waitForCollectionCallback) {
-                    // Skip individual key changes for collection callbacks during collection updates
-                    // to prevent duplicate callbacks - the collection update will handle this properly
+                if (OnyxKeys.isCollectionKey(subscriber.key)) {
+                    // Skip individual key changes during collection updates to prevent duplicate
+                    // callbacks - the collection update will handle this properly.
                     if (isProcessingCollectionUpdate) {
                         continue;
                     }
@@ -698,13 +675,8 @@ function keyChanged<TKey extends OnyxKey>(
                         cachedCollection = getCachedCollection(subscriber.key);
                         cachedCollections[subscriber.key] = cachedCollection;
                     }
-                    lastConnectionCallbackData.set(subscriber.subscriptionID, {
-                        value: cachedCollection,
-                        matchedKey: subscriber.key,
-                    });
-                    subscriber.callback(cachedCollection, subscriber.key, {
-                        [key]: value,
-                    });
+                    lastConnectionCallbackData.set(subscriber.subscriptionID, {value: cachedCollection, matchedKey: subscriber.key});
+                    (subscriber.callback as CollectionConnectCallback<OnyxKey>)(cachedCollection, subscriber.key, {[key]: value});
                     continue;
                 }
 
@@ -738,10 +710,10 @@ function sendDataToConnection<TKey extends OnyxKey>(mapping: CallbackToStateMapp
     }
 
     // Always read the latest value from cache to avoid stale or duplicate data.
-    // For collection subscribers with waitForCollectionCallback, read the full collection.
+    // For collection-root subscribers, read the full collection.
     // For individual key subscribers, read just that key's value.
     let value: OnyxValue<TKey> | undefined;
-    if (OnyxKeys.isCollectionKey(mapping.key) && mapping.waitForCollectionCallback) {
+    if (OnyxKeys.isCollectionKey(mapping.key)) {
         const collection = getCachedCollection(mapping.key);
         value = Object.keys(collection).length > 0 ? (collection as OnyxValue<TKey>) : undefined;
     } else {
@@ -759,7 +731,7 @@ function sendDataToConnection<TKey extends OnyxKey>(mapping: CallbackToStateMapp
         return;
     }
 
-    (mapping as DefaultConnectOptions<TKey>).callback?.(value, matchedKey as TKey);
+    (mapping.callback as DefaultConnectCallback<TKey> | undefined)?.(value, matchedKey as TKey);
 }
 
 /**
@@ -1180,30 +1152,19 @@ function subscribeToKey<TKey extends OnyxKey>(connectOptions: ConnectOptions<TKe
                     cache.addNullishStorageKey(mapping.key);
                 }
 
-                const matchedKey = OnyxKeys.isCollectionKey(mapping.key) && mapping.waitForCollectionCallback ? mapping.key : undefined;
+                const matchedKey = OnyxKeys.isCollectionKey(mapping.key) ? mapping.key : undefined;
 
                 // Here we cannot use batching because the nullish value is expected to be set immediately for default props
                 // or they will be undefined.
                 sendDataToConnection(mapping, matchedKey);
                 return;
             }
 
-            // When using a callback subscriber we will either trigger the provided callback for each key we find or combine all values
-            // into an object and just make a single call. The latter behavior is enabled by providing a waitForCollectionCallback key
-            // combined with a subscription to a collection key.
+            // When using a callback subscriber, a subscription to a collection key combines all matching
+            // member values into a single object and makes one call with the whole collection object.
             if (typeof mapping.callback === 'function') {
                 if (OnyxKeys.isCollectionKey(mapping.key)) {
-                    if (mapping.waitForCollectionCallback) {
-                        getCollectionDataAndSendAsObject(matchingKeys, mapping);
-                        return;
-                    }
-
-                    // We did not opt into using waitForCollectionCallback mode so the callback is called for every matching key.
-                    multiGet(matchingKeys).then(() => {
-                        for (const key of matchingKeys) {
-                            sendDataToConnection(mapping, key as TKey);
-                        }
-                    });
+                    getCollectionDataAndSendAsObject(matchingKeys, mapping);
                     return;
                 }
 
@@ -1463,7 +1424,7 @@ function multiSetWithRetry(data: OnyxMultiSetInput, retryAttempt?: number): Prom
             // and subscriber state, matching the original per-key notification semantics.
             cache.set(key, value);
             // Skip subscriber notification on retry — already notified on attempt 0.
-            // waitForCollectionCallback subscribers re-fire on every keyChanged by contract.
+            // Collection-root subscribers re-fire on every keyChanged by contract.
             if (!retryAttempt) {
                 keyChanged(key, value);
             }
@@ -1554,7 +1515,7 @@ function setCollectionWithRetry<TKey extends CollectionKeyBase>({collectionKey,
         for (const [key, value] of keyValuePairs) cache.set(key, value);
 
         // Skip subscriber notification on retry — already notified on attempt 0.
-        // waitForCollectionCallback subscribers re-fire on every keysChanged by contract.
+        // Collection-root subscribers re-fire on every keysChanged by contract.
         if (!retryAttempt) {
             keysChanged(collectionKey, mutableCollection, previousCollection);
         }
@@ -1704,7 +1665,7 @@ function mergeCollectionWithPatches<TKey extends CollectionKeyBase>(
                 const previousCollection = getCachedCollection(collectionKey, existingKeys);
                 cache.merge(finalMergedCollection);
                 // Skip subscriber notification on retry — already notified on attempt 0.
-                // waitForCollectionCallback subscribers re-fire on every keysChanged by contract.
+                // Collection-root subscribers re-fire on every keysChanged by contract.
                 if (!retryAttempt) {
                     keysChanged(collectionKey, finalMergedCollection, previousCollection);
                 }
@@ -1797,7 +1758,7 @@ function partialSetCollection<TKey extends CollectionKeyBase>({collectionKey, co
         for (const [key, value] of keyValuePairs) cache.set(key, value);
 
         // Skip subscriber notification on retry — already notified on attempt 0.
-        // waitForCollectionCallback subscribers re-fire on every keysChanged by contract.
+        // Collection-root subscribers re-fire on every keysChanged by contract.
         if (!retryAttempt) {
             keysChanged(collectionKey, mutableCollection, previousCollection);
         }