[DataGrid] Prevent React state update before mount (@arminmeh)

packages/x-data-grid/src/hooks/utils/useGridSelector.test.tsx

@@ -0,0 +1,113 @@
+import * as React from 'react';
+import { describe, it, expect } from 'vitest';
+import { act, createRenderer, screen } from '@mui/internal-test-utils';
+import { Store } from '@mui/x-internals/store';
+import { useGridSelector } from './useGridSelector';
+
+type TestState = {
+  selectorTestState: {
+    value?: number;
+    a?: number;
+    b?: number;
+  };
+};
+
+type TestApiRef = {
+  current: {
+    state: TestState;
+    store: Store<TestState>;
+  };
+};
+
+const createApiRef = (initialState: TestState): TestApiRef => {
+  const store = Store.create(initialState);
+
+  return {
+    current: {
+      state: store.state,
+      store,
+    },
+  };
+};
+
+const setApiRefState = (apiRef: TestApiRef, newState: TestState) => {
+  apiRef.current.state = newState;
+  apiRef.current.store.update(newState);
+};
+
+const selectorValue = (apiRef: any) =>
+  (apiRef.current.state as TestState).selectorTestState.value ?? 0;
+
+const selectorValueA = (apiRef: any) =>
+  (apiRef.current.state as TestState).selectorTestState.a ?? 0;
+
+const selectorValueB = (apiRef: any) =>
+  (apiRef.current.state as TestState).selectorTestState.b ?? 0;
+
+// Run in both StrictMode on and off
+describe.each([true, false])('useGridSelector (strict: %s)', (strict) => {
+  const { render } = createRenderer({ strict });
+
+  it('should catch store updates fired before the selector subscription is attached', () => {
+    const apiRef = createApiRef({ selectorTestState: { value: 0 } });
+    let didUpdate = false;
+
+    function SelectorProbe() {
+      const value = useGridSelector(apiRef as any, selectorValue);
+
+      const handleRef = React.useCallback((node: HTMLDivElement | null) => {
+        if (node && !didUpdate) {
+          didUpdate = true;
+          setApiRefState(apiRef, {
+            ...apiRef.current.state,
+            selectorTestState: { value: 1 },
+          });
+        }
+      }, []);
+
+      return (
+        <div ref={handleRef} data-testid="selector-probe">
+          {value}
+        </div>
+      );
+    }
+
+    render(<SelectorProbe />);
+
+    expect(screen.getByTestId('selector-probe').textContent).to.equal('1');
+  });
+
+  it('should batch updates of multiple selectors from the same store change into a single commit', async () => {
+    const apiRef = createApiRef({
+      selectorTestState: { a: 0, b: 0 },
+    });
+
+    let commits: string[] = [];
+
+    const SelectorProbe = React.memo(function SelectorProbe() {
+      const valueA = useGridSelector(apiRef as any, selectorValueA);
+      const valueB = useGridSelector(apiRef as any, selectorValueB);
+
+      React.useEffect(() => {
+        commits.push(`${valueA}:${valueB}`);
+      });
+
+      return <div data-testid="selector-probe">{`${valueA}:${valueB}`}</div>;
+    });
+
+    render(<SelectorProbe />);
+
+    // reset before the update
+    commits = [];
+
+    await act(async () => {
+      setApiRefState(apiRef, {
+        ...apiRef.current.state,
+        selectorTestState: { a: 1, b: 1 },
+      });
+    });
+
+    expect(screen.getByTestId('selector-probe').textContent).to.equal('1:1');
+    expect(commits).to.deep.equal(['1:1']);
+  });
+});

packages/x-data-grid/src/hooks/utils/useGridSelector.ts

@@ -1,10 +1,11 @@
 'use client';
 import * as React from 'react';
+import useEnhancedEffect from '@mui/utils/useEnhancedEffect';
 import type { RefObject } from '@mui/x-internals/types';
 import { fastObjectShallowCompare } from '@mui/x-internals/fastObjectShallowCompare';
 import { warnOnce } from '@mui/x-internals/warning';
-import { useSyncExternalStore } from 'use-sync-external-store/shim';
 import type { GridApiCommon } from '../../models/api/gridApiCommon';
+import type { GridStateCommunity } from '../../models/gridStateCommunity';
 import { useLazyRef } from './useLazyRef';
 
 const defaultCompare = Object.is;
