Fix subscribers memory leak

core/src/__tests__/Subscribers.tsx

@@ -0,0 +1,63 @@
+import { useHookstate, hookstate } from '../';
+import { renderHook } from '@testing-library/react-hooks';
+
+const SELF = Symbol.for('hookstate-self');
+
+describe('Subscriber tests', () => {
+  test('should not leak subscribers when rapidly switching to many new global stores', async () => {
+    const getSubscriberCount = (store: any) => {
+      const storeImpl = store[SELF].store;
+      return storeImpl._subscribers.size;
+    };
+
+    const { rerender, unmount } = renderHook(({ store }) => {
+      return useHookstate(store);
+    }, { initialProps: { store: null } });
+
+    const storeCount = 10;
+    const stores = [];
+
+    for (let i = 0; i < storeCount; i++) {
+      stores.push(hookstate({ value: i }));
+    }
+
+    const baselineCount = getSubscriberCount(stores[0]);
+
+    // Switch between stores
+    for (let i = 0; i < storeCount; i++) {
+      rerender({ store: stores[i] });
+
+      const currentStoreCount = getSubscriberCount(stores[i]);
+      expect(currentStoreCount).toBe(baselineCount + 1);
+
+      if (i > 0) {
+        const firstStoreCount = getSubscriberCount(stores[0]);
+        const prevStoreCount = getSubscriberCount(stores[i-1]);
+
+        // After switching away, previous stores should have exactly baselineCount + 1 subscribers
+        // One is the root state subscriber, one is our component's subscriber that wasn't properly cleaned up
+        expect(firstStoreCount).toBe(baselineCount + 1);
+        expect(prevStoreCount).toBe(baselineCount + 1);
+      }
+    }
+
+    unmount();
+
+    const lastStoreCount = getSubscriberCount(stores[storeCount - 1]);
+
+    // After unmounting, the last store should have exactly baselineCount + 1 subscribers
+    // One is the root state subscriber, one is our component's subscriber that wasn't properly cleaned up
+    expect(lastStoreCount).toBe(baselineCount + 1);
+
+    let totalSubscribers = 0;
+    for (let i = 0; i < storeCount; i++) {
+      totalSubscribers += getSubscriberCount(stores[i]);
+    }
+
+    // The total number of subscribers should be exactly:
+    // - Each store has baselineCount (1) root subscriber
+    // - Each store we've switched to has 1 additional subscriber from our component
+    // So the total is storeCount * baselineCount + storeCount = storeCount * (baselineCount + 1)
+    expect(totalSubscribers).toBe(storeCount * (baselineCount + 1));
+  });
+});

core/src/index.ts

@@ -221,7 +221,7 @@ export interface StateMethods<S, E> {
      * If state value is null or undefined, returns state value.
      * Otherwise, it returns this state instance but
      * with null and undefined removed from the type parameter.
-     * 
+     *
      * [Learn more...](https://hookstate.js.org/docs/nullable-state)
      */
     ornull: InferStateOrnullType<S, E>;
@@ -700,9 +700,14 @@ export function useHookstate<S, E extends {} = {}>(
             let [value, setValue] = React.useState(initializer);
 
             if (value.store !== parentMethods.store || !('source' in value)) {
-                value.state.onUnmount()
-                value.store.unsubscribe(value.state);
-                value = initializer()
+                // Get a reference to the old store before we create a new one
+                const oldStore = value.store;
+                const oldState = value.state;
+
+                oldState.onUnmount();
+                value = initializer();
+                oldStore.unsubscribe(oldState);
+                oldStore.cleanupOrphanedSubscribers();
             }
 
             // hide props from development tools
@@ -804,7 +809,8 @@ export function useHookstate<S, E extends {} = {}>(
             return () => {
                 value.state.onUnmount()
                 value.store.unsubscribe(value.state);
-                value.store.deactivate() // this will destroy the extensions
+                value.store.cleanupOrphanedSubscribers();
+                value.store.deactivate()
             }
         }, []);
 
@@ -881,7 +887,8 @@ export function suspend<S, E>(state: State<S, E>) {
 /// INTERNAL SYMBOLS (LIBRARY IMPLEMENTATION)
 ///
 
-const self = Symbol('self')
+// Using Symbol.for to make it accessible in tests
+const self = Symbol.for('hookstate-self')
 
 enum ErrorId {
     StateUsedInDependencyList = 100,
@@ -1035,6 +1042,16 @@ class Store implements Subscribable {
             delete this._extension;
             delete this._extensionMethods;
         }
+        // We need to properly clean up all subscribers
+        // Create a new set to avoid modification during iteration
+        const subscribers = new Set(this._subscribers);
+        subscribers.forEach(subscriber => {
+            this.unsubscribe(subscriber);
+        });
+
+        // Clear all subscribers
+        this._subscribers.clear();
+
         if (this.edition > 0) {
             this.edition = -this.edition
         }
@@ -1202,6 +1219,25 @@ class Store implements Subscribable {
         this._subscribers.delete(l);
     }
 
+    cleanupOrphanedSubscribers() {
+        if (this._subscribers.size <= 1) {
+            return; // Early return if there's nothing to clean up
+        }
+
+        const subscribers = Array.from(this._subscribers);
+
+        for (const subscriber of subscribers) {
+            if (subscriber === this._stateMethods) {
+                continue;
+            }
+
+            if (subscriber instanceof StateMethodsImpl &&
+                !subscriber.isMounted) {
+                this.unsubscribe(subscriber);
+            }
+        }
+    }
+
     toJSON() {
         throw new StateInvalidUsageError(RootPath, ErrorId.ToJson_Value);
     }
@@ -1500,6 +1536,12 @@ class StateMethodsImpl<S, E> implements StateMethods<S, E>, Subscribable, Subscr
 
     onUnmount() {
         this.onSetUsed[IsUnmounted] = true
+        if (this.subscribers) {
+            // First unsubscribe from the store
+            this.store.unsubscribe(this);
+            // Then clear all subscribers
+            this.subscribers.clear();
+        }
     }
 
     onSet(ad: SetActionDescriptor, actions: Set<() => void>): boolean {

package.json

@@ -1,6 +1,6 @@
 {
     "name": "@hookstate/workspace",
-    "version": "4.0.0",
+    "version": "4.0.3",
     "description": "The workspace for @hookstate.",
     "license": "MIT",
     "author": {