[Onyx Audit] Batching mechanism feels overcomplicated

[fabioh8010]

Overview

This issue is part of our current effort to audit the performance and scalability of Onyx.

During the Step 1: Objective Analysis we gathered a list of Perceived Problems we have exclusively about Onyx, that is, the problems in Onyx we imagine or have some evidences that that are pain points / performance bottenecks to Onyx, but we need more analysis to confirm if they really are and, if so, how much of a problem they are currently.

The main goal right now is to deep dive into each Perceived Problem, investigate/audit it and come up with findings, which we will use for the next discussions.

Problem

For Onyx.set() and Onyx.merge() functions, the process is basically this one:

1. Changes are broadcasted to subscribers with OnyxUtils.broadcastUpdate().
2. At the same, Storage.setItem()/Storage.mergeItem() is called to store the changes (according to the code this approach prioritizes fast UI changes without waiting for data to be stored in device storage). If the value is being removed, OnyxUtils.remove() will be called instead and do the same steps below.
3. OnyxUtils.broadcastUpdate() notifies subscribers and writes current value to cache. OnyxUtils.scheduleSubscriberUpdate() is called to schedule the subscriber update.
4. OnyxUtils.scheduleSubscriberUpdate() will first call Onyx.keyChanged() in the next tick to notify the Onyx.connect() subscribers. Then, it will use OnyxUtils.batchUpdates() to batch a Onyx.keyChanged() call that notify the withOnyx()/useOnyx() subscribers. Finally, it waits for both OnyxUtils.maybeFlushBatchUpdates() and the first Onyx.keyChanged() to resolve.

The major difference is that Onyx.merge() will queue multiple changes to the same Onyx key, and later batch them into just one.

For Onyx.update(), SET and MERGE operations are queued by their Onyx keys. Collection key changes are grouped to be used into Onyx.mergeCollection() calls.

This overall system of broadcasting and merging looks too complex at a first glance, and I'm not really sure we still need it with React concurrent and useOnyx().

Action Items

1. Analyse how other state management libraries handle batching/setting/merging of multiple store updates, and compare with ours.
2. Understand why we have all these complexities in our internals and evaluate what we could change or simplify.

Relevant Links

Step 1 Canvas: https://expensify.slack.com/docs/T02P6RS2ZK7/F09B668M1MG

Slack thread: https://expensify.slack.com/archives/C05LX9D6E07/p1757416584331869

[VickyStash]

Hi, I'm Viktoryia from Callstack - expert contributor group - and I would like to work on this issue!
It would be the great opportunity for me to continue these investigations: #66799 (comment)
I should have space to work on it in ~1-2 weeks

[VickyStash]

First of all, I want to say that it feels like the current batching doesn't make sense anymore, after withOnyx HOC was removed.

The batching logic was always working only for withOnyx subscribers!

---

Let's take step back and see how it worked before withOnyx was removed.
See how scheduleSubscriberUpdate looked:

function scheduleSubscriberUpdate<TKey extends OnyxKey>(
    key: TKey,
    value: OnyxValue<TKey>,
    previousValue: OnyxValue<TKey>,
    canUpdateSubscriber: (subscriber?: Mapping<OnyxKey>) => boolean = () => true,
): Promise<void> {
    const promise = Promise.resolve().then(() => keyChanged(key, value, previousValue, canUpdateSubscriber, true, false));
    batchUpdates(() => keyChanged(key, value, previousValue, canUpdateSubscriber, false, true));
    return Promise.all([maybeFlushBatchUpdates(), promise]).then(() => undefined);
}

1. In this line:

const promise = Promise.resolve().then(() => keyChanged(key, value, previousValue, canUpdateSubscriber, true, false));

The keyChanged was called with notifyConnectSubscribers=true and notifyWithOnyxSubscribers=false, so connect subscribers, including useOnyx subscribers, are notified right away.

2. And then in this line:

 batchUpdates(() => keyChanged(key, value, previousValue, canUpdateSubscriber, false, true));

The keyChanged was called with notifyConnectSubscribers=false and notifyWithOnyxSubscribers=true.
So notifying of the withOnyx subscribers was batched!

---

In comparison, right now:
https://github.com/Expensify/react-native-onyx/blob/413cb1ad4f23b46aa3c35990c82e5bbe37bf15a1/lib/OnyxUtils.ts#L821-L823

1. We call keyChanged with notifyConnectSubscribers=true, so notifying connect subscribers right away

const promise = Promise.resolve().then(() => keyChanged(key, value, canUpdateSubscriber, true));

2. And then batch call keyChanged with notifyConnectSubscribers=false.

batchUpdates(() => keyChanged(key, value, canUpdateSubscriber, false));

