aws / aws/graph-explorer

Unbounded atomFamily growth: families keyed on freshly allocated arrays and objects are never released

Open
#2,116 0 comments 0 reactions 0 assignees View on GitHub
internal performance tech debt
Dominant language
TypeScript
Stars
481
Forks
108
Avg merge
6d 8h
Merged PRs (30d)
5

Description

`atomFamily` from `jotai-family` caches one atom per parameter identity and never evicts unless `remove` or `setShouldRemove` is called. Nothing under `packages/graph-explorer/src` calls either — verified, zero matches. That is fine for a family keyed on a stable branded ID, where the key set is bounded by the data. It is not fine for a family keyed on a value that is freshly allocated on every recomputation: each recomputation interns a new entry that can never be reached again, and each entry retains whatever its derived atom computed. Retained memory then grows with the number of state mutations rather than with the size of the data.

Neighbor expansion is the app's primary interaction and it mutates `nodesAtom`/`edgesAtom` on every expand, so this grows during ordinary use. It is not a crash and there is no recovery path short of closing the tab.

## Why it happens

`useAllNeighbors` is the clearest case. It derives the key array from the node map, so the array is a new object on every `nodesAtom` change:

```ts
// packages/graph-explorer/src/core/StateProvider/neighbors.ts:98-103
export function useAllNeighbors() {
const vertices = useAtomValue(nodesAtom);
const vertexIds = useMemo(() => vertices.keys().toArray(), [vertices]);
...
const fetchedNeighbors = useAtomValue(allFetchedNeighborsSelector(vertexIds));
```

The `useMemo` stabilises the array across renders but not across store mutations, which is exactly the axis that matters here. Every expand interns a fresh entry in `allFetchedNeighborsSelector`, and that entry holds a `Map` over the whole graph as it stood at that moment.

## Sites

All keyed on a value allocated fresh per recomputation:

1. `core/StateProvider/neighbors.ts:229` — `allFetchedNeighborsSelector = atomFamily((ids: VertexId[]) => …)`. **Highest priority**: same shape and severity as the one already fixed, in the neighbor path, driven by the same interaction, and each retained entry is proportional to the whole graph.
2. `core/StateProvider/displayEdge.ts:67` — `displayEdgeSelector = atomFamily((edge: Edge) => …)`, keyed on object identity. A re-derived `Edge` interns a second entry for an edge that already had one. Reached from three call sites, including the map over the full edge set at `displayEdge.ts:129`.
3. `core/ConfigurationProvider/useConfiguration.ts:50` — `vertexTypeConfigsSelector = atomFamily((vertexTypes?: VertexType[]) => …)`.
4. `core/ConfigurationProvider/useConfiguration.ts:80` — `edgeTypeConfigsSelector = atomFamily((edgeTypes?: EdgeType[]) => …)`. Same shape as 3; both retain a config array per caller-supplied array identity.

The remaining families in the codebase are already correctly keyed on a branded ID or on a `ConfigurationId | null`, and are not in scope: `vertexTypeConfigSelector`, `edgeTypeConfigSelector`, `vertexStyleByTypeAtom`, `edgeStyleByTypeAtom`, `displayVertexTypeConfigSelector`, `displayEdgeTypeConfigSelector`, `nodeSelector`, `edgeSelector`, `fetchedNeighborsSelector`, `fetchedNeighborIdsAtom`, `schemaByIdAtom`.

## Fix

Prefer deleting the array-keyed layer over adding eviction — bounding the cache keeps the concept, removing the key retires it. `setShouldRemove` is a last resort, not the goal.