@@ -27,20 +28,24 @@ export const argsEqual = (prev: any, curr: any) => {
   return fn(prev, curr);
 };
 
-const createRefs = () => ({ state: null, equals: null, selector: null, args: undefined }) as any;
+const createRefs = () =>
+  ({ state: null, equals: null, selector: null, args: undefined, storeState: null }) as any;
 
 const EMPTY = [] as unknown[];
 
 type Refs<T> = {
+  // `state` holds the value this hook currently returns.
+  // `storeState` remembers which store state that value was computed from.
+  // The store creates a new state object on every update, so comparing `storeState`
+  // with the current `store.state` allows `updateState` to skip needless selector
+  // calls and to catch updates that happened before the hook subscribed to the store.
   state: T;
+  storeState: GridStateCommunity | null;
   equals: <U = T>(a: U, b: U) => boolean;
   selector: Function;
   args: any;
-  subscription: undefined | (() => void);
 };
 
-const emptyGetSnapshot = () => null;
-
 export function useGridSelector<Api extends GridApiCommon, T>(
   apiRef: RefObject<Api>,
   selector: (apiRef: RefObject<Api>) => T,
@@ -75,6 +80,9 @@ export function useGridSelector<Api extends GridApiCommon, Args, T>(
   );
 
   refs.current.state = state;
+  if (!didInit) {
+    refs.current.storeState = apiRef.current.store.state;
+  }
   refs.current.equals = equals;
   refs.current.selector = selector;
   const prevArgs = refs.current.args;
@@ -86,44 +94,40 @@ export function useGridSelector<Api extends GridApiCommon, Args, T>(
       refs.current.state = newState;
       setState(newState);
     }
+    refs.current.storeState = apiRef.current.store.state;
   }
 
