fix calling set() in callbacks potentially causing multiple updates

integration-tests/create-state/state-modifications.test.ts

@@ -202,6 +202,42 @@ describe("state modifications", () => {
       expect(quadrupled.get()).toBe(8);
     });
 
+    test("coalesces derived notifications when a subscriber updates another source", () => {
+      const source = new StateCore(1);
+      const trigger = new StateCore(false);
+      const doubled = source.map((value) => value * 2);
+      const combined = trigger.combine(doubled);
+
+      expect(combined.get()).toEqual([false, 2]);
+
+      const emissions: Array<[boolean, number]> = [];
+      source.on(() => trigger.set(true));
+      combined.on((value) => emissions.push(value as [boolean, number]));
+
+      source.set(2);
+
+      expect(emissions).toEqual([[true, 4]]);
+      expect(combined.get()).toEqual([true, 4]);
+    });
+
+    test("does not deliver an older derived notification after a reentrant update", () => {
+      const source = new StateCore(1);
+      const doubled = source.map((value) => value * 2);
+
+      expect(doubled.get()).toBe(2);
+
+      const emissions: number[] = [];
+      source.on((value) => {
+        if (value === 2) source.set(3);
+      });
+      doubled.on((value) => emissions.push(value));
+
+      source.set(2);
+
+      expect(emissions).toEqual([6]);
+      expect(doubled.get()).toBe(6);
+    });
+
     test("does not emit an intermediate value in a deep, uneven diamond graph", () => {
       const source = new StateCore(1);
       const short = source.map((value) => value + 100);

src/create-state/state-core.ts

@@ -1,4 +1,7 @@
 import { MinHeap } from "./min-heap";
+import { runWithStateScheduler } from "./state-scheduler";
+
+import type { StateScheduler } from "./state-scheduler";
 
 // explicit symbol for empty value, so it is 100% unique
 const emptyValue = Symbol("veles-state-core-empty");
@@ -22,8 +25,6 @@ type CoreOptions<T> = {
 
 const defaultEquality: EqualityFn<any> = (value1, value2) => value1 === value2;
 
-type PendingNotification = { node: StateCore<any>; value: any; prevValue: any };
-
 class StateCore<T> {
   private _value: CoreValue<T>;
   private _prevValue: CoreValue<T> = emptyValue;
@@ -66,24 +67,20 @@ class StateCore<T> {
   set(newValue: T) {
     if (this._equality(this._value, newValue)) return;
 
-    const prevValue = this._value;
-    this._prevValue = prevValue;
-    this._value = newValue;
-
-    this._children.forEach((child) => {
-      child._dirty = true;
-    });
+    runWithStateScheduler((scheduler) => {
+      const prevValue = this._value;
+      this._prevValue = prevValue;
+      this._value = newValue;
 
-    const pendingNotifications = this.flush();
+      // Schedule the directly updated state before its derived children. The
+      // scheduler only starts notification after the graph has been recomputed.
+      this.scheduleNotification(scheduler, newValue, prevValue);
 
-    // technically, we can notify subscribers earlier, but then synchronous reading
-    // would return stale state values.
-    this.notifySubscribers(newValue, prevValue);
+      this._children.forEach((child) => {
+        child._dirty = true;
+      });
 
-    // we intentionally process derived notifications after direct subscribers to
-    // the signal
-    pendingNotifications.forEach(({ node, value, prevValue }) => {
-      node.notifySubscribers(value, prevValue);
+      this.flush(scheduler);
     });
   }
 
@@ -216,9 +213,7 @@ class StateCore<T> {
     return { changed };
   }
 
-  private flush(): PendingNotification[] {
-    const pendingNotifications: PendingNotification[] = [];
-
+  private flush(scheduler: StateScheduler) {
     /**
      * We utilize a min-heap to ensure that we don't trigger notifications
      * before all the previous necessary work is done. Otherwise it can
@@ -269,18 +264,28 @@ class StateCore<T> {
 
       if (!changed) continue;
 
-      pendingNotifications.push({
-        node: child,
-        prevValue: child._prevValue,
-        value,
-      });
+      child.scheduleNotification(scheduler, value, child._prevValue);
       child._children.forEach((grandchild) => {
         grandchild._dirty = true;
         enqueue(grandchild);
       });
     }
+  }
 
-    return pendingNotifications;
+  private scheduleNotification(
+    scheduler: StateScheduler,
+    value: CoreValue<T>,
+    previousValue: CoreValue<T>,
+  ) {
+    scheduler.scheduleNotification({
+      key: this,
+      rank: this._rank,
+      value,
+      previousValue,
+      notify: (notificationValue, notificationPreviousValue) => {
+        this.notifySubscribers(notificationValue, notificationPreviousValue);
+      },
+    });
   }
 
   private notifySubscribers(value: CoreValue<T>, prevValue: CoreValue<T>) {

src/create-state/state-scheduler.ts

@@ -0,0 +1,78 @@
+import { MinHeap } from "./min-heap";
+
+type PendingNotification = {
+  key: object;
+  rank: number;
+  value: any;
+  previousValue: any;
+  notify: (value: any, previousValue: any) => void;
+};
+
+/**
+ * Coordinates notifications produced by one synchronous state transaction.
+ *
+ * A subscriber can synchronously update another state. That nested update adds
+ * its notifications to this same scheduler, allowing an older notification for
+ * the same state to be replaced before it is delivered.
+ */
+class StateScheduler {
+  private _notifications = new MinHeap<PendingNotification>(
+    (notification1, notification2) => notification1.rank - notification2.rank,
+  );
+  private _notificationsByKey = new Map<object, PendingNotification>();
+
+  scheduleNotification(notification: PendingNotification) {
+    const existingNotification = this._notificationsByKey.get(notification.key);
+
+    if (existingNotification) {
+      // Keep the previous value from the first change in this transaction, but
+      // deliver only the latest value produced before notification.
+      existingNotification.value = notification.value;
+      return;
+    }
+
+    this._notificationsByKey.set(notification.key, notification);
+    this._notifications.push(notification);
+  }
+
+  deliverNotifications() {
+    while (this._notifications.size > 0) {
+      const notification = this._notifications.pop()!;
+
+      // A notification can only be replaced in place today. Keeping this check
+      // makes stale heap entries safe if cancellation/replacement is added later.
+      if (this._notificationsByKey.get(notification.key) !== notification) {
+        continue;
+      }
+
+      this._notificationsByKey.delete(notification.key);
+      notification.notify(notification.value, notification.previousValue);
+    }
+  }
+}
+
+let activeScheduler: StateScheduler | undefined;
+
+/**
+ * Top-level state updates own the scheduler and synchronously drain it. Updates
+ * made by subscribers reuse the active scheduler, so they cannot deliver a
+ * nested notification batch ahead of notifications already waiting to run.
+ */
+function runWithStateScheduler(runUpdate: (scheduler: StateScheduler) => void) {
+  if (activeScheduler) {
+    runUpdate(activeScheduler);
+    return;
+  }
+
+  const scheduler = new StateScheduler();
+  activeScheduler = scheduler;
+
+  try {
+    runUpdate(scheduler);
+    scheduler.deliverNotifications();
+  } finally {
+    activeScheduler = undefined;
+  }
+}
+
+export { StateScheduler, runWithStateScheduler };