1. **Build the collection from a per-id family** rather than interning the whole collection under one array key. For `allFetchedNeighborsSelector` the per-id family it needs (`fetchedNeighborsSelector`, `neighbors.ts:181`) already exists and is already correctly keyed, so the array-keyed wrapper can go away entirely.
2. **Key every family on a branded ID** — `VertexId`, `EdgeId`, `VertexType`, `EdgeType` — never on a freshly allocated object or array. For `displayEdgeSelector`, key on `EdgeId` and read the `Edge` inside via `edgeSelector`, mirroring what `displayVertexSelector` now does.
3. **Where a public hook must keep an array parameter for its callers**, have it read a non-family context atom and map over the input, interning nothing. `useDisplayVerticesFromVertices` + `displayVertexContextSelector` (`displayVertex.ts:52-56`, `:80`) is the pattern to copy; it fits `useVertexTypeConfigs` / `useEdgeTypeConfigs` directly, since both already read an all-configs atom and then index into it.

## Reference fix

`core/StateProvider/displayVertex.ts` is the worked example, landed on the unmerged branch `schema-view-style-perf`. On `origin/main` that file still has both bad forms — `displayVertexSelector` keyed on `Vertex` object identity (`:73`) and `displayVerticesSelector` keyed on a freshly allocated `Vertex[]` (`:136`). The branch replaces them with a family keyed on `VertexId` plus a plain `displayVertexContextSelector` atom holding the shared derivation, and records the reason in a comment on the family. The four sites above are untouched by that branch.

## Testing note

Retention is not directly assertable — `atomFamily` exposes no size or introspection API. The observable proxy is identity stability: after an unrelated mutation, an entity that did not change should keep the same derived object (`toBe`, not `toEqual`). A test that reads a derived value, mutates a sibling entity, reads again, and asserts referential equality of the untouched one will fail today and pass after the fix.

## Convention

`docs/agents/react.md:20-23` (added on the same branch) documents this under "Client state (Jotai)": prefer a derived atom when several pipelines consume a derivation, and never key a family on a freshly allocated object or array. These four sites are the remaining violations of that rule.

## Also in scope, separable — blocked on `schema-view-style-perf` merging

Not an `atomFamily` problem; a dependency-width problem in the same neighbourhood, worth doing as its own commit.

`canvasVerticesAtom` (`core/StateProvider/renderedEntities.ts:56`) depends on `displayVerticesInCanvasSelector`, which resolves display labels through the vertex style lookup (`displayVertexContextSelector` reads `vertexStyleAtom`). So any user style write recomputes the whole canvas visibility and vertex pipeline, even though the filter predicate needs only `id` and `types` from the raw `Vertex`. Sourcing the predicate from `nodesAtom` would decouple visibility from styling. There is already a comment recording this at `renderedEntities.ts:52-54` — note it names `vertexStyleByTypeAtom`, which is stale; the actual path is `vertexStyleAtom`. Fix the comment along with it.

## Notes

- Pre-existing on `main` for all four sites; none was introduced by an in-flight branch.
- No user-visible symptom to reproduce deterministically. The evidence is structural: no eviction call anywhere in `src`, plus keys that are provably fresh per recomputation.

## Related Issues

- Related to #2104 — same style/canvas hot path, but CPU in the Cytoscape style engine rather than retention.
- Related to #1890 — reference-stable canvas index; the churn it addresses is what drives the atom churn here.
- Related to #1887 — canvas rebuild-per-change; adjacent churn, different layer.
- Related to #263 — the only open report mentioning out-of-memory (RDF). Unconfirmed whether this contributes; worth re-checking once the sites are fixed.

> [!IMPORTANT]
> Internal only — this issue is maintained by the core team and is not accepting external contributions.

Contributor guide

Open the contributing guide

Research direction

Start with the four atomFamily sites in core/StateProvider/neighbors.ts, core/StateProvider/displayEdge.ts, and core/ConfigurationProvider/useConfiguration.ts, then compare core/StateProvider/displayVertex.ts with the schema-view-style-perf branch. Add identity-stability tests that mutate a sibling entity and assert an untouched derived value remains referentially equal; done means the array- and object-keyed families are removed or keyed by stable IDs without changing behavior.

Written by the indexing model from the issue text.

Assessment

Tech stack
react, typescript
Domain
frontend, performance
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Clearly specified
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.