Skip to content

fix: DH-21144: usePersistentState flaky with React 18 #2590

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
mattrunyon merged 6 commits into from
Jan 5, 2026
Merged

fix: DH-21144: usePersistentState flaky with React 18 #2590

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
a1a85f4
fix: DH-21144: usePersistentState flaky with React 18
mattrunyon Dec 10, 2025
54ec7c2
Change to key-based state instead of render order
mattrunyon Dec 15, 2025
e44eef4
Fix types
mattrunyon Dec 16, 2025
91f5f31
Fix test
mattrunyon Dec 16, 2025
f83aeff
Add test
mattrunyon Dec 16, 2025
2d5ed4a
Add migrateState docs
mattrunyon Jan 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
60 changes: 36 additions & 24 deletions packages/plugin/src/PersistentStateContext.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,12 +29,11 @@ export type PersistentStateContextType = {
* It is expected that each hook will call this once to get its initial state.
* @param id A unique ID generated by the caller.
* This is used to prevent one caller from consuming multiple values from the initial state.
* @returns The iterator result containing the initial state value (state, version, type) and a done flag.
* @returns The iterator result containing the initial state value (state, version, type) if it exists.
*/
getInitialState: <S>(id: string) => {
value: { state: S; version: number; type: string };
done: boolean;
};
getInitialState: <S>(
id: string
) => { state: S; version: number; type: string } | undefined;
};

/**
Expand All @@ -44,26 +43,32 @@ export const PersistentStateContext =
createContext<PersistentStateContextType | null>(null);
PersistentStateContext.displayName = 'PersistentStateContext';

export type PersistentState<S = unknown> = {
state: S;
version: number;
type: string;
};

export type PersistentStateProviderProps = React.PropsWithChildren<{
/**
* The initial state of all calls to usePersistentState.
* If there are more calls to usePersistentState than there are elements in this array,
* the state initializer of the usePersistentState call will be used for the rest.
*/
initialState: unknown[];
initialState: PersistentState[];

/**
* Called when the state changes.
* The state is passed as an array of the values of all calls to usePersistentState.
* The order of the values is the same as the order of the calls to usePersistentState.
* @param state The state of all calls to usePersistentState.
*/
onChange: (state: unknown[]) => void;
onChange: (state: PersistentState[]) => void;
}>;

