Merge pull request #770 from callstack-internal/fix/improve-onyxConnectionManager

Author: mjasikowski

API-INTERNAL.md

@@ -145,6 +145,9 @@ that this internal function allows passing an additional <code>mergeReplaceNullP
 Any existing collection members not included in the new data will not be removed.
 Retries on failure.</p>
 </dd>
+<dt><a href="#getCallbackToStateMapping">getCallbackToStateMapping()</a></dt>
+<dd><p>Getter - returns the callback to state mapping, useful in test environments.</p>
+</dd>
 <dt><a href="#clearOnyxUtilsInternals">clearOnyxUtilsInternals()</a></dt>
 <dd><p>Clear internal variables used in this file, useful in test environments.</p>
 </dd>
@@ -519,6 +522,12 @@ Retries on failure.
 | params.collection | Object collection keyed by individual collection member keys and values |
 | retryAttempt | retry attempt |
 
+<a name="getCallbackToStateMapping"></a>
+
+## getCallbackToStateMapping()
+Getter - returns the callback to state mapping, useful in test environments.
+
+**Kind**: global function  
 <a name="clearOnyxUtilsInternals"></a>
 
 ## clearOnyxUtilsInternals()

lib/OnyxUtils.ts

@@ -1053,6 +1053,24 @@ function subscribeToKey<TKey extends OnyxKey>(connectOptions: ConnectOptions<TKe
     callbackToStateMapping[subscriptionID] = mapping as CallbackToStateMapping<OnyxKey>;
     callbackToStateMapping[subscriptionID].subscriptionID = subscriptionID;
 
+    // If the subscriber is attempting to connect to a collection member whose ID is skippable (e.g. "undefined", "null", etc.)
+    // we suppress wiring the subscription fully to avoid unnecessary callback emissions such as for "report_undefined".
+    // We still return a valid subscriptionID so callers can disconnect safely.
+    try {
+        const skippableIDs = getSkippableCollectionMemberIDs();
+        if (skippableIDs.size) {
+            const [, collectionMemberID] = OnyxKeys.splitCollectionMemberKey(mapping.key);
+            if (skippableIDs.has(collectionMemberID)) {
+                // Clean up the provisional mapping to avoid retaining unused subscribers.
+                cache.addNullishStorageKey(mapping.key);
+                delete callbackToStateMapping[subscriptionID];
+                return subscriptionID;
+            }
+        }
+    } catch (e) {
+        // Not a collection member key, proceed as usual.
+    }
+
     // 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 subscriptionID
     // 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 subscriptionIDs to access the specific list of subscriptionIDs.
@@ -1659,6 +1677,13 @@ function logKeyRemoved(onyxMethod: Extract<OnyxMethod, 'set' | 'merge'>, key: On
     Logger.logInfo(`${onyxMethod} called for key: ${key} => null passed, so key was removed`);
 }
 
+/**
+ * Getter - returns the callback to state mapping, useful in test environments.
+ */
+function getCallbackToStateMapping(): Record<number, CallbackToStateMapping<OnyxKey>> {
+    return callbackToStateMapping;
+}
+
 /**
  * Clear internal variables used in this file, useful in test environments.
  */
@@ -1718,6 +1743,7 @@ const OnyxUtils = {
     setWithRetry,
     multiSetWithRetry,
     setCollectionWithRetry,
+    getCallbackToStateMapping,
 };
 
 export type {OnyxMethod};

tests/unit/onyxUtilsTest.ts

@@ -95,6 +95,74 @@ describe('OnyxUtils', () => {
 
     afterEach(() => jest.clearAllMocks());
 
+    describe('skippable member subscriptions', () => {
+        const BASE = ONYXKEYS.COLLECTION.TEST_KEY;
+
+        beforeEach(() => {
+            // Enable skipping of undefined member IDs for these tests
+            OnyxUtils.setSkippableCollectionMemberIDs(new Set(['undefined']));
+        });
+
+        afterEach(() => {
+            // Restore to no skippable IDs to avoid affecting other tests
+            OnyxUtils.setSkippableCollectionMemberIDs(new Set());
+        });
+
+        it('does not emit initial callback for report_undefined member', async () => {
+            const key = `${BASE}undefined`;
+            const callback = jest.fn();
+            Onyx.connect({key, callback});
+
+            // Flush async subscription flow
+            await act(async () => waitForPromisesToResolve());
+
+            // No initial data should be sent for a skippable member
+            expect(callback).not.toHaveBeenCalled();
+        });
+
+        it('still emits for valid member keys', async () => {
+            const key = `${BASE}123`;
+            await Onyx.set(key, {id: 123});
+
+            const callback = jest.fn();
+            Onyx.connect({key, callback});
+            await act(async () => waitForPromisesToResolve());
+            expect(callback).toHaveBeenCalledTimes(1);
+            expect(callback).toHaveBeenCalledWith({id: 123}, key);
+        });
+
+        it('omits skippable members from base collection', async () => {
+            const undefinedKey = `${BASE}undefined`;
+            const validKey = `${BASE}1`;
+
+            await Onyx.set(undefinedKey, {bad: true});
+            await Onyx.set(validKey, {ok: true});
+
+            let received: Record<string, unknown> | undefined;
+            Onyx.connect({
+                key: BASE,
+                waitForCollectionCallback: true,
+                callback: (value) => {
+                    received = value as Record<string, unknown>;
+                },
+            });
+            await act(async () => waitForPromisesToResolve());
+            expect(received).toEqual({[validKey]: {ok: true}});
+            expect(Object.keys(received ?? {})).not.toContain(undefinedKey);
+        });
+
+        it('does not register an active subscription in callbackToStateMapping for a skippable member', async () => {
+            const skippableKey = `${BASE}undefined`;
+            Onyx.connect({key: skippableKey, callback: jest.fn()});
+
+            await act(async () => waitForPromisesToResolve());
+
+            const mappings = OnyxUtils.getCallbackToStateMapping();
+            const hasActiveSubscription = Object.values(mappings).some((m) => m.key === skippableKey);
+            expect(hasActiveSubscription).toBe(false);
+        });
+    });
+
     describe('partialSetCollection', () => {
         beforeEach(() => {
             Onyx.clear();

tests/unit/useOnyxTest.ts

@@ -1110,6 +1110,54 @@ describe('useOnyx', () => {
             expect(result.current[0]).toBeUndefined();
             expect(result.current[1].status).toEqual('loaded');
         });
+
+        it('should return undefined and loaded state when switching from a valid key to a skippable one', async () => {
+            await Onyx.set(`${ONYXKEYS.COLLECTION.TEST_KEY}1`, {id: '1'});
+            // Seed a value directly in storage for the skippable key.
+            // If the subscription is NOT skipped, Onyx would load this and return it.
+            // Asserting undefined below proves the subscription was actually suppressed.
+            await StorageMock.setItem(`${ONYXKEYS.COLLECTION.TEST_KEY}skippable-id`, {id: 'skippable'});
+
+            const {result, rerender} = renderHook((key: string) => useOnyx(key), {initialProps: `${ONYXKEYS.COLLECTION.TEST_KEY}1` as string});
+
+            await act(async () => waitForPromisesToResolve());
+
+            expect(result.current[0]).toEqual({id: '1'});
+            expect(result.current[1].status).toEqual('loaded');
+
+            await act(async () => {
+                rerender(`${ONYXKEYS.COLLECTION.TEST_KEY}skippable-id`);
+            });
+
+            await act(async () => waitForPromisesToResolve());
+
+            expect(result.current[0]).toBeUndefined();
+            expect(result.current[1].status).toEqual('loaded');
+        });
+
+        it('should transition through loading and return value when switching from a skippable key to a valid one', async () => {
+            // Seed a value for the skippable key — must stay invisible to the hook
+            await StorageMock.setItem(`${ONYXKEYS.COLLECTION.TEST_KEY}skippable-id`, {id: 'skippable'});
+            // Seed the target valid key in storage only (not in cache) so the switch goes through loading
+            await StorageMock.setItem(`${ONYXKEYS.COLLECTION.TEST_KEY}1`, {id: '1'});
+
+            const {result, rerender} = renderHook((key: string) => useOnyx(key), {initialProps: `${ONYXKEYS.COLLECTION.TEST_KEY}skippable-id` as string});
+
+            await act(async () => waitForPromisesToResolve());
+
+            expect(result.current[0]).toBeUndefined();
+            expect(result.current[1].status).toEqual('loaded');
+
+            // Switch to a valid key whose value is in storage but not in cache — should transition through loading
+            rerender(`${ONYXKEYS.COLLECTION.TEST_KEY}1`);
+
+            expect(result.current[1].status).toEqual('loading');
+
+            await act(async () => waitForPromisesToResolve());
+
+            expect(result.current[0]).toEqual({id: '1'});
+            expect(result.current[1].status).toEqual('loaded');
+        });
     });
 
     describe('RAM-only keys', () => {