I have the following simplified custom hook:

function useSpecialState(defaultValue, key) {
  const { stateStore, setStateStore } = useContext(StateStoreContext);

  const [state, setState] = useState(
    stateStore[key] !== undefined ? stateStore[key] : defaultValue
  );

  const stateRef = useRef(state);

  useEffect(() => {
    stateRef.current = state;
  }, [state]);

  useEffect(() => {
    return () => {
      setStateStore((prevStateStore) => ({
        ...prevStateStore,
        [key]: stateRef.current,
      }));
    };
  }, []);

  return [state, setState];
}

The goal would be to save to a context on unmount, however, this code does not work. Putting state in the dependency array of the useEffect which is responsible for saving to context would not be a good solution, because then it would be saved on every state change, which is grossly unnecessary.

The context:

const StateStoreContext = createContext({
  stateStore: {},
  setStateStore: () => {},
});

The parent component:

function StateStoreComponent(props) {
  const [stateStore, setStateStore] = useState({});

  return (
    <StateStoreContext.Provider value={{ stateStore, setStateStore }}>
      {props. Children}
    </StateStoreContext.Provider>
  );
}

TL;DR

The code you have is fine and technically correct, the observed behavior is caused by the React.StrictMode component double-mounting components in non-production builds. In other words, the code & logic should behave as you expect in normal production builds you deploy. This is, or should be, all expected behavior.

Explanation

The code you have is fine and technically correct. The reason it appears that it is not working is because you are rendering the app within the React StrictMode component which executes additional behavior in non-production builds. Specifically in this case it's the double-mounting of components as part of React's check for Ensuring Reusable State or Fixing bugs found by re-running Effects if you prefer the current docs.

Strict Mode can also help find bugs in Effects.

Every Effect has some setup code and may have some cleanup code. Normally, React calls setup when the component mounts (is added to the screen) and calls cleanup when the component unmounts (is removed from the screen). React then calls cleanup and setup again if its dependencies changed since the last render.

When Strict Mode is on, React will also run one extra setup+cleanup cycle in development for every Effect. This may feel surprising, but it helps reveal subtle bugs that are hard to catch manually.

Any component rendered within a React.StrictMode component and using the custom useSpecialState hook will be mounted, unmounted and run the second useEffect hook's cleanup function which will update the state in the context, and then mount again the component.

Here's a small demo toggling the mounting of identical components that use the useSpecialState hook, where only one of them is mounted within a React.StrictMode component. Notice that "Component A" updates the context state each time when it is mounted and unmounted, while "Component B" updates the context state only when it unmounts.

Steps:

1. App mounts, context render 0
2. Toggle A mounted: observe mount/unmount/mount, state update A, context render 1
3. Toggle B mounted: observe mount, no state update
4. Toggle A unmounted: observe unmount, state update A, context render 2
5. Toggle B unmounted: observe unmount, state update B, context render 3

Sandbox Code:

const MyComponent = ({ label }) => {
  const [count, setCount] = useSpecialState(0, "count" + label);

  return (
    <>
      <h1>Component{label}</h1>
      <div>Count: {count}</div>
      <button type="button" onClick={() => setCount((c) => c + 1)}>
        +
      </button>
    </>
  );
};

export default function App() {
  const [mountA, setMountA] = useState(false);
  const [mountB, setMountB] = useState(false);
  return (
    <StateStoreComponent>
      <div className="App">
        <h1>Hello CodeSandbox</h1>
        <h2>Start editing to see some magic happen!</h2>

        <div>
          <button type="button" onClick={() => setMountA((mount) => !mount)}>
            {mountA ? "Unmount" : "Mount"} A
          </button>
          <button type="button" onClick={() => setMountB((mount) => !mount)}>
            {mountB ? "Unmount" : "Mount"} B
          </button>
        </div>
        <StrictMode>{mountA && <MyComponent label="A" />}</StrictMode>
        {mountB && <MyComponent label="B" />}
      </div>
    </StateStoreComponent>
  );
}