refactor(utils): extract useReduxConnection for less repetition

src/utils/redux-extension/createReduxConnection.ts

@@ -5,9 +5,6 @@ import { ReduxExtension } from './getReduxExtension';
 type ConnectResponse = ReturnType<NonNullable<ReduxExtension>['connect']>;
 
 export type Connection = {
-  /** Mark the connection as not initiated, so it can be initiated before using it. */
-  shouldInit?: boolean;
-
   /** Initiate the connection and add it to the extension connections.
    *  Should only be executed once in the live time of the connection.
    */
@@ -40,7 +37,5 @@ export const createReduxConnection = (
   if (!extension) return undefined;
   const connection = extension.connect({ name });
 
-  return Object.assign(connection, {
-    shouldInit: true,
-  }) as Connection;
+  return connection as Connection;
 };

src/utils/redux-extension/useReduxConnection.ts

@@ -0,0 +1,34 @@
+import { useEffect, useRef } from 'react';
+import { Connection, createReduxConnection } from './createReduxConnection';
+import { getReduxExtension } from './getReduxExtension';
+
+interface ConnectionConfig<T> {
+  name: string;
+  enabled: boolean | undefined;
+  initialValue: T;
+  disconnectAllOnCleanup?: boolean;
+}
+
+export const useReduxConnection = <T>({
+  name,
+  initialValue,
+  enabled,
+  disconnectAllOnCleanup,
+}: ConnectionConfig<T>) => {
+  const connectionRef = useRef<Connection>();
+  const firstValue = useRef(initialValue);
+
+  useEffect(() => {
+    const extension = getReduxExtension(enabled);
+
+    const connection = createReduxConnection(extension, name);
+    connection?.init(firstValue.current);
+    connectionRef.current = connection;
+
+    return () => {
+      if (disconnectAllOnCleanup) extension?.disconnect?.();
+    };
+  }, [name, enabled, disconnectAllOnCleanup]);
+
+  return connectionRef;
+};

src/utils/useAtomDevtools.ts

@@ -1,11 +1,8 @@
 import { useEffect, useRef } from 'react';
 import { useAtom } from 'jotai/react';
 import type { Atom, WritableAtom } from 'jotai/vanilla';
-import {
-  Connection,
-  createReduxConnection,
-} from './redux-extension/createReduxConnection';
-import { getReduxExtension } from './redux-extension/getReduxExtension';
+import { useReduxConnection } from './redux-extension/useReduxConnection';
+import { useDidMount } from './useDidMount';
 
 type DevtoolOptions = Parameters<typeof useAtom>[1] & {
   name?: string;
@@ -17,21 +14,24 @@ export function useAtomDevtools<Value, Result>(
   options?: DevtoolOptions,
 ): void {
   const { enabled, name } = options || {};
-
-  const extension = getReduxExtension(enabled);
+  const didMount = useDidMount();
 
   const [value, setValue] = useAtom(anAtom, options);
 
   const lastValue = useRef(value);
   const isTimeTraveling = useRef(false);
-  const devtools = useRef<Connection>();
 
   const atomName = name || anAtom.debugLabel || anAtom.toString();
 
+  const connection = useReduxConnection({
+    name: atomName,
+    enabled,
+    initialValue: value,
+  });
+
   useEffect(() => {
-    if (!extension) {
-      return;
-    }
+    if (!connection.current) return;
+
     const setValueIfWritable = (value: Value) => {
       if (typeof setValue === 'function') {
         (setValue as (value: Value) => void)(value);
@@ -43,9 +43,7 @@ export function useAtomDevtools<Value, Result>(
       );
     };
 
-    devtools.current = createReduxConnection(extension, atomName);
-
-    const unsubscribe = devtools.current?.subscribe((message) => {
+    const unsubscribe = connection.current.subscribe((message) => {
       if (message.type === 'ACTION' && message.payload) {
         try {
           setValueIfWritable(JSON.parse(message.payload));
@@ -68,7 +66,7 @@ export function useAtomDevtools<Value, Result>(
         message.type === 'DISPATCH' &&
         message.payload?.type === 'COMMIT'
       ) {
-        devtools.current?.init(lastValue.current);
+        connection.current?.init(lastValue.current);
       } else if (
         message.type === 'DISPATCH' &&
         message.payload?.type === 'IMPORT_STATE'
@@ -78,7 +76,7 @@ export function useAtomDevtools<Value, Result>(
 
         computedStates.forEach(({ state }: { state: Value }, index: number) => {
           if (index === 0) {
-            devtools.current?.init(state);
+            connection.current?.init(state);
           } else {
             setValueIfWritable(state);
           }
@@ -87,23 +85,19 @@ export function useAtomDevtools<Value, Result>(
     });
 
     return unsubscribe;
-  }, [anAtom, extension, atomName, setValue]);
+  }, [anAtom, connection, setValue]);
 
   useEffect(() => {
-    if (!devtools.current) {
-      return;
-    }
+    if (!connection.current || !didMount) return;
+
     lastValue.current = value;
-    if (devtools.current.shouldInit) {
-      devtools.current.init(value);
-      devtools.current.shouldInit = false;
-    } else if (isTimeTraveling.current) {
+    if (isTimeTraveling.current) {
       isTimeTraveling.current = false;
     } else {
-      devtools.current.send(
+      connection.current.send(
         `${atomName} - ${new Date().toLocaleString()}` as any,
         value,
       );
     }
-  }, [anAtom, extension, atomName, value]);
+  }, [atomName, connection, didMount, value]);
 }

src/utils/useAtomsDevtools.ts

@@ -1,11 +1,8 @@
 import { useEffect, useRef } from 'react';
 import { AnyAtom, AnyAtomValue, AtomsSnapshot, Options } from '../types';
-import {
-  Connection,
-  createReduxConnection,
-} from './redux-extension/createReduxConnection';
-import { getReduxExtension } from './redux-extension/getReduxExtension';
+import { useReduxConnection } from './redux-extension/useReduxConnection';
 import { useAtomsSnapshot } from './useAtomsSnapshot';
+import { useDidMount } from './useDidMount';
 import { useGotoAtomsSnapshot } from './useGotoAtomsSnapshot';
 
 const atomToPrintable = (atom: AnyAtom) =>
@@ -35,23 +32,26 @@ export function useAtomsDevtools(
   options?: DevtoolsOptions,
 ): void {
   const { enabled } = options || {};
-
-  const extension = getReduxExtension(enabled);
+  const didMount = useDidMount();
 
   // This an exception, we don't usually use utils in themselves!
   const atomsSnapshot = useAtomsSnapshot(options);
   const goToSnapshot = useGotoAtomsSnapshot(options);
 
   const isTimeTraveling = useRef(false);
   const isRecording = useRef(true);
-  const devtools = useRef<Connection>();
-
   const snapshots = useRef<AtomsSnapshot[]>([]);
 
+  const connection = useReduxConnection({
+    name,
+    enabled,
+    initialValue: undefined,
+    disconnectAllOnCleanup: true,
+  });
+
   useEffect(() => {
-    if (!extension) {
-      return;
-    }
+    if (!connection.current) return;
+
     const getSnapshotAt = (index = snapshots.current.length - 1) => {
       // index 0 is @@INIT, so we need to return the next action (0)
       const snapshot = snapshots.current[index >= 0 ? index : 0];
@@ -61,9 +61,7 @@ export function useAtomsDevtools(
       return snapshot;
     };
 
-    devtools.current = createReduxConnection(extension, name);
-
-    const devtoolsUnsubscribe = devtools.current?.subscribe((message) => {
+    const unsubscribe = connection.current.subscribe((message) => {
       switch (message.type) {
         case 'DISPATCH':
           switch (message.payload?.type) {
@@ -72,7 +70,7 @@ export function useAtomsDevtools(
               break;
 
             case 'COMMIT':
-              devtools.current?.init(getDevtoolsState(getSnapshotAt()));
+              connection.current?.init(getDevtoolsState(getSnapshotAt()));
               snapshots.current = [];
               break;
 
@@ -89,32 +87,23 @@ export function useAtomsDevtools(
       }
     });
 
-    return () => {
-      extension?.disconnect?.();
-      devtoolsUnsubscribe?.();
-    };
-  }, [extension, goToSnapshot, name]);
+    return unsubscribe;
+  }, [connection, goToSnapshot]);
 
   useEffect(() => {
-    if (!devtools.current) {
-      return;
-    }
-    if (devtools.current.shouldInit) {
-      devtools.current.init(undefined);
-      devtools.current.shouldInit = false;
-      return;
-    }
+    if (!connection.current || !didMount) return;
+
     if (isTimeTraveling.current) {
       isTimeTraveling.current = false;
     } else if (isRecording.current) {
       snapshots.current.push(atomsSnapshot);
-      devtools.current.send(
+      connection.current.send(
         {
           type: `${snapshots.current.length}`,
           updatedAt: new Date().toLocaleString(),
         } as any,
         getDevtoolsState(atomsSnapshot),
       );
     }
-  }, [atomsSnapshot]);
+  }, [atomsSnapshot, connection, didMount]);
 }

src/utils/useDidMount.ts

@@ -0,0 +1,9 @@
+import { useEffect, useRef } from 'react';
+
+export const useDidMount = () => {
+  const didMount = useRef(false);
+  useEffect(() => {
+    didMount.current = true;
+  }, []);
+  return didMount.current;
+};