keyChange update key/connectionID for faster lookup

lib/Onyx.ts

@@ -99,6 +99,11 @@ function connect<TKey extends OnyxKey>(connectOptions: ConnectOptions<TKey>): nu
     callbackToStateMapping[connectionID] = mapping as Mapping<OnyxKey>;
     callbackToStateMapping[connectionID].connectionID = connectionID;
 
+    // When keyChanged is called, the key is passed and looks through all the callbackToStateMapping the right key
+    // to avoid having to loop through all the keys all the time (even when just one connection belongs to one key),
+    // we store the key by connectionID, to access the specific list of connectionIDs
+    OnyxUtils.storeKeyByConnections(mapping.key, callbackToStateMapping[connectionID].connectionID);
+
     if (mapping.initWithStoredValues === false) {
         return connectionID;
     }
@@ -199,6 +204,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 - a connection ID of the subscriber
+ * @param key - a 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.
@@ -381,6 +410,24 @@
 
     evictionBlocklist[key].push(connectionID);
 }
+/**
+ * It extracts the pure, non-numeric part of a given key.
+ *
+ * For example:
+ * - `getPureKey("report_123")` would return "report_"
+ * - `getPureKey("report")` would return "report"
+ * - `getPureKey("report_")` would return "report_"
+ * - `getPureKey(null)` would return ""
+ *
+ * @param {OnyxKey} key - The key to process.
+ * @return {string} The pure key without any numeric
+ */
+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
@@ -442,6 +489,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) {
@@ -626,10 +674,21 @@
         removeLastAccessedKey(key);
     }
 
-    // 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
+    // We get the subscribers 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);
+    // Given the amount of times this function is called we need to make sure we are not iterating over all subscribers every time. On the other hand, we don't need to
+    // do the same in keysChanged, because we only call that function when a collection key changes, and it doesn't happen that often.
+    let stateMappingKeys = onyxKeyToConnectionIDs.get(key);
+
+    if (!stateMappingKeys) {
+        // Getting the collection key from the specific key because 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 +1167,8 @@
     prepareKeyValuePairsForStorage,
     applyMerge,
     initializeWithDefaultKeyStates,
+    storeKeyByConnections,
+    deleteKeyByConnections,
 };
 
 export default OnyxUtils;