-  const subscribe = React.useCallback(
+  const updateState = React.useCallback(
     () => {
-      if (refs.current.subscription) {
-        return null;
-      }
+      const storeState = apiRef.current.store.state;
 
-      refs.current.subscription = apiRef.current.store.subscribe(() => {
+      if (refs.current.storeState !== storeState) {
         const newState = refs.current.selector(apiRef, refs.current.args) as T;
+        refs.current.storeState = storeState;
+
         if (!refs.current.equals(refs.current.state, newState)) {
           refs.current.state = newState;
           setState(newState);
         }
-      });
-
-      return null;
+      }
     },
     // eslint-disable-next-line react-hooks/exhaustive-deps
     EMPTY,
   );
 
-  const unsubscribe = React.useCallback(() => {
-    // Fixes issue in React Strict Mode, where getSnapshot is not called
-    if (!refs.current.subscription) {
-      subscribe();
-    }
-
-    return () => {
-      if (refs.current.subscription) {
-        refs.current.subscription();
-        refs.current.subscription = undefined;
-      }
-    };
+  // Why subscribe in an effect instead of during render: a component can render without
+  // ever mounting (e.g. when it suspends during hydration). If it subscribed during render,
+  // it could receive a store update and call `setState` before being mounted (#17077).
+  // Effects only run for mounted components, so subscribing here is safe.
+  //
+  // Using a layout effect because the store may already have changed
+  // between render and mount (e.g. from a child's ref callback or layout effect).
+  // `updateState()` picks up such changes, so the corrected value is shown right away instead of in a second frame.
+  useEnhancedEffect(() => {
+    updateState();
+    return apiRef.current.store.subscribe(updateState);
     // eslint-disable-next-line react-hooks/exhaustive-deps
   }, EMPTY);
 
-  useSyncExternalStore(unsubscribe, subscribe, emptyGetSnapshot);
-
   return state;
 }

packages/x-data-grid/src/tests/layout.DataGrid.test.tsx

@@ -128,6 +128,34 @@ describe('<DataGrid /> - Layout & warnings', () => {
         expect(ref.current).to.equal(container.firstChild?.firstChild);
       });
 
+      it('mounts the root element during the first SPA commit', () => {
+        let rootElementInLayoutEffect: HTMLDivElement | null | undefined;
+        let renderCount = 0;
+
+        function TestCase() {
+          const apiRef = useGridApiRef();
+
+          React.useLayoutEffect(() => {
+            renderCount += 1;
+            if (renderCount === 1) {
+              rootElementInLayoutEffect = apiRef.current!.rootElementRef.current;
+            }
+          }, [apiRef]);
+
+          return (
+            <div style={{ width: 300, height: 300 }}>
+              <DataGrid apiRef={apiRef} {...baselineProps} />
+            </div>
+          );
+        }
+
+        render(<TestCase />);
+
+        // Assert against the actual mounted root node because `rootElementInLayoutEffect` starts as `undefined`.
+        // A plain null check would pass if the effect never captured the node.
+        expect(rootElementInLayoutEffect).to.equal(document.querySelector(`.${gridClasses.root}`));
+      });
+
       describe('`classes` prop', () => {
         it("should apply the `root` rule name's value as a class to the root grid component", () => {
           const classes = {

packages/x-data-grid/src/tests/ssr.DataGrid.test.tsx

@@ -0,0 +1,100 @@
+import * as React from 'react';
+import * as ReactDOMServer from 'react-dom/server';
+import { describe, it, expect } from 'vitest';
+import { act, reactMajor } from '@mui/internal-test-utils';
+import { DataGrid, useGridApiContext, useGridSelector } from '@mui/x-data-grid';
+import { isJSDOM } from 'test/utils/skipIf';
+
+const COLUMNS = [{ field: 'id' }, { field: 'name', width: 200 }];
+
+const ROWS = [
+  { id: 1, name: 'Alpha' },
+  { id: 2, name: 'Beta' },
+  { id: 3, name: 'Gamma' },
+];
+
+const rowsStateSelector = (apiRef: ReturnType<typeof useGridApiContext>) =>
+  apiRef.current.state.rows;
+
+describe('<DataGrid /> - SSR', () => {
+  // Reproduces https://github.com/mui/mui-x/issues/17077 in a minimal way.
+  // In the Next.js reproduction, the grid synchronously updates its store during hydration
+  // while selector subscribers registered during render are not committed yet. The suspended
+  // footer below mimics that timing: it subscribes with `useGridSelector`, suspends before
+  // mount, and then triggers a grid state update while React still considers the fiber unmounted.
+  // `react-dom/client` only exists in React 18+, so it is imported dynamically inside the test
+  // to avoid breaking module resolution on the React 17 lane.
+  it.skipIf(isJSDOM || reactMajor < 18)(
+    'should not notify grid selector subscribers before they have mounted',
+    async () => {
+      const ReactDOMClient = await import('react-dom/client');
+      let shouldSuspend = false;
+      let didSuspend = false;
+      let promise: Promise<void> | null = null;
+
+      function SuspendedSelectorFooter() {
+        const apiRef = useGridApiContext();
+        useGridSelector(apiRef, rowsStateSelector);
+
+        if (shouldSuspend && !didSuspend) {
+          promise ??= Promise.resolve().then(() => {
+            didSuspend = true;
+            apiRef.current.setState((state) => ({
+              ...state,
+              rows: {
+                ...state.rows,
+              },
+            }));
+          });
+
+          throw promise;
+        }
+
+        return null;
+      }
+
+      const tree = (
+        <React.StrictMode>
+          <React.Suspense fallback={null}>
+            <div style={{ height: 300, width: 400 }}>
+              <DataGrid rows={ROWS} columns={COLUMNS} slots={{ footer: SuspendedSelectorFooter }} />
+            </div>
+          </React.Suspense>
+        </React.StrictMode>
+      );
+
+      const container = document.createElement('div');
+      document.body.appendChild(container);
+      container.innerHTML = ReactDOMServer.renderToString(tree);
+      shouldSuspend = true;
+
+      const errors: string[] = [];
+      const originalConsoleError = console.error;
+      const interceptor = (...args: any[]) => {
+        errors.push(args.map(String).join(' '));
+      };
+      // Direct assignment shadows the property; the test runner's
+      // vitest-fail-on-console wrapper assigned in beforeEach is bypassed.
+      console.error = interceptor;
+
+      let root: ReturnType<typeof ReactDOMClient.hydrateRoot> | undefined;
+      try {
+        await act(async () => {
+          root = ReactDOMClient.hydrateRoot(container, tree);
+        });
+        await act(async () => {
+          root?.unmount();
+        });
+      } finally {
+        console.error = originalConsoleError;
+        document.body.removeChild(container);
+      }
+
+      const hasStateUpdateWarning = errors.some((error) =>
+        /Can't perform a React state update on a component that hasn't mounted yet/.test(error),
+      );
+
+      expect(hasStateUpdateWarning).to.equal(false);
+    },
+  );
+});