keyChange update key/connectionID for faster lookup

lib/Onyx.ts

@@ -102,6 +102,11 @@ function connect<TKey extends OnyxKey>(connectOptions: ConnectOptions<TKey>): nu
     callbackToStateMapping[connectionID] = mapping as Mapping<OnyxKey>;
     callbackToStateMapping[connectionID].connectionID = connectionID;
 
+    // When keyChanged is called, a key is passed and the method looks through all the Subscribers in callbackToStateMapping for the matching key to get the connectionID
+    // to avoid having to loop through all the Subscribers all the time (even when just one connection belongs to one key),
+    // We create a mapping from key to lists of connectionIDs to access the specific list of connectionIDs.
+    OnyxUtils.storeKeyByConnections(mapping.key, callbackToStateMapping[connectionID].connectionID);
+
     if (mapping.initWithStoredValues === false) {
         return connectionID;
     }
@@ -202,6 +207,8 @@ function disconnect(connectionID: number, keyToRemoveFromEvictionBlocklist?: Ony
         OnyxUtils.removeFromEvictionBlockList(keyToRemoveFromEvictionBlocklist, connectionID);
     }
 
+    OnyxUtils.deleteKeyByConnections(lastConnectionID);
+
     delete callbackToStateMapping[connectionID];
 }
 

lib/OnyxUtils.ts

@@ -51,6 +51,9 @@ const callbackToStateMapping: Record<string, Mapping<OnyxKey>> = {};
 // 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[] = [];
 
@@ -289,7 +292,7 @@ function isCollectionMemberKey<TCollectionKey extends CollectionKeyBase>(collect
  * @returns A tuple where the first element is the collection part and the second element is the ID part.
  */
 function splitCollectionMemberKey<TKey extends CollectionKey>(key: TKey): [TKey extends `${infer Prefix}_${string}` ? `${Prefix}_` : never, string] {
-    const underscoreIndex = key.indexOf('_');
+    const underscoreIndex = key.lastIndexOf('_');
 
     if (underscoreIndex === -1) {
         throw new Error(`Invalid ${key} key provided, only collection keys are allowed.`);
@@ -311,6 +314,33 @@ function isSafeEvictionKey(testKey: OnyxKey): boolean {
     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)) {
+        const updatedConnectionIDs = onyxKeyToConnectionIDs.get(subscriber.key).filter((id: number) => id !== connectionID);
+        onyxKeyToConnectionIDs.set(subscriber.key, updatedConnectionIDs);
+    }
+}
+
 /**
  * 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.
@@ -397,6 +427,26 @@ function addToEvictionBlockList(key: OnyxKey, connectionID: number): void {
 
     evictionBlocklist[key].push(connectionID);
 }
+/**
+ * It extracts the non-numeric collection identifier of a given key.
+ *
+ * For example:
+ * - `getCollectionKey("report_123")` would return "report_"
+ * - `getCollectionKey("report")` would return "report"
+ * - `getCollectionKey("report_")` would return "report_"
+ *
+ * @param {OnyxKey} key - The key to process.
+ * @return {string} The pure key without any numeric
+ */
+function getCollectionKey(key: OnyxKey): string {
+    const underscoreIndex = key.lastIndexOf('_');
+
+    if (underscoreIndex === -1) {
+        return key;
+    }
+
+    return key.substring(0, underscoreIndex + 1);
+}
 
 /**
  * Take all the keys that are safe to evict and add them to
@@ -462,6 +512,7 @@ function keysChanged<TKey extends CollectionKeyBase>(
     // 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) {
@@ -654,10 +705,24 @@ function keyChanged<TKey extends OnyxKey>(
         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.
+    // For performance reason, we look for the given key and later if don't find it we look for the collection key, instead of checking if it is a collection key first.
+    let stateMappingKeys = onyxKeyToConnectionIDs.get(key) ?? [];
+    const collectionKey = getCollectionKey(key);
+    const plainCollectionKey = collectionKey.lastIndexOf('_') !== -1 ? collectionKey : undefined;
+
+    if (plainCollectionKey) {
+        // Getting the collection key from the specific key because only collection keys were stored in the mapping.
+        stateMappingKeys = [...stateMappingKeys, ...(onyxKeyToConnectionIDs.get(plainCollectionKey) ?? [])];
+        if (stateMappingKeys.length === 0) {
+            return;
+        }
+    }
+
     for (let i = 0; i < stateMappingKeys.length; i++) {
         const subscriber = callbackToStateMapping[stateMappingKeys[i]];
         if (!subscriber || !isKeyMatch(subscriber.key, key) || !canUpdateSubscriber(subscriber)) {
@@ -1140,6 +1205,7 @@ const OnyxUtils = {
     keyChanged,
     sendDataToConnection,
     addKeyToRecentlyAccessedIfNeeded,
+    getCollectionKey,
     getCollectionDataAndSendAsObject,
     scheduleSubscriberUpdate,
     scheduleNotifyCollectionSubscribers,
@@ -1152,6 +1218,8 @@ const OnyxUtils = {
     prepareKeyValuePairsForStorage,
     applyMerge,
     initializeWithDefaultKeyStates,
+    storeKeyByConnections,
+    deleteKeyByConnections,
     getSnapshotKey,
 };
 

tests/unit/onyxUtilsTest.ts

@@ -0,0 +1,44 @@
+import OnyxUtils from '../../lib/OnyxUtils';
+
+describe('OnyxUtils', () => {
+    it('splitCollectionMemberKey should return correct values', () => {
+        const dataResult: Record<string, [string, string]> = {
+            test_: ['test_', ''],
+            test_level_: ['test_level_', ''],
+            test_level_1: ['test_level_', '1'],
+            test_level_2: ['test_level_', '2'],
+            test_level_last_3: ['test_level_last_', '3'],
+        };
+
+        Object.keys(dataResult).forEach((key) => {
+            const [collectionKey, id] = OnyxUtils.splitCollectionMemberKey(key);
+            expect(collectionKey).toEqual(dataResult[key][0]);
+            expect(id).toEqual(dataResult[key][1]);
+        });
+    });
+
+    it('splitCollectionMemberKey should throw error if key does not contain underscore', () => {
+        expect(() => {
+            OnyxUtils.splitCollectionMemberKey('test');
+        }).toThrowError('Invalid test key provided, only collection keys are allowed.');
+        expect(() => {
+            OnyxUtils.splitCollectionMemberKey('');
+        }).toThrowError('Invalid  key provided, only collection keys are allowed.');
+    });
+
+    it('getCollectionKey should return correct values', () => {
+        const dataResult: Record<string, string> = {
+            test: 'test',
+            test_: 'test_',
+            test_level_: 'test_level_',
+            test_level_1: 'test_level_',
+            test_level_2: 'test_level_',
+            test_level_last_3: 'test_level_last_',
+        };
+
+        Object.keys(dataResult).forEach((key) => {
+            const collectionKey = OnyxUtils.getCollectionKey(key);
+            expect(collectionKey).toEqual(dataResult[key]);
+        });
+    });
+});