Fix pull scheduler dirty propagation and dependents

Author: seefeldb

packages/runner/src/scheduler.ts

@@ -115,6 +115,7 @@ export class Scheduler {
   private effects = new Set<Action>();
   private computations = new Set<Action>();
   private dependents = new WeakMap<Action, Set<Action>>();
+  private reverseDependencies = new WeakMap<Action, Set<Action>>();
   // Track which actions are effects persistently (survives unsubscribe/re-subscribe)
   private isEffectAction = new WeakMap<Action, boolean>();
   private dirty = new Set<Action>();
@@ -319,6 +320,18 @@ export class Scheduler {
     this.cancels.delete(action);
     this.dependencies.delete(action);
     this.pending.delete(action);
+    const dependencies = this.reverseDependencies.get(action);
+    if (dependencies) {
+      for (const dependency of dependencies) {
+        const dependents = this.dependents.get(dependency);
+        dependents?.delete(action);
+        if (dependents && dependents.size === 0) {
+          this.dependents.delete(dependency);
+        }
+      }
+      this.reverseDependencies.delete(action);
+    }
+    this.dependents.delete(action);
     // Clean up effect/computation tracking
     this.effects.delete(action);
     this.computations.delete(action);
@@ -632,7 +645,20 @@ export class Scheduler {
    * For each action that writes to paths this action reads, add this action as a dependent.
    */
   private updateDependents(action: Action, log: ReactivityLog): void {
+    const previousDependencies = this.reverseDependencies.get(action);
+    if (previousDependencies) {
+      for (const dependency of previousDependencies) {
+        const dependents = this.dependents.get(dependency);
+        dependents?.delete(action);
+        if (dependents && dependents.size === 0) {
+          this.dependents.delete(dependency);
+        }
+      }
+      this.reverseDependencies.delete(action);
+    }
+
     const reads = log.reads;
+    const newDependencies = new Set<Action>();
 
     // For each read of the new action, find other actions that write to it
     for (const read of reads) {
@@ -658,11 +684,16 @@ export class Scheduler {
                 this.dependents.set(otherAction, deps);
               }
               deps.add(action);
+              newDependencies.add(otherAction);
             }
           }
         }
       }
     }
