fix(react-context-selector): avoid useState eager-bailout pitfall (#36002)

Author: layershifter

change/@fluentui-react-context-selector-3a434472-dd03-454e-82d2-431c6a6b7366.json

@@ -0,0 +1,7 @@
+{
+  "type": "patch",
+  "comment": "fix: rewrite useContextSelector to avoid React's eager-bailout pitfall on memoized consumers",
+  "packageName": "@fluentui/react-context-selector",
+  "email": "olfedias@microsoft.com",
+  "dependentChangeType": "patch"
+}

packages/react-components/react-context-selector/etc/react-context-selector.api.md

@@ -17,24 +17,18 @@ export type ContextSelector<Value, SelectedValue> = (value: Value) => SelectedVa
 
 // @internal (undocumented)
 export type ContextValue<Value> = {
-    listeners: ((payload: readonly [ContextVersion, Value]) => void)[];
-    value: React_2.MutableRefObject<Value>;
-    version: React_2.MutableRefObject<ContextVersion>;
+    listeners: ((payload: Value) => void)[];
+    value: {
+        current: Value;
+    };
+    isDefault?: boolean;
 };
 
-// @internal (undocumented)
-export type ContextValues<Value> = ContextValue<Value> & {
-    listeners: ((payload: readonly [ContextVersion, Record<string, Value>]) => void)[];
-};
-
-// @internal (undocumented)
-export type ContextVersion = number;
-
 // @internal (undocumented)
 export const createContext: <Value>(defaultValue: Value) => Context<Value>;
 
 // @internal
-export const useContextSelector: <Value, SelectedValue>(context: Context<Value>, selector: ContextSelector<Value, SelectedValue>) => SelectedValue;
+export const useContextSelector: <Value, SelectedValue>(context: Context<Value>, selectorFn: ContextSelector<Value, SelectedValue>) => SelectedValue;
 
 // @internal
 export function useHasParentContext<Value>(context: Context<Value>): boolean;

packages/react-components/react-context-selector/src/createContext.test.tsx

@@ -1,9 +1,10 @@
 import { createContext } from './createContext';
-import * as ReactIs from 'react-is';
+import { isValidElementType } from 'react-is';
 
 describe('createContext', () => {
   it('creates a Provider component', () => {
     const Context = createContext(null);
-    expect(ReactIs.isValidElementType(Context.Provider)).toBeTruthy();
+
+    expect(isValidElementType(Context.Provider)).toBeTruthy();
   });
 });

packages/react-components/react-context-selector/src/createContext.ts

@@ -10,27 +10,23 @@ const createProvider = <Value>(Original: React.Provider<ContextValue<Value>>) =>
   const Provider: React.FC<React.ProviderProps<Value>> = props => {
     // Holds an actual "props.value"
     const valueRef = React.useRef(props.value);
-    // Used to sync context updates and avoid stale values, can be considered as render/effect counter of Provider.
-    const versionRef = React.useRef(0);
 
     // A stable object, is used to avoid context updates via mutation of its values.
     const contextValue = React.useRef<ContextValue<Value>>(null);
 
     if (!contextValue.current) {
       contextValue.current = {
         value: valueRef,
-        version: versionRef,
         listeners: [],
       };
     }
 
     useIsomorphicLayoutEffect(() => {
       valueRef.current = props.value;
-      versionRef.current += 1;
 
       runWithPriority(NormalPriority, () => {
         (contextValue.current as ContextValue<Value>).listeners.forEach(listener => {
-          listener([versionRef.current, props.value]);
+          listener(props.value);
         });
       });
     }, [props.value]);
@@ -53,8 +49,8 @@ export const createContext = <Value>(defaultValue: Value): Context<Value> => {
   // eslint-disable-next-line @fluentui/no-context-default-value
   const context = React.createContext<ContextValue<Value>>({
     value: { current: defaultValue },
-    version: { current: -1 },
     listeners: [],
+    isDefault: true,
   });
 
   context.Provider = createProvider<Value>(context.Provider);

packages/react-components/react-context-selector/src/index.ts

@@ -2,4 +2,4 @@ export { createContext } from './createContext';
 export { useContextSelector } from './useContextSelector';
 export { useHasParentContext } from './useHasParentContext';
 // eslint-disable-next-line @fluentui/ban-context-export
-export type { Context, ContextSelector, ContextValue, ContextValues, ContextVersion } from './types';
+export type { Context, ContextSelector, ContextValue } from './types';

packages/react-components/react-context-selector/src/types.ts

@@ -10,31 +10,16 @@ export type Context<Value> = React.Context<Value> & {
 
 export type ContextSelector<Value, SelectedValue> = (value: Value) => SelectedValue;
 
-/**
- * @internal
- */
-export type ContextVersion = number;
-
 /**
  * @internal
  */
 export type ContextValue<Value> = {
   /** Holds a set of subscribers from components. */
-  listeners: ((payload: readonly [ContextVersion, Value]) => void)[];
+  listeners: ((payload: Value) => void)[];
 
   /** Holds an actual value of React's context that will be propagated down for computations. */
-  value: // eslint-disable-next-line @typescript-eslint/no-deprecated
-  React.MutableRefObject<Value>;
+  value: { current: Value };
 
-  /** A version field is used to sync a context value and consumers. */
-  version: // eslint-disable-next-line @typescript-eslint/no-deprecated
-  React.MutableRefObject<ContextVersion>;
-};
-
-/**
- * @internal
- */
-export type ContextValues<Value> = ContextValue<Value> & {
-  /** List of listners to publish changes */
-  listeners: ((payload: readonly [ContextVersion, Record<string, Value>]) => void)[];
+  /** Indicates if a context holds default value. */
+  isDefault?: boolean;
 };

packages/react-components/react-context-selector/src/useContextSelector.test.tsx

@@ -98,4 +98,99 @@ describe('useContextSelector', () => {
     expect(document.querySelector<HTMLElement>('.test-component')?.dataset.active).toBe('true');
     expect(onUpdate).toHaveBeenCalledTimes(2);
   });
+
+  it('memoized consumers re-render only when their selected slice changes', () => {
+    const MemoizedTestComponent = React.memo(TestComponent);
+
+    const onUpdates: Record<number, jest.Mock> = {
+      1: jest.fn(),
+      2: jest.fn(),
+      3: jest.fn(),
+      4: jest.fn(),
+    };
+
+    render(
+      <TestProvider>
+        <MemoizedTestComponent index={1} onUpdate={onUpdates[1]} />
+        <MemoizedTestComponent index={2} onUpdate={onUpdates[2]} />
+        <MemoizedTestComponent index={3} onUpdate={onUpdates[3]} />
+        <MemoizedTestComponent index={4} onUpdate={onUpdates[4]} />
+      </TestProvider>,
+      { container: container as HTMLElement },
+    );
+
+    // Initial render. TestProvider starts at index=0, so none are active. Each memoized item has committed once — the
+    // `onUpdate` effect fires once per item on mount.
+    Object.values(onUpdates).forEach(m => expect(m).toHaveBeenCalledTimes(1));
+
+    // Click 1: 0 → 1. Only index=1 flips (active: false → true).
+    jest.clearAllMocks();
+    act(() => {
+      document.querySelector<HTMLElement>('.test-provider')?.click();
+    });
+
+    expect(onUpdates[1]).toHaveBeenCalledTimes(1); // true => false
+    expect(onUpdates[2]).toHaveBeenCalledTimes(0);
+    expect(onUpdates[3]).toHaveBeenCalledTimes(0);
+    expect(onUpdates[4]).toHaveBeenCalledTimes(0);
+
+    // Click 2: 1 → 2. index=1 flips (active: true → false), index=2 flips (active: false → true).
+    // Items 3 and 4 did not change and must not re-render.
+    jest.clearAllMocks();
+    act(() => {
+      document.querySelector<HTMLElement>('.test-provider')?.click();
+    });
+
+    expect(onUpdates[1]).toHaveBeenCalledTimes(1); // true => false
+    expect(onUpdates[2]).toHaveBeenCalledTimes(1); // false => true
+    expect(onUpdates[3]).toHaveBeenCalledTimes(0); // ← was 2 under the old eager-bailout with `useState()`
+    expect(onUpdates[4]).toHaveBeenCalledTimes(0);
+
+    // Click 3: 2 → 3.
+    jest.clearAllMocks();
+    act(() => {
+      document.querySelector<HTMLElement>('.test-provider')?.click();
+    });
+
+    expect(onUpdates[1]).toHaveBeenCalledTimes(0);
+    expect(onUpdates[2]).toHaveBeenCalledTimes(1); // true => false
+    expect(onUpdates[3]).toHaveBeenCalledTimes(1); // false => true
+    expect(onUpdates[4]).toHaveBeenCalledTimes(0);
+  });
+
+  it('a single consumer throw does not starve later subscribers of context updates', () => {
+    const ThrowingConsumer: React.FC<{ threshold: number }> = ({ threshold }) => {
+      useContextSelector(TestContext, v => {
+        if (v.index > threshold) {
+          throw new Error('selector cannot handle this value');
+        }
+        return v.index;
+      });
+
+      return null;
+    };
+
+    const onUpdate = jest.fn();
+
+    render(
+      <TestProvider>
+        <ThrowingConsumer threshold={0} />
+        <TestComponent index={1} onUpdate={onUpdate} />
+      </TestProvider>,
+      { container: container as HTMLElement },
+    );
+
+    expect(onUpdate).toHaveBeenCalledTimes(1);
+
+    // Click: index goes 0 → 1. The first consumer's selector will throw
+    // (1 > 0). The downstream TestComponent must still receive the update
+    // because listener error isolation prevents `listeners.forEach` from
+    // short-circuiting.
+    act(() => {
+      document.querySelector<HTMLElement>('.test-provider')?.click();
+    });
+
+    expect(document.querySelector<HTMLElement>('.test-component')?.dataset.active).toBe('true');
+    expect(onUpdate).toHaveBeenCalledTimes(2);
+  });
 });

packages/react-components/react-context-selector/src/useContextSelector.ts

@@ -1,9 +1,9 @@
 'use client';
 
-import { useEventCallback, useIsomorphicLayoutEffect } from '@fluentui/react-utilities';
+import { useIsomorphicLayoutEffect } from '@fluentui/react-utilities';
 import * as React from 'react';
 
-import type { Context, ContextSelector, ContextValue, ContextVersion } from './types';
+import type { Context, ContextSelector, ContextValue } from './types';
 
 /**
  * This hook returns context selected value by selector.
@@ -14,74 +14,56 @@ import type { Context, ContextSelector, ContextValue, ContextVersion } from './t
  */
 export const useContextSelector = <Value, SelectedValue>(
   context: Context<Value>,
-  selector: ContextSelector<Value, SelectedValue>,
+  selectorFn: ContextSelector<Value, SelectedValue>,
 ): SelectedValue => {
   const contextValue = React.useContext(context as unknown as Context<ContextValue<Value>>);
+  const { value: valueRef, listeners } = contextValue;
 
-  const {
-    value: { current: value },
-    version: { current: version },
-    listeners,
-  } = contextValue;
-  const selected = selector(value);
+  // Read valueRef during render and return selector(value) directly. This is analogous to `useSyncExternalStore`'s
+  // `getSnapshot` and is the only way to select a slice from a shared ref-based store without re-rendering every
+  // consumer on every provider update.
+  const valueAtRender = selectorFn(valueRef.current);
+  const [, forceUpdate] = React.useReducer((x: number) => x + 1, 0);
 
-  const [state, setState] = React.useState<readonly [Value, SelectedValue]>([value, selected]);
-  const dispatch = (
-    payload:
-      | undefined // undefined from render below
-      | readonly [ContextVersion, Value], // from provider effect
-  ) => {
-    setState(prevState => {
-      if (!payload) {
-        // early bail out when is dispatched during render
-        return [value, selected] as const;
-      }
-
-      if (payload[0] <= version) {
-        if (Object.is(prevState[1], selected)) {
-          return prevState; // bail out
-        }
+  // Refs holding the current selector and the most-recently-returned slice.
+  // Updated in a layout effect (ordering: children first, then provider) so
+  // they are current by the time the provider's listener loop fires.
+  const selectorFnRef = React.useRef<ContextSelector<Value, SelectedValue>>(selectorFn);
+  const lastValueAtRender = React.useRef<SelectedValue>(valueAtRender);
 
-        return [value, selected] as const;
-      }
+  useIsomorphicLayoutEffect(() => {
+    selectorFnRef.current = selectorFn;
+    lastValueAtRender.current = valueAtRender;
+  });
 
+  useIsomorphicLayoutEffect(() => {
+    const listener = (payload: Value) => {
+      // Selectors can throw on transiently-inconsistent inputs (stale props vs. newer context value). Swallow so a
+      // single consumer's throw doesn't abort the provider's `listeners.forEach`.
       try {
-        if (Object.is(prevState[0], payload[1])) {
-          return prevState; // do not update
-        }
-
-        const nextSelected = selector(payload[1]);
+        const nextSelectedValue = selectorFnRef.current(payload);
 
-        if (Object.is(prevState[1], nextSelected)) {
-          return prevState; // do not update
+        if (!Object.is(lastValueAtRender.current, nextSelectedValue)) {
+          forceUpdate();
         }
-
-        return [payload[1], nextSelected] as const;
-      } catch (e) {
-        // ignored (stale props or some other reason)
+      } catch {
+        // ignored (stale props or similar — heals on the next parent-driven render)
       }
+    };
 
-      // explicitly spread to enforce typing
-      return [prevState[0], prevState[1]] as const; // schedule update
-    });
-  };
-
-  if (!Object.is(state[1], selected)) {
-    // schedule re-render
-    // this is safe because it's self contained
-    dispatch(undefined);
-  }
+    listeners.push(listener);
 
-  const stableDispatch = useEventCallback(dispatch);
-
-  useIsomorphicLayoutEffect(() => {
-    listeners.push(stableDispatch);
+    // Effect-fixup: catch updates that occurred between render and effect run (Relay's useFragmentInternal pattern).
+    listener(valueRef.current);
 
     return () => {
-      const index = listeners.indexOf(stableDispatch);
-      listeners.splice(index, 1);
+      const index = listeners.indexOf(listener);
+
+      if (index !== -1) {
+        listeners.splice(index, 1);
+      }
     };
-  }, [stableDispatch, listeners]);
+  }, [listeners, valueRef]);
 
-  return state[1] as SelectedValue;
+  return valueAtRender;
 };

packages/react-components/react-context-selector/src/useHasParentContext.ts

@@ -14,9 +14,5 @@ import type { Context, ContextValue } from './types';
 export function useHasParentContext<Value>(context: Context<Value>): boolean {
   const contextValue = React.useContext(context as unknown as Context<ContextValue<Value>>);
 
-  if (contextValue.version) {
-    return contextValue.version.current !== -1;
-  }
-
-  return false;
+  return !contextValue.isDefault;
 }