perf: optimize atom store for unmounted atom reads and DAG traversal by edkimmel · Pull Request #3259 · pmndrs/jotai

edkimmel · 2026-03-07T15:07:16Z

Related Bug Reports or Discussions

Fixes discussion # #2334

Summary

The discussion linked is two years old now, time flies.
My old suggestions and approaches no longer work with the big store rewrite, and I do not believe they ever worked with async atoms. My goals are the same, albeit I did use AI to figure out how exactly it could be re-implemented in this new jotai world.

React SSR and React Hydration are slow due to repeated tree traversal reading unmounted atom values. During the synchronous react render, there aren't any mutations to the store to warrant re-traversal, so a caching strategy was desired.
There has been some slight performance regressions over time that have added up in the dependency updates flow since the initial conversion to topological sort a few years ago.

Two key performance optimizations:

Store-level epoch cache for unmounted atoms: Repeated store.get() calls on unmounted derived atoms no longer recursively walk the entire dependency chain when no mutations have occurred. A per-store epoch counter (incremented on each store.set()) allows storeGet to skip the expensive dep walk entirely when the epoch is unchanged.
Inline getMountedOrPendingDependents: The hot-path functions invalidateDependents and recomputeInvalidatedAtoms now iterate mounted.t and atomState.p directly, avoiding a Set allocation on every call during DAG traversal.

Check List

pnpm run fix for formatting and linting code and docs

Benchmark Results (vs main)

Two key performance optimizations: 1. Store-level epoch cache for unmounted atoms: Repeated store.get() calls on unmounted derived atoms no longer recursively walk the entire dependency chain when no mutations have occurred. A per-store epoch counter (incremented on each store.set()) allows storeGet to skip the expensive dep walk entirely when the epoch is unchanged. 2. Inline getMountedOrPendingDependents: The hot-path functions invalidateDependents and recomputeInvalidatedAtoms now iterate mounted.t and atomState.p directly, avoiding a Set allocation on every call during DAG traversal. https://claude.ai/code/session_01Urj5ZjLjsGhU5Ujb3oE8or

vercel · 2026-03-07T15:07:20Z

The latest updates on your projects. Learn more about Vercel for GitHub.

Project	Deployment	Actions	Updated (UTC)
jotai	Ready	Preview, Comment	Mar 7, 2026 3:07pm

codesandbox-ci · 2026-03-07T15:07:50Z

This pull request is automatically built and testable in CodeSandbox.

To see build info of the built libraries, click here or the icon next to each commit SHA.

dai-shi · 2026-03-07T15:14:01Z

I'm almost forgetting the past discussion, but if I remember correctly, we don't plan to merge this as this is only beneficial for some use cases. Thanks to new building blocks, I think you can create an ecosystem library to implement this capability (if not, we should fix the building blocks). I can invite you to https://github.com/jotaijs, let me know if you are interested.

edkimmel · 2026-03-07T15:24:56Z

@dai-shi Ah, that's a fun development! I'll have to take a look.

edkimmel · 2026-03-07T22:54:36Z

@dai-shi
I will take a look at whether it's more realistic for us to maintain an ecosystem library or a fork for ourselves like we have been. Given where the changes are, we may have to re-implement more as an ecosystem than maintaining a fork.

The current benchmarks do not test wide and deep dependency trees. Here's an example addition
https://github.com/edkimmel/jotai/blob/performance-optimizations/benchmarks/derived-read.ts

I would still encourage looking at the problem of redundant tree traversal during that first render before effects mount in core jotai, independent of my proposals. Every read scales linearly with depth/width of the dependency tree, but could be cut down to O(1) with a guarantee that data has not changed since the last read.

dai-shi · 2026-03-08T01:05:29Z

@edkimmel

Thanks for the follow-up.

The current benchmarks do not test wide and deep dependency trees. Here's an example addition https://github.com/edkimmel/jotai/blob/performance-optimizations/benchmarks/derived-read.ts

Ah, I think I misunderstood something. This should be the general improvement for deep/wide use cases (at the cost of complexity).

Please send another PR to add the benchmark file.

I would still encourage looking at the problem of redundant tree traversal during that first render before effects mount in core jotai, independent of my proposals.

Thanks. Sure, I'll take another deeper look soon. Meanwhile, do you think you can reproduce with a failing test (without relying on "time")?

Every read scales linearly with depth/width of the dependency tree, but could be cut down to O(1) with a guarantee that data has not changed since the last read.

Does the guarantee hold? Okay, I think I get it. I see, it's store epoch to guarantee. I think it's valid. Clever. But, I'm not sure if it covers 100%. We can update store async.
Using the WeakMap solves my concern about GC. #2334 (comment)
Did you solve it two years ago? I didn't notice it.

The implementation needs to be re-designed. I'll take care of it as I can't tell my preference easily (and it may end up with something like yours). If you can help the work, the failing test should be really helpful.

