How to save to context in custom hook on unmount in React?

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>
  );
}

I tried your useSpecialState hook and It seems to work fine as expected.
From what I understand that you needed to save the last state value in the context when the component un-mounts.
I created a child component that increment the counter and another toggle button which mounts and un-mounts that CounterComponent while persisting the counter value from the last un-mount.

Here is how I tried it. Working codesandbox

import React, {
useState,
useRef,
useEffect,
useContext,
createContext
} from "react";
import "./styles.css";
const StateStoreContext = createContext({
stateStore: {},
setStateStore: () => {}
});
function StateStoreComponent(props) {
const [stateStore, setStateStore] = useState({});
return (
<StateStoreContext.Provider value={{ stateStore, setStateStore }}>
{props.children}
</StateStoreContext.Provider>
);
}
function useSpecialState(defaultValue, key) {
const { stateStore, setStateStore } = useContext(StateStoreContext);
const [state, setState] = useState(
stateStore[key] !== undefined ? stateStore[key] : defaultValue
);
const stateRef = useRef(state);
useEffect(() => {
console.log("child state change ---> ");
stateRef.current = state;
}, [state]);
useEffect(() => {
console.log("mount ---> ");
return () => {
console.log("un-mount <--- ");
console.log("prevStateStore ", stateStore);
console.log("persistState :: ", key , stateRef.current);
setStateStore((prevStateStore) => ({
...prevStateStore,
[key]: stateRef.current
}));
};
}, []);
return [state, setState];
}
const CounterComponent = () => {
const [counter, setCounter] = useSpecialState(0, "counter");
return (
<>
<h1>Hello CodeSandbox Counter: {counter}</h1>
<button onClick={() => setCounter((prev) => prev + 1)}>increment</button>
<br />
<br />
<br />
</>
);
};
export default function App() {
const [visible, setVisible] = useSpecialState(true, "visible");
return (
<div className="App">
<StateStoreComponent>
{visible && <CounterComponent />}
<button onClick={() => setVisible(!visible)}>toggle</button>
</StateStoreComponent>
</div>
);
}