Fix deferred subscription lifecycle

Assisted-By: devx/e0620a8d-8624-4748-b9c7-ce390238fc37

Author: JoviDeCroock

packages/react/runtime/src/index.ts

@@ -59,6 +59,8 @@ interface Effect {
 	_dispose(): void;
 }
 
+type SubscriptionNode = Parameters<Signal["_subscribe"]>[0];
+
 /**
  * Use this flag to represent a bare `useSignals` call that doesn't manually
  * close its effect store and relies on auto-closing when the next useSignals is
@@ -102,38 +104,57 @@ export interface EffectStore {
 	getSnapshot(): number;
 	/** startEffect - begin tracking signals used in this component */
 	_start(): void;
-	_subscribers: Array<{ signal: Signal; node: any }>;
-	_sub: typeof Signal.prototype._subscribe;
+	_subscribers: Array<{ signal: Signal; node: SubscriptionNode }>;
 	/** finishEffect - stop tracking the signals used in this component */
 	f(): void;
 	[symDispose](): void;
 }
 
 let currentStore: EffectStore | undefined;
 
-const realSubscribe = Signal.prototype._subscribe;
 function startComponentEffect(
 	prevStore: EffectStore | undefined,
 	nextStore: EffectStore
 ) {
-	nextStore._sub = prevStore ? prevStore._sub : realSubscribe;
-	Signal.prototype._subscribe = function (this: Signal, node: any) {
-		nextStore._subscribers.push({ signal: this, node });
+	const previousSubscribe = Signal.prototype._subscribe;
+	Signal.prototype._subscribe = function (
+		this: Signal,
+		node: SubscriptionNode
+	) {
+		if ((node as { _target: Effect })._target === nextStore.effect) {
+			nextStore._subscribers.push({ signal: this, node });
+		}
 	};
-	const endEffect = nextStore.effect._start();
+
+	let endEffect: () => void;
+	try {
+		endEffect = nextStore.effect._start();
+	} catch (error) {
+		Signal.prototype._subscribe = previousSubscribe;
+		throw error;
+	}
 	currentStore = nextStore;
 
-	return finishComponentEffect.bind(nextStore, prevStore, endEffect);
+	return finishComponentEffect.bind(
+		nextStore,
+		prevStore,
+		endEffect,
+		previousSubscribe
+	);
 }
 
 function finishComponentEffect(
 	this: EffectStore,
 	prevStore: EffectStore | undefined,
-	endEffect: () => void
+	endEffect: () => void,
+	previousSubscribe: typeof Signal.prototype._subscribe
 ) {
-	Signal.prototype._subscribe = prevStore ? prevStore._sub : realSubscribe;
-	endEffect();
-	currentStore = prevStore;
+	Signal.prototype._subscribe = previousSubscribe;
+	try {
+		endEffect();
+	} finally {
+		currentStore = prevStore;
+	}
 }
 
 /**
@@ -180,15 +201,15 @@ function createEffectStore(
 		if (onChangeNotifyReact) onChangeNotifyReact();
 	};
 
-	return {
+	const store: EffectStore = {
 		_usage,
 		_subscribers: [],
-		_sub: realSubscribe,
 		effect: effectInstance,
 		subscribe(onStoreChange) {
 			onChangeNotifyReact = onStoreChange;
 
 			return function () {
+				store._subscribers = [];
 				/**
 				 * Rotate to next version when unsubscribing to ensure that components are re-run
 				 * when subscribing again.
@@ -310,16 +331,18 @@ function createEffectStore(
 		},
 		[symDispose]() {
 			this.f();
+			this._subscribers = [];
 		},
 	};
+
+	return store;
 }
 
 const noop = () => {};
 
 function createEmptyEffectStore(): EffectStore {
 	return {
 		_subscribers: [],
-		_sub: noop as any,
 		_usage: UNMANAGED,
 		effect: {
 			_sources: undefined,
@@ -385,10 +408,17 @@ export function _useSignalsImplementation(
 	if (_usage === UNMANAGED) useIsomorphicLayoutEffect(cleanupTrailingStore);
 
 	useIsomorphicLayoutEffect(() => {
-		store._subscribers.forEach(({ signal, node }) => {
-			realSubscribe.call(signal, node);
-		});
+		const subscribers = store._subscribers;
 		store._subscribers = [];
+
+		subscribers.forEach(({ signal, node }) => {
+			if (
+				(node as { _target: Effect })._target === store.effect &&
+				(node as { _version: number })._version !== -1
+			) {
+				signal._subscribe(node);
+			}
+		});
 	});
 
 	return store;