keyChange update key/connectionID for faster lookup

lib/Onyx.ts

@@ -99,6 +99,8 @@ function connect<TKey extends OnyxKey>(connectOptions: ConnectOptions<TKey>): nu
     callbackToStateMapping[connectionID] = mapping as Mapping<OnyxKey>;
     callbackToStateMapping[connectionID].connectionID = connectionID;
 
+    OnyxUtils.storeKeyByConnections(mapping.key, callbackToStateMapping[connectionID].connectionID);
+
     if (mapping.initWithStoredValues === false) {
         return connectionID;
     }
@@ -199,6 +201,8 @@ function disconnect(connectionID: number, keyToRemoveFromEvictionBlocklist?: Ony
         OnyxUtils.removeFromEvictionBlockList(keyToRemoveFromEvictionBlocklist, connectionID);
     }
 
+    OnyxUtils.deleteKeyByConnections(lastConnectionID);
+
     delete callbackToStateMapping[connectionID];
 }
 

lib/OnyxUtils.ts

@@ -51,6 +51,9 @@
 // Keeps a copy of the values of the onyx collection keys as a map for faster lookups
 let onyxCollectionKeyMap = new Map<OnyxKey, OnyxValue<OnyxKey>>();
 
+// Holds a mapping of the connected key to the connectionID for faster lookups
+const onyxKeyToConnectionIDs = new Map();
+
 // Holds a list of keys that have been directly subscribed to or recently modified from least to most recent
 let recentlyAccessedKeys: OnyxKey[] = [];
 
@@ -295,6 +298,32 @@
     return evictionAllowList.some((key) => isKeyMatch(key, testKey));
 }
 
+/**
+ * Stores a connection ID associated with a given key.
+ *
+ * @param connectionID connectionID of the subscriber
+ * @param key key that the subscriber is connected to
+ */
+function storeKeyByConnections(key: OnyxKey, connectionID: number) {
+    if (!onyxKeyToConnectionIDs.has(key)) {
+        onyxKeyToConnectionIDs.set(key, []);
+    }
+    onyxKeyToConnectionIDs.get(key).push(connectionID);
+}
+
+/**
+ * Deletes a connection ID associated with its corresponding key.
+ *
+ * @param {number} connectionID - The connection ID to be deleted.
+ */
+function deleteKeyByConnections(connectionID: number) {
+    const subscriber = callbackToStateMapping[connectionID];
+
+    if (subscriber && onyxKeyToConnectionIDs.has(subscriber.key)) {
+        onyxKeyToConnectionIDs.set(subscriber.key, _.without(onyxKeyToConnectionIDs.get(subscriber.key), connectionID));
+    }
+}
+
 /**
  * Tries to get a value from the cache. If the value is not present in cache it will return the default value or undefined.
  * If the requested key is a collection, it will return an object with all the collection members.
@@ -382,6 +411,13 @@
     evictionBlocklist[key].push(connectionID);
 }
 
+function getPureKey(key: OnyxKey): string {
+    if (!key) {
+        return '';
+    }
+    return key.replace(/\d+/g, '');
+}
+
 /**
  * Take all the keys that are safe to evict and add them to
  * the recently accessed list when initializing the app. This
@@ -442,6 +478,7 @@
     // 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 (let i = 0; i < stateMappingKeys.length; i++) {
         const subscriber = callbackToStateMapping[stateMappingKeys[i]];
         if (!subscriber) {
@@ -629,7 +666,16 @@
     // We are iterating over all subscribers to see if they are interested in the key that has just changed. If the subscriber's  key is a collection key then we will
     // notify them if the key that changed is a collection member. Or if it is a regular key notify them when there is an exact match. Depending on whether the subscriber
     // was connected via withOnyx we will call setState() directly on the withOnyx instance. If it is a regular connection we will pass the data to the provided callback.
-    const stateMappingKeys = Object.keys(callbackToStateMapping);
+    let stateMappingKeys = onyxKeyToConnectionIDs.get(key);
+
+    if (!stateMappingKeys) {
+        // Getting the collection key from the specific key 'cos only collection keys were stored in the mapping.
+        stateMappingKeys = onyxKeyToConnectionIDs.get(getPureKey(key));
+        if (!stateMappingKeys) {
+            return;
+        }
+    }
+
     for (let i = 0; i < stateMappingKeys.length; i++) {
         const subscriber = callbackToStateMapping[stateMappingKeys[i]];
         if (!subscriber || !isKeyMatch(subscriber.key, key) || !canUpdateSubscriber(subscriber)) {
@@ -1108,6 +1154,8 @@
     prepareKeyValuePairsForStorage,
     applyMerge,
     initializeWithDefaultKeyStates,
+    storeKeyByConnections,
+    deleteKeyByConnections,
 };
 
 export default OnyxUtils;