/**
* Tracks all calls to the usePersistentState hook below this provider.
* Keeps track of the state in call order so and calls onChange when the state changes.
* Keeps track of the state in call order and calls onChange when the state changes.
*/
export function PersistentStateProvider(
props: PersistentStateProviderProps
Expand All @@ -78,12 +83,12 @@ export function PersistentStateProvider(
// while updating its state if it changed in that re-render.
const persistentData = useRef({
initial: initialState,
initialStateMap: new Map<string, { value: unknown; done: boolean }>(),
state: new Map<string, unknown>(),
initialStateMap: new Map<string, PersistentState | undefined>(),
state: new Map<string, PersistentState>(),
isTracking: true, // We want to start tracking on the first render
});

const addState = useCallback((id: string, state: unknown) => {
const addState = useCallback((id: string, state: PersistentState) => {
if (persistentData.current.isTracking) {
persistentData.current.state.set(id, state);
}
Expand All @@ -93,7 +98,6 @@ export function PersistentStateProvider(
// Don't trigger again if we are already tracking a render
if (!persistentData.current.isTracking) {
persistentData.current.isTracking = true;
persistentData.current.state = new Map<string, unknown>();
setUpdateId(prev => prev + 1);
}
}, []);
Expand All @@ -102,17 +106,16 @@ export function PersistentStateProvider(
persistentData.current.initial[Symbol.iterator]()
);

const getInitialState = useCallback(function getState<S>(id: string): {
value: S;
done: boolean;
} {
const getInitialState = useCallback(function getState<S>(
id: string
): PersistentState<S> | undefined {
// Prevents a component re-rendering multiple times in the same render cycle from taking multiple values from the iterator
const initialStateForId = persistentData.current.initialStateMap.get(id);
if (initialStateForId) {
return initialStateForId as { value: S; done: boolean };
return initialStateForId as PersistentState<S>;
}
const { value, done } = initialStateIterator.current.next();
const stateVal = { value: value as S, done: done ?? false };
const { value } = initialStateIterator.current.next();
const stateVal = value as PersistentState<S> | undefined;
persistentData.current.initialStateMap.set(id, stateVal);
return stateVal;
}, []);
Expand All @@ -127,15 +130,24 @@ export function PersistentStateProvider(
[updateId, onChange]
);

const contextValue = useMemo(
() => ({
// updateId is used to force consumers to re-render when a state update is scheduled.
// This way we keep states up-to-date and in order.
// For example, components with a key being re-ordered would not have new IDs, so we can't just
// track by ID since we will need the order to rehydrate the state correctly.
const contextValue = useMemo(() => {
// Clear tracking here to avoid a React 18 batched updates issue where
// we cleared the map, then the changed component re-rendered its usePersistentState hook
// putting the states out of order.
if (persistentData.current.isTracking) {
persistentData.current.state = new Map();
}
Copy link
Member

@mofojed mofojed Dec 11, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Mmmm why is this in the useMemo hook? Don't like that there's a side effect executing in a useMemo.

Copy link
Collaborator Author

@mattrunyon mattrunyon Dec 11, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ya it feels gross to me too. The other option was a useEffect which triggers another render pass. That caused some unit tests to fail on the expected number of calls to onChange because it triggers an onChange and then collects the actual state and triggers again.

The bug seems to be related to batched state updates from what I can tell. So if I have 2 tables in 1 panel and filter/sort the 2nd table, it is doing something like this

  1. Schedule a persistent data update. This clears the map to collect the data
  2. Something changed in React 18 (batched updates I think) results in the 2nd table running its render function before the updateVersion state is committed in the provider
  3. The 2nd table running its render function now calls addState which makes it the 1st state in the map
  4. The updateVersion state is committed and rendered. At this point, the context value is supposed to change to force all subscribers to the context to run again so we can collect the data in render order; however, the 2nd table has already run once, so it just updates its value in the map.
  5. The 1st table renders (by force) and adds its state to the map. It is now 2nd in the map.
  6. The map is saved w/ the data from the 2nd table as the 1st entry because it logged itself first.

These can be further out of order b/c there is also a table plugin state for each, but that is undefined. It just potentially messes with the order further.

So this fix is to push that map collection until after we know the version has changed.

I am leaning towards my idea w/ __dhId + type to persist the states. The limitation of 1 state type per panel/widget seems acceptable especially if it means less possibility of breaking from the order of render calls

return {
addState,
scheduleStateUpdate,
getInitialState,
updateId,
}),
[addState, getInitialState, scheduleStateUpdate, updateId]
);
};
}, [addState, getInitialState, scheduleStateUpdate, updateId]);

return (
<PersistentStateContext.Provider value={contextValue}>
Expand Down
11 changes: 7 additions & 4 deletions packages/plugin/src/usePersistentState.test.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,10 @@ import React from 'react';
import { render, screen } from '@testing-library/react';
import userEvent from '@testing-library/user-event';
import { TestUtils } from '@deephaven/test-utils';
import { PersistentStateProvider } from './PersistentStateContext';
import {
type PersistentState,
PersistentStateProvider,
} from './PersistentStateContext';
import usePersistentState, {
type PersistentStateMigration,
} from './usePersistentState';
Expand Down Expand Up @@ -53,8 +56,8 @@ function createWrapper({
initialState = [],
onChange = jest.fn(),
}: {
initialState?: unknown[];
onChange?: (state: unknown[]) => void;
initialState?: PersistentState[];
onChange?: (state: PersistentState[]) => void;
}) {
return function Wrapper({ children }: { children: React.ReactNode }) {
return (
Expand Down Expand Up @@ -172,7 +175,7 @@ describe('usePersistentState', () => {
</>
);

expect(mockOnChange).toHaveBeenCalledWith([
expect(mockOnChange).toHaveBeenLastCalledWith([
expect.objectContaining({ state: 'updated-foo' }),
expect.objectContaining({ state: 'default-baz' }),
]);
Expand Down
22 changes: 9 additions & 13 deletions packages/plugin/src/usePersistentState.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -84,16 +84,10 @@ export function usePersistentState<S>(
// Otherwise, usePersistentState might be called twice by the same component in the same render cycle before flushing in the provider.
const [id] = useState(() => nanoid());
const context = useContext(PersistentStateContext);
const { value: persistedData, done } = context?.getInitialState<S>(id) ?? {
value: undefined,
done: true,
};

// If not done, then we can use the persisted state
// Otherwise, we have exhausted the persisted state
// By checking done, we are able to explicitly save undefined as a state value
const persistedData = context?.getInitialState<S>(id);

const [state, setState] = useState<S>(() => {
if (persistedData == null || done) {
if (persistedData == null) {
return typeof initialState === 'function'
? (initialState as () => S)()
: initialState;
Expand Down Expand Up @@ -124,23 +118,25 @@ export function usePersistentState<S>(

context?.addState(id, stateWithConfig);

const scheduleStateUpdate = context?.scheduleStateUpdate;

// This won't cause unnecessary renders on initial mount because the state is already tracking,
// so calls to scheduleStateUpdate will be no-ops since tracking finishes in an effect at the provider after this effect.
// When a component mounts after the parents have already rendered, this will trigger a re-render to track the new state immediately.
useEffect(
function scheduleUpdateOnMountAndChange() {
context?.scheduleStateUpdate();
scheduleStateUpdate?.();
},
[context, state]
[scheduleStateUpdate, state]
);

useEffect(
function scheduleUpdateOnUnmount() {
return () => {
context?.scheduleStateUpdate();
scheduleStateUpdate?.();
};
},
[context]
[scheduleStateUpdate]
);

return [state, setState];
Expand Down
Loading