But keyChanged notify no subscibers in this case (cause all of our subscribers are connect subscribers).
https://github.com/Expensify/react-native-onyx/blob/413cb1ad4f23b46aa3c35990c82e5bbe37bf15a1/lib/OnyxUtils.ts#L730-L732

[VickyStash]

UPD:
Today I had some capacity and decided to experiment and get rid of the current batching at all (based on my comment). Aaaand, I see some tests failing, which means I'm missing something, so I'm trying to understand all of the logic and what's happening under the hood.

The test which is failing for me: should replace the old value after a null merge in the top-level object when batching updates.

I want to confirm one thing, is it expected that the Onyx.connect callback is called twice with the same data when waitForCollectionCallback is true. You can check with this simple jest test (it will fail):

    it("shouldn't call a collection connection twice when setting a connection value", async () => {
        const mockCallback = jest.fn();
        connection = Onyx.connect({
            key: ONYX_KEYS.COLLECTION.TEST_UPDATE,
            waitForCollectionCallback: true,
            callback: mockCallback,
        });

        await Onyx.set(`${ONYX_KEYS.COLLECTION.TEST_UPDATE}1`, 'test');

        expect(mockCallback).toHaveBeenCalledTimes(1);
    });

[VickyStash]

Updates:

• Finally figured out the reason for failing tests after removing the batching mechanism. It's not due to the removal of the batching mechanism on its own, but due to the removal of 0 setTimout. Without this "macrotask barrier" the outdated cache results are returned in callbacks in some cases.

I've prepared a Draft PR Expensify/react-native-onyx#689 with the removed batch mechanism to simplify the testing and other experiments.

TODOs:

• I want to try to figure out why we need to have a macrotask to make the cache mechanism work as expected.
• Check performance tests
• Test the E/App over onyx without batching mechanism

[VickyStash]

Updates:

Finally figured out the reason for failing tests after removing the batching mechanism. It's not due to the removal of the batching mechanism on its own, but due to the removal of 0 setTimout. Without this "macrotask barrier" the outdated cache results are returned in callbacks in some cases.

I've investigated more why cache doesn't work as expected if I remove the 0 setTimeout inside the scheduleSubscriberUpdate/scheduleNotifyCollectionSubscribers.

1. The reason is that when the Onyx.connect is called we call the subscribeToKey.
2. Inside the subscribeToKey function we have deffered logic
3. In this logic if waitForCollectionCallback is true we call the callback with cahced data if it exists.

Problem: without extra setTimout in the scheduleSubscriberUpdate/scheduleNotifyCollectionSubscribers there is a race condition, when the connection subscriber and cache is firstly updated by the data from scheduleSubscriberUpdate and then with outdated data from subscribeToKey callback.

So, to sum up what I have for now:

• It looks like current batching logic doesn't used in the way it was created to be used (for withOnyx). More details here.
• setTimeout inside maybeFlushBatchUpdates helps to prevent cache race condition (we can use just setTimeout without any bacthing logic, the effect seems to be the same).

TODOs:

• Check performance tests
• Test the E/App over onyx without batching mechanism

[VickyStash]

• I've tested the E/App over the branch without a batching mechanism, and overall there is no huge difference (which I think is expected), sometimes the duration is a little longer/shorter over some flows, but mostly it feels the same. Example, app restart:

• All jest tests have passed.

[VickyStash]

Summary of what I have for now:

• Current batching logic that exists inside the onyx lib was created for withOnyx HOC, and has no batching effect right now. More details here.
• setTimeout inside maybeFlushBatchUpdates helps to prevent cache race condition (we can use just setTimeout without any bacthing logic, the effect seems to be the same). More details here.
• As following from the above, removing of the current batching doesn't affect the E/App a lot, it tests pretty the same: details here.

TODOs:
Right now, useOnyx hook subscribers have no additional batching. In the E/App the concurrent rendering is enabled, which means some updates get batched automatically. But it not always happens for the complex Onyx.udates operations, so I want to:

• define cases where the lack of additional batching for useOnyx subscribers is causing problems
• consider, how it can be solved and if we should implement the custom batching for useOnyx.

[fabioh8010]

• I've tested the E/App over the branch without a batching mechanism, and overall there is no huge difference (which I think is expected), sometimes the duration is a little longer/shorter over some flows, but mostly it feels the same. Example, app restart:

* All jest tests have passed.

Yes, I was discussing with @VickyStash and I think this is pretty much the direction we want to move – strip away this whole complex batching logic from Onyx and just leave to React to handle the updates itself.

I think if it’s like >90% of things are working and we just have a few issues, maybe it’s just something we can adjust in E/App
If the macro tasks hacks are needed only for tests – even better.

We should be able to remove this, I'm pretty sure the other state management libraries don't need this whole stuff.