fix(core): treat exceptions in equal as part of computation (#55818)

leonsenft · AndrewKushnir · commit d12a186d531b · 2025-01-16T08:10:10.000-08:00
Prevent leaking signal reads and exceptions from a custom `equal` function of a producer `computed()` to a consumer. Upstream tc39/proposal-signals#90 with a notable change: Angular does **not** track reactive reads in custom `equal` implementations. PR Close #55818
diff --git a/packages/core/primitives/signals/src/computed.ts b/packages/core/primitives/signals/src/computed.ts
@@ -14,6 +14,7 @@ import {
   producerUpdateValueVersion,
   REACTIVE_NODE,
   ReactiveNode,
+  setActiveConsumer,
   SIGNAL,
 } from './graph';
 
@@ -120,21 +121,25 @@ const COMPUTED_NODE = /* @__PURE__ */ (() => {
 
       const prevConsumer = consumerBeforeComputation(node);
       let newValue: unknown;
+      let wasEqual = false;
       try {
         newValue = node.computation();
+        // We want to mark this node as errored if calling `equal` throws; however, we don't want
+        // to track any reactive reads inside `equal`.
+        setActiveConsumer(null);
+        wasEqual =
+          oldValue !== UNSET &&
+          oldValue !== ERRORED &&
+          newValue !== ERRORED &&
+          node.equal(oldValue, newValue);
       } catch (err) {
         newValue = ERRORED;
         node.error = err;
       } finally {
         consumerAfterComputation(node, prevConsumer);
       }
 
-      if (
-        oldValue !== UNSET &&
-        oldValue !== ERRORED &&
-        newValue !== ERRORED &&
-        node.equal(oldValue, newValue)
-      ) {
+      if (wasEqual) {
         // No change to `valueVersion` - old and new values are
         // semantically equivalent.
         node.value = oldValue;
diff --git a/packages/core/test/signals/computed_spec.ts b/packages/core/test/signals/computed_spec.ts
@@ -7,7 +7,7 @@
  */
 
 import {computed, signal} from '@angular/core';
-import {createWatch, ReactiveNode, SIGNAL} from '@angular/core/primitives/signals';
+import {createWatch, ReactiveNode, SIGNAL, defaultEquals} from '@angular/core/primitives/signals';
 
 describe('computed', () => {
   it('should create computed', () => {
@@ -201,4 +201,120 @@ describe('computed', () => {
     ] as ReactiveNode;
     expect(node.debugName).toBe('computedSignal');
   });
+
+  describe('with custom equal', () => {
+    it('should cache exceptions thrown by equal', () => {
+      const s = signal(0);
+
+      let computedRunCount = 0;
+      let equalRunCount = 0;
+      const c = computed(
+        () => {
+          computedRunCount++;
+          return s();
+        },
+        {
+          equal: () => {
+            equalRunCount++;
+            throw new Error('equal');
+          },
+        },
+      );
+
+      // equal() isn't run for the initial computation.
+      expect(c()).toBe(0);
+      expect(computedRunCount).toBe(1);
+      expect(equalRunCount).toBe(0);
+
+      s.set(1);
+
+      // Error is thrown by equal().
+      expect(() => c()).toThrowError('equal');
+      expect(computedRunCount).toBe(2);
+      expect(equalRunCount).toBe(1);
+
+      // Error is cached; c throws again without needing to rerun computation or equal().
+      expect(() => c()).toThrowError('equal');
+      expect(computedRunCount).toBe(2);
+      expect(equalRunCount).toBe(1);
+    });
+
+    it('should not track signal reads inside equal', () => {
+      const value = signal(1);
+      const epsilon = signal(0.5);
+
+      let innerRunCount = 0;
+      let equalRunCount = 0;
+      const inner = computed(
+        () => {
+          innerRunCount++;
+          return value();
+        },
+        {
+          equal: (a, b) => {
+            equalRunCount++;
+            return Math.abs(a - b) < epsilon();
+          },
+        },
+      );
+
+      let outerRunCount = 0;
+      const outer = computed(() => {
+        outerRunCount++;
+        return inner();
+      });
+
+      // Everything runs the first time.
+      expect(outer()).toBe(1);
+      expect(innerRunCount).toBe(1);
+      expect(outerRunCount).toBe(1);
+
+      // Difference is less than epsilon().
+      value.set(1.2);
+
+      // `inner` reruns because `value` was changed, and `equal` is called for the first time.
+      expect(outer()).toBe(1);
+      expect(innerRunCount).toBe(2);
+      expect(equalRunCount).toBe(1);
+      // `outer does not rerun because `equal` determined that `inner` had not changed.
+      expect(outerRunCount).toBe(1);
+
+      // Previous difference is now greater than epsilon().
+      epsilon.set(0.1);
+
+      // While changing `epsilon` would change the outcome of the `inner`, we don't rerun it
+      // because we intentionally don't track reactive reads in `equal`.
+      expect(outer()).toBe(1);
+      expect(innerRunCount).toBe(2);
+      expect(equalRunCount).toBe(1);
+      // Equally important is that the signal read in `equal` doesn't leak into the outer reactive
+      // context either.
+      expect(outerRunCount).toBe(1);
+    });
+
+    it('should recover from exception', () => {
+      let shouldThrow = true;
+      const source = signal(0);
+      const derived = computed(source, {
+        equal: (a, b) => {
+          if (shouldThrow) {
+            throw new Error('equal');
+          }
+          return defaultEquals(a, b);
+        },
+      });
+
+      // Initial read doesn't throw because it doesn't call `equal`.
+      expect(derived()).toBe(0);
+
+      // Update `source` to begin throwing.
+      source.set(1);
+      expect(() => derived()).toThrowError('equal');
+
+      // Stop throwing and update `source` to cause `derived` to recompute.
+      shouldThrow = false;
+      source.set(2);
+      expect(derived()).toBe(2);
+    });
+  });
 });
