react 18 changes - getting registerStateContainer from context

src/ContainerRoot/ContainerRoot.tsx

@@ -26,13 +26,12 @@ import React, { PropsWithChildren, useEffect, useRef, useState } from "react";
 import {
   Action,
   getUniqueId,
-  registerStateContainer,
+  registerStateContainer as libRegisterStateContainer,
   unregisterStateContainer,
 } from "key-value-state-container";
 
 import { ContainerRootContext } from "./context";
 import { ContainerRootProps } from "./props";
-import { reRegister } from "./re-register";
 
 export const ContainerRoot = <TState extends Object, TAction extends Action>(
   props: PropsWithChildren<ContainerRootProps<TState, TAction>>
@@ -48,10 +47,9 @@ export const ContainerRoot = <TState extends Object, TAction extends Action>(
     persistence,
     reducer,
   } = props;
-  const runningFirstTimeRef = useRef<boolean>(true);
-  if (runningFirstTimeRef.current) {
-    runningFirstTimeRef.current = false;
-    registerStateContainer({
+
+  const registerStateContainer = () => {
+    libRegisterStateContainer({
       autoActions,
       autoState,
       config,
@@ -60,36 +58,51 @@ export const ContainerRoot = <TState extends Object, TAction extends Action>(
       persistence,
       reducer,
     });
-  }
+  };
 
-  useEffect(() => {
+  const runningFirstTimeRef = useRef<boolean>(true);
+  if (runningFirstTimeRef.current) {
+    runningFirstTimeRef.current = false;
     /**
-     * Checking for unregistered container
-     * Container might be unregistered due to `useEffect` return function
-     * called twice in React 18 Strict Mode
+     * We need to register the container here, with the first render. 
+     * 
+     * Assuming the `console.log()` invocations are printed out for 
+     * mount/unmount component lifecycle events (by `useEffect()`),
+     * the printout for React 18 Strict Mode would look as follows:
+     *
+     * useSelector mounted ⭐️
+     * ContainerRoot mounted ⭐️⭐️
+     * useSelector unmounted
+     * ContainerRoot unmounted
+     * useSelector mounted ⭐️⭐️⭐️
+     * ContainerRoot mounted
+     * 
+     * Please note is is way too late to register the container in the ⭐️⭐️
+     * (as there was a `useSelector` call in ⭐️, that wanted to register 
+     * callback listeners with an unregistered container!)
+     *
+     * The same story goes with `useSelector`, marked by ⭐️⭐️⭐️,
+     * that needs a container to be registered.
+     * 
+     * Keep in mind the biggest possible memory issues for a container 
+     * are the callback listeners mentioned above, 
+     * as these can easily allocate GBs of memory
+     * (btw. this problem is handled by the `useSelector` code)
      */
-    if (reRegister({ containerId })) {
-      registerStateContainer({
-        autoActions,
-        autoState,
-        config,
-        containerId,
-        initialState,
-        persistence,
-        reducer,
-      });
-    }
+    registerStateContainer();
+  }
+
+  useEffect(() => {
     return () => {
-      unregisterStateContainer({
-        containerId,
-      });
+      unregisterStateContainer({ containerId });
     };
   }, []);
 
   return (
     <ContainerRootContext.Provider
       value={{
         containerId,
+        registerStateContainer,
       }}
     >
       {children}

src/ContainerRoot/context.ts

@@ -27,10 +27,12 @@ import { getUniqueId } from "key-value-state-container";
 
 export interface ContainerRootContextProps {
   containerId: string;
+  registerStateContainer: () => void;
 }
 
 export const ContainerRootContext = React.createContext<
   ContainerRootContextProps
 >({
   containerId: getUniqueId(),
+  registerStateContainer: () => { },
 });

src/use-get-context.ts

@@ -0,0 +1,41 @@
+/**
+  The MIT License (MIT)
+
+  Copyright Tomasz Szatkowski and WealthArc https://www.wealtharc.com (c) 2023 
+
+  Permission is hereby granted, free of charge, to any person obtaining a copy
+  of this software and associated documentation files (the "Software"), to deal
+  in the Software without restriction, including without limitation the rights
+  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+  copies of the Software, and to permit persons to whom the Software is
+  furnished to do so, subject to the following conditions:
+
+  The above copyright notice and this permission notice shall be included in all
+  copies or substantial portions of the Software.
+
+  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+  SOFTWARE.
+*/
+
+import { useContext } from "react";
+
+import {
+  ContainerRootContext,
+  ContainerRootContextProps,
+} from "./ContainerRoot/context";
+
+type Result = {
+  containerIdFromContext: string;
+  registerStateContainer: () => void;
+}
+
+export const useGetContext = (): Result => {
+  const { containerId, registerStateContainer } =
+    useContext<ContainerRootContextProps>(ContainerRootContext);
+  return { containerIdFromContext: containerId, registerStateContainer };
+};

src/use-selector.ts

@@ -2,14 +2,15 @@ import { useEffect, useRef, useState } from "react";
 import {
   Action,
   ClientNotificationCallbackArgs,
+  containerRegistered,
   getContainer,
   getUniqueId,
   registerStateChangedCallback,
   TKnownStatePath,
   unregisterStateChangedCallback,
 } from "key-value-state-container";
-import { useGetContainerId } from "./use-get-container-id";
 import { RendersWithContainerId } from "./types/contracts";
+import { useGetContext } from "./use-get-context";
 
 interface Args<TState extends Object, TAction extends Action = Action>
   extends Partial<RendersWithContainerId> {
@@ -67,8 +68,8 @@ export const useSelector = <
     `${listenerTag ? `${listenerTag}:` : ""}${getUniqueId()}`
   );
   const unmountedRef = useRef<boolean>(false);
-  const containerFromHook = useGetContainerId();
-  const containerId = containerIdFromProps || containerFromHook;
+  const { containerIdFromContext, registerStateContainer } = useGetContext();
+  const containerId = containerIdFromProps || containerIdFromContext;
   const initialState =
     getContainer<TState>({
       containerId,
@@ -78,15 +79,24 @@ export const useSelector = <
   const lateInvoke =
     typeof lateInvokeFromProps === "boolean" ? lateInvokeFromProps : true;
   const [currentState, setCurrentState] = useState(initialState);
+
   useEffect(() => {
-    return () => {
-      unmountedRef.current = true;
-    };
+    /**
+     * Component gets mounted
+     */
+    unmountedRef.current = false;
   }, []);
+
   useEffect(() => {
     if (switchOff) {
       return;
     }
+    /**
+     * Awkward, but necessary logic, explained in ContextRoot.tsx component.
+     */
+    if (!containerRegistered({ containerId })) {
+      registerStateContainer();
+    }
     const statePaths =
       typeof statePath === "string"
         ? [statePath]
@@ -138,5 +148,14 @@ export const useSelector = <
     };
   }, [setCurrentState]);
 
+  useEffect(() => {
+    return () => {
+      /**
+       * component gets unmounted
+       */
+      unmountedRef.current = true;
+    };
+  }, []);
+
   return currentState;
 };