dmaskasky · 2026-03-14T17:10:48Z

Can you refactor this to use Building Blocks? I would start by moving the store epoch to buildingBlocks and verified atomEpoch to atomState. We want both to be accessible by stores that compose these building blocks.

We don't rely on derived stores anymore, since stores are now composed by building blocks, so I would drop that requirement and leave it up to the developer to integrate with store cache - but we have to ensure the default behavior is always correctness so if they choose to do nothing, it will just result in extra readAtomState traversals.

dmaskasky · 2026-03-14T17:14:39Z

src/vanilla/internals.ts

    }
    visiting.add(a)
    // Push unvisited dependents onto the stack
-    for (const d of getMountedOrPendingDependents(a, aState, mountedMap)) {


Nice. Can we remove getMountedOrPendingDependents completely or is it still used somewhere?

dmaskasky · 2026-03-14T17:24:33Z

src/vanilla/internals.ts

+      const entry = epochState.verified.get(atom)
+      if (
+        entry &&
+        entry[0] === epochState.epoch &&
+        entry[1] === cachedState.n
+      ) {


Suggested change

const entry = epochState.verified.get(atom)

if (

entry &&

entry[0] === epochState.epoch &&

entry[1] === cachedState.n

) {

const [storeEpoch, atomEpoch] = epochState.verified.get(atom) || []

if (

storeEpoch === epochState.epoch &&

atomEpoch === cachedState.n

) {

dmaskasky · 2026-03-14T17:35:26Z

Can you revisit the benchmarks? I would like to understand the performance of various atom trees:

wide (one atom having 10,000 dependencies)
tall (one atom having one dependency recursive 10,000 times)
balanced (the same width as height)

dmaskasky · 2026-03-14T17:42:27Z

src/vanilla/internals.ts

+    const atomStateMap = buildingBlocks[0]
+    const cachedState = atomStateMap.get(atom)
+    if (cachedState && ('v' in cachedState || 'e' in cachedState)) {


Suggested change

const atomStateMap = buildingBlocks[0]

const cachedState = atomStateMap.get(atom)

if (cachedState && ('v' in cachedState || 'e' in cachedState)) {

const ensureAtomState = buildingBlocks[11]

const atomState = ensureAtomState(atom)

if (isAtomStateInitialized(atomState)) {

dai-shi · 2026-03-15T14:51:03Z

Note that I'll revisit this and #3265 when I get time (maybe next months). I hope to find a minimal solution that should be comfortable. Meanwhile, feel free to continue experiments, but current plan is not merge this PR.

Can you refactor this to use Building Blocks?

Sounds good as an experiment. And, if it can be done with a third-party library, it is very nice as an experiment too.

pkg-pr-new · 2026-03-20T08:54:54Z

More templates

npm i https://pkg.pr.new/jotai@3259

commit: ad4cad6

github-actions · 2026-03-20T08:55:07Z

Preview in LiveCodes

Latest commit: ad4cad6
Last updated: Mar 8, 2026 1:05am (UTC)

Playground	Link
React demo	https://livecodes.io?x=id/FNZ5ZBBJE

See documentations for usage instructions.

vercel bot deployed to Preview March 7, 2026 15:07 View deployment

dai-shi closed this Mar 7, 2026

dai-shi reopened this Mar 8, 2026

dai-shi marked this pull request as draft March 8, 2026 06:39

dmaskasky mentioned this pull request Mar 14, 2026

add performance test to evaluate store.get performance #3265

Open

dmaskasky reviewed Mar 14, 2026

View reviewed changes

Uh oh!

Conversation

edkimmel commented Mar 7, 2026

Related Bug Reports or Discussions

Summary

Check List

Benchmark Results (vs main)

Uh oh!

vercel bot commented Mar 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

codesandbox-ci bot commented Mar 7, 2026

Uh oh!

dai-shi commented Mar 7, 2026

Uh oh!

edkimmel commented Mar 7, 2026

Uh oh!

edkimmel commented Mar 7, 2026

Uh oh!

dai-shi commented Mar 8, 2026

Uh oh!

dmaskasky commented Mar 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

dmaskasky Mar 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dmaskasky Mar 14, 2026

Choose a reason for hiding this comment

Uh oh!

dmaskasky commented Mar 14, 2026

Uh oh!

dmaskasky Mar 14, 2026

Choose a reason for hiding this comment

Uh oh!

dai-shi commented Mar 15, 2026

Uh oh!

pkg-pr-new bot commented Mar 20, 2026

Uh oh!

github-actions bot commented Mar 20, 2026

Preview in LiveCodes

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

vercel bot commented Mar 7, 2026 •

edited

Loading

dmaskasky commented Mar 14, 2026 •

edited

Loading

dmaskasky Mar 14, 2026 •

edited

Loading