How to create complex trees of dependencies with jotai?

[ivankoleda]

I'm trying to perform a spike using Jotai for creating large trees of derived values.

Spike/Benchmark

Cells of each next layer are calculated from the previous layer cells + the value of the current cell. Clicking on cell increments its value and should update all the dependants.
A2 Derived=A2+B1
B2 Derived=B2+A1-C1
C2 Derived=C2+B1+D1
D2 Derived=D2+C1

It's based on this benchmark for reactive computations: https://github.com/Riim/cellx?tab=readme-ov-file#benchmark

This is what it looks like:

jotai_tree.mp4

https://stackblitz.com/edit/vitejs-vite-szeeg9?file=src%2FApp.tsx,src%2Findex.css&terminal=dev

Issue

I cannot go to up to 30 layers with jotai. Creating the tree just freezes the page on initial load. I'm a jotai newbie so I might be missing something obvious in the documentation.

[dai-shi]

Thanks for reporting.
I don't find anything obvious. 30 layers mean 120 atoms? That doesn't sound too many.

I wonder where's the bottleneck. Would anyone like to dive into it?

[DavidJFelix]

Atom creation takes no time at all from what I can tell. The cost is when they're mounted and subbed. I have an example below where I sub the bottom-most layer, which triggers all subsequent layers to sub and it seems a little faster but still rather slow. I think react's rendering will sub in reverse where the base layer is subbed first, then the dependent layer then that layers dependent layer and so on. This seems to be really slow and it's not clear to me why.

• Why is react render-triggered-mounting slower than subbing the bottom most layer?
  • My guess here is that batch mounting the atoms is saving us time rather than looping through creation with react.
• Why does each layer seem to cost more?
  • The only thing I can think is that a sub that triggers a sub is exploding exponentially and while its pretty cheap to sub an atom, having a depth of 30 means registering thousands of dependent subscriptions at the bottom. I'm not sure if this theory tracks with the internals of jotai though.
• What is the recommended remedy for performanc?
  • I think maybe an "atom-of-atoms" pattern here could help, but I haven't tried. Since these dependent atoms are a grouping of data, it might save time on mount if they're all mounted via one atom of atoms? I'll explore some ideas i have soon.

[dai-shi]

which triggers all subsequent layers to sub and it seems a little faster but still rather slow

How about this?

https://stackblitz.com/edit/vitejs-vite-xtndpz?file=src%2FApp.tsx&terminal=dev

// force mounting the last layer
store[store.length - 1].forEach((a) => {
  const unsub = getDefaultStore().sub(a, () => {});
  // we need to call unsub() if no longer used, so this isn't for production
});

[DavidJFelix]

That's basically what I did below. Like I said, it's a little faster than allowing react to do it (you have to bump it up to 25+ layers to see the delay) but it's still pretty slow and each layer seems to be significantly more costly than the last. try:

for (let i = 1; i < 35; i++) {
    layers.push(createLayer(layers[i - 1]));
 }

[DavidJFelix]

https://stackblitz.com/edit/vitejs-vite-92dmdb?file=src%2FApp.tsx updated

[dai-shi]

That's basically what I did below.

Ah, I see.

The only thing I can think is that a sub that triggers a sub is exploding exponentially and while its pretty cheap to sub an atom, having a depth of 30 means registering thousands of dependent subscriptions at the bottom.

Alright, so I think we could mount atoms in an efficient order:

Here we go: https://stackblitz.com/edit/vitejs-vite-65jnfz?file=src%2FApp.tsx

store.forEach((layer) =>
  layer.forEach((a) => {
    const unsub = getDefaultStore().sub(a, () => {});
    // we need to call unsub() if no longer used, so this isn't for production
  })
);

[dai-shi]

Is #2334 related?

[DavidJFelix]

I don't think it's the number of atoms that's the problem, but rather the function creating the atoms. I managed to get 31 to work but I do have to click "wait" when the page is reported as unresponsive.

[DavidJFelix]

Here's a picture of 31. It seems fine when running changes, but the initial setup does take awhile.

[DavidJFelix]

Here's the react profile -- the final line takes 6 seconds. It's not clear to me if this is atoms setting up a subscription depth of 30 or what.

It took 16 seconds to do 31 layers. My guess is 32 would be something like 23 seconds on my machine?

[DavidJFelix]

Here's 30, which I'm also finding to be the point where iterating slows down. I'll see if I can dig into what's actually causing issues.
https://github.com/user-attachments/assets/726f41e0-66a7-477d-8989-eb98cc768c9c

[DavidJFelix]

I changed the code a little and managed to get 35. The timing code says it takes 46 seconds just to mount the bottom-most atoms with 35 layers, but actually rendering the component is quick because the atoms are already loaded. My guess is the reason this is slow is because it's chaining observations in a call-stack. You're seeing it in the react component because the atoms are first mounted by react, but we can pull that into vanilla by mounting them up front.

function createAtoms() {
  const layers: Array<AtomLayer> = [
    [
      createSimpleIncrementAtom(),
      createSimpleIncrementAtom(),
      createSimpleIncrementAtom(),
      createSimpleIncrementAtom(),
    ],
  ];

  for (let i = 1; i < 35; i++) {
    layers.push(createLayer(layers[i - 1]));
  }

  return layers;
}

const store = createAtoms();

const jotaiStore = createStore();
// Pre warm the bottom-most layer by subbing it
const t0 = Date.now();
jotaiStore.sub(store[store.length - 1][0], () => {});
jotaiStore.sub(store[store.length - 1][1], () => {});
jotaiStore.sub(store[store.length - 1][2], () => {});
jotaiStore.sub(store[store.length - 1][3], () => {});
const t1 = Date.now();
console.log(`Time: ${t1 - t0}`);

const Counter = ({ counterAtom }: { counterAtom: CounterAtom }) => {
  const [count, increment] = useAtom(counterAtom);

  return <button onClick={increment}>{count}</button>;
};

const Layer = ({ layer }: { layer: AtomLayer }) => {
  return (
    <div>
      {layer.map((counterAtom, i) => (
        <Counter counterAtom={counterAtom} key={i} />
      ))}
    </div>
  );
};

const App = () => {
  return (
    <Provider store={jotaiStore}>
      <div>
        <p>A2 Derived=A2+B1</p>
        <p>B2 Derived=B2+A1-C1</p>
        <p>C2 Derived=C2+B1+D1</p>
        <p>D2 Derived=D2+C1</p>
        {store.map((layer, i) => (
          <Layer key={i} layer={layer} />
        ))}
      </div>
    </Provider>
  );
};

export default App;

[DavidJFelix]

Times I saw with different layer counts:

28 -> 1.668
29 -> 2.762 (+1.1s from 28)
30 -> 4.388 (+1.6s from 29)
31 -> 6.961 (+2.6s from 30)
32 -> 11.335 (+4.4s from 31)
33 -> 18.084 (+7s from 32)
34 -> 30.172s (+12s from 33)
35 -> 46.984s (+16s from 34)

It seems like each layer costs more than the last to set up, which is why you're seeing this.

[dmaskasky]

I got 3000. Is this cheating?

https://stackblitz.com/edit/vitejs-vite-pvkhtl?file=src%2FApp.tsx

[dai-shi]

Interesting use of unstable_derive.

See also: #2754 (reply in thread)

Actually, setTimeout(unsub) works: https://stackblitz.com/edit/vitejs-vite-g6yavh?file=src%2FApp.tsx