+
+    if (newDependencies.size > 0) {
+      this.reverseDependencies.set(action, newDependencies);
+    }
   }
 
   /**
@@ -734,6 +765,8 @@ export class Scheduler {
     if (this.dirty.has(action)) return; // Already dirty, avoid infinite recursion
 
     this.dirty.add(action);
+    // Treat dirty computations as triggered so they bypass filtering
+    this.pushTriggered.add(action);
 
     // Propagate to dependents transitively
     const deps = this.dependents.get(action);
@@ -1106,6 +1139,7 @@ export class Scheduler {
     findEffects(computation);
 
     for (const effect of toSchedule) {
+      this.pushTriggered.add(effect);
       this.scheduleWithDebounce(effect);
     }
   }
@@ -1386,6 +1420,9 @@ export class Scheduler {
 
     // In pull mode, filter based on pushTriggered
     if (this.pullMode) {
+      if (this.dirty.has(action)) {
+        return false;
+      }
       // If not triggered by actual storage changes, filter it out
       if (!this.pushTriggered.has(action)) {
         return true;

packages/runner/test/scheduler.test.ts

@@ -1715,6 +1715,223 @@ describe("pull-based scheduling", () => {
     expect(effectRuns).toBeGreaterThan(initialEffectRuns);
   });
 
+  it("should recompute multi-hop chains before running effects in pull mode", async () => {
+    runtime.scheduler.enablePullMode();
+
+    const source = runtime.getCell<number>(
+      space,
+      "pull-multihop-source",
+      undefined,
+      tx,
+    );
+    source.set(1);
+    const intermediate1 = runtime.getCell<number>(
+      space,
+      "pull-multihop-mid-1",
+      undefined,
+      tx,
+    );
+    intermediate1.set(0);
+    const intermediate2 = runtime.getCell<number>(
+      space,
+      "pull-multihop-mid-2",
+      undefined,
+      tx,
+    );
+    intermediate2.set(0);
+    const effectResult = runtime.getCell<number>(
+      space,
+      "pull-multihop-effect",
+      undefined,
+      tx,
+    );
+    effectResult.set(0);
+    await tx.commit();
+    tx = runtime.edit();
+
+    let comp1Runs = 0;
+    let comp2Runs = 0;
+    let effectRuns = 0;
+
+    const computation1: Action = (actionTx) => {
+      comp1Runs++;
+      const val = source.withTx(actionTx).get();
+      intermediate1.withTx(actionTx).send(val + 1);
+    };
+
+    const computation2: Action = (actionTx) => {
+      comp2Runs++;
+      const val = intermediate1.withTx(actionTx).get();
+      intermediate2.withTx(actionTx).send(val * 2);
+    };
+
+    const effect: Action = (actionTx) => {
+      effectRuns++;
+      const val = intermediate2.withTx(actionTx).get();
+      effectResult.withTx(actionTx).send(val - 3);
+    };
+
+    runtime.scheduler.subscribe(
+      computation1,
+      {
+        reads: [source.getAsNormalizedFullLink()],
+        writes: [intermediate1.getAsNormalizedFullLink()],
+      },
+      { scheduleImmediately: true },
+    );
+    await runtime.idle();
+
+    runtime.scheduler.subscribe(
+      computation2,
+      {
+        reads: [intermediate1.getAsNormalizedFullLink()],
+        writes: [intermediate2.getAsNormalizedFullLink()],
+      },
+      { scheduleImmediately: true },
+    );
+    await runtime.idle();
+
+    runtime.scheduler.subscribe(
+      effect,
+      {
+        reads: [intermediate2.getAsNormalizedFullLink()],
+        writes: [effectResult.getAsNormalizedFullLink()],
+      },
+      { scheduleImmediately: true, isEffect: true },
+    );
+    await runtime.idle();
+
+    expect(effectResult.get()).toBe((1 + 1) * 2 - 3);
+    expect(comp2Runs).toBe(1);
+    expect(effectRuns).toBe(1);
+
+    const tx2 = runtime.edit();
+    source.withTx(tx2).send(5);
+    await tx2.commit();
+    tx = runtime.edit();
+    await runtime.idle();
+
+    expect(comp1Runs).toBe(2);
+    expect(comp2Runs).toBe(2);
+    expect(effectRuns).toBe(2);
+    expect(effectResult.get()).toBe((5 + 1) * 2 - 3);
+  });
+
+  it("should drop stale dependents when computation changes inputs", async () => {
+    runtime.scheduler.enablePullMode();
+
+    const sourceA = runtime.getCell<number>(
+      space,
+      "pull-deps-source-a",
+      undefined,
+      tx,
+    );
+    sourceA.set(2);
+    const sourceB = runtime.getCell<number>(
+      space,
+      "pull-deps-source-b",
+      undefined,
+      tx,
+    );
+    sourceB.set(7);
+    const selector = runtime.getCell<boolean>(
+      space,
+      "pull-deps-selector",
+      undefined,
+      tx,
+    );
+    selector.set(false);
+    const intermediate = runtime.getCell<number>(
+      space,
+      "pull-deps-intermediate",
+      undefined,
+      tx,
+    );
+    intermediate.set(0);
+    const effectResult = runtime.getCell<number>(
+      space,
+      "pull-deps-effect",
+      undefined,
+      tx,
+    );
+    effectResult.set(0);
+    await tx.commit();
+    tx = runtime.edit();
+
+    let effectRuns = 0;
+
+    const computation: Action = (actionTx) => {
+      const useB = selector.withTx(actionTx).get();
+      const value = useB
+        ? sourceB.withTx(actionTx).get()
+        : sourceA.withTx(actionTx).get();
+      intermediate.withTx(actionTx).send(value * 10);
+    };
+
+    const effect: Action = (actionTx) => {
+      effectRuns++;
+      const value = intermediate.withTx(actionTx).get();
+      effectResult.withTx(actionTx).send(value);
+    };
+
+    runtime.scheduler.subscribe(
+      computation,
+      {
+        reads: [
+          selector.getAsNormalizedFullLink(),
+          sourceA.getAsNormalizedFullLink(),
+        ],
+        writes: [intermediate.getAsNormalizedFullLink()],
+      },
+      { scheduleImmediately: true },
+    );
+    await runtime.idle();
+
+    runtime.scheduler.subscribe(
+      effect,
+      {
+        reads: [intermediate.getAsNormalizedFullLink()],
+        writes: [effectResult.getAsNormalizedFullLink()],
+      },
+      { scheduleImmediately: true, isEffect: true },
+    );
+    await runtime.idle();
+
+    expect(effectRuns).toBe(1);
+    expect(effectResult.get()).toBe(20);
+
+    // Switch computation to sourceB
+    const toggleTx = runtime.edit();
+    selector.withTx(toggleTx).send(true);
+    await toggleTx.commit();
+    tx = runtime.edit();
+    await runtime.idle();
+
+    expect(effectRuns).toBe(2);
+    expect(effectResult.get()).toBe(70);
+
+    // Updating sourceA should not dirty the computation any more
+    const tx3 = runtime.edit();
+    sourceA.withTx(tx3).send(999);
+    await tx3.commit();
+    tx = runtime.edit();
+    await runtime.idle();
+
+    expect(effectRuns).toBe(2);
+    expect(effectResult.get()).toBe(70);
+    expect(runtime.scheduler.isDirty(computation)).toBe(false);
+
+    // Updating sourceB should still run the computation
+    const tx4 = runtime.edit();
+    sourceB.withTx(tx4).send(6);
+    await tx4.commit();
+    tx = runtime.edit();
+    await runtime.idle();
+
+    expect(effectRuns).toBe(3);
+    expect(effectResult.get()).toBe(60);
+  });
+
   it("should track getStats with dirty count", async () => {
     runtime.scheduler.enablePullMode();