readAtomState is called many times before atoms are mounted

[edkimmel]

We had a bad performance regression in our project recently. The root cause was the addition of a deep atom nesting structure to an already deep nested atom structure.

The call count for readAtomState during SSR jumped from 60k (Which is already a lot) to 170k due to the introduction of a few new atoms at the top of the tree. This is a situation where all data has been fetched and we are bootstrapping and rendering a large stateful page in one render.

Update - a new content rendering flow was found to call readAtomState 1,050,000 times during SSR due to a very heavy atom tree. This was brought down to 30k with a local patch to jotai similar to the proposal here.

#2333 was created to showcase the issue and hopefully find a solution acceptable to be put back upstream.

The root of the issue appears to be that jotai cannot make assumptions about the value of an atom until after the atoms are mounted. This is due to useHydrateAtoms and other functionality which is able to mutate the state of an atom immediately, expecting all future reads in the same render pass to return the new state.

One proposal is the ability to tell a jotai store that an atom will be immutable prior to mounting, allowing it to exit without iterating through the dependency tree. This has the benefit of being backwards compatible and not changing the core interface most users of jotai would work with.

[dai-shi]

Thanks for reporting. I'll take a closer look eventually, but it may take time.
Meanwhile, do you know if any previous version work better than the latest one?

[edkimmel]

I looked through the commit history of v2, and the last v1 tag. In all iterations of store.ts, it appears to always recurse through the entire dependency chain while unmounted.

I still believe the safe approach is an opt-in approach of specifying which atoms can be considered immutable until mounted. This is a bit cumbersome to flag in practice, but did yield great results (1 million down to 30k by flagging 4 of our worst offenders).

If we flip it around and attempt to consider atoms "pure until proven otherwise", I can think of two options right now

First option

• A call to store.set on an unmounted atom flags that atom as impure.
• Atoms are also impure if they depend on an impure atom.
• Pure atoms are able to return the pre-existing atomState while unmounted, without recursing through dependencies. Impure atoms must always recurse.
• This method requires us to still recurse in readAtomState to check for impure atoms. However, this is just reading data and not heavy object creation.

Second option

• When store.set is called, flag every known atom with an existing state in the store as impure
• As atoms go through readAtomState, they clear their own impure status. Pure atoms are still safe to return the pre-existing values and impure atoms must recurse.
• This means a mostly harmless store.set on a leaf atom somewhere will cause every atom in the application to re-evaluate once again. However, the state of code today has every atom re-evaluating every call anyway, so this should still be a net win.

[edkimmel]

@dai-shi I took some time to implement the second option offered above in our core project.

On our worst possible atom-heavy path, we've gone from 1,050,000 readAtomState to 3100. Around 300 of those reads were an additional recursion path due to invalidations due to a store.set on an unmounted atom.

I'll update the attached PR to showcase the new approach today.

[dai-shi]

Okay, I have some time and read your description again. (My previous comment was based on my misunderstanding.)

First of all, yes I understand the behavior and I had also wondered if there's any way to avoid re-evaluation for unmounted atoms before.
But, to make atom state garbage collectable, it's unavoidable.

Now, what I wasn't aware is this becomes a real issue with SSR. Because, in SSR, atoms are never mounted, atom states are never cached, right? That's something to consider.

I thought about something like your second approach too. Yeah, WeakSet is a nice idea. But, it complicates the implementation. I don't think it makes sense other than SSR scenario.

That said, if this is SSR only issue, I think providing a wrapper util or even non-subscribeable cache-only store creator makes sense. Or, SSR-aware Provider. In any case, I think it's a better idea to not modify the core implementation, but do things outside.
(If there's a real issue on client side, without SSR, I would need to re-consider.)

[edkimmel]

This also affects the initial hydration render inside of the browser. The initial hydration tends to be websites' longest blocking task for the sake of Lighthouse Total Blocking Time score. Having a solution in place helps to keep that blocking task minimal.

[dai-shi]

I see. Interesting. I still think we should do this outside, at least until we learn more use cases and how some approaches improve them.
You have two options:

1. Create an "unstable_" util in this repo. (The release cycle is slow.)
2. Create a new package in Jotai Labs. (You can do as you like.)

Be aware that some of packages Jotai Labs are really important ones: jotai-scope, jotai-effect, etc.

[eigilsagafos]

@edkimmel Did you find a solution here? We have the exact same issue. We are moving from Recoil to Jotai and have a case where loading with Recoil takes 100ms and Jotai takes 10 000ms + (worst cases crashes the browser)

[edkimmel]

Yes, we maintain a yarn patch on jotai that adds an unmounted atom value cache to the store. It's a cleaner version of the example PR opened and linked to from the first post on this discussion.

[dai-shi]

Nice. @edkimmel Can you try the PR build?
https://ci.codesandbox.io/status/pmndrs/jotai/pr/2333
☝️ Find "Local Install Instructions"

[jarandmi]

Thanks, this helped us. @edkimmel: Did you update the hack to work with the new store implementation? Would you be willing to share the patch with us, or is it the same as in the PR above?