[FIX] reactivity: memory leak

The deletion of a key in an observed object did clear the observers of
that key but did not clear the atoms created (if any) when the key existed
(e.g. on the key value if it was trackable). This can lead to a growing
set of atoms that are useless but kept in memory if a lot of keys are
added/deleted. The same thing can happen if a key value is changed many
times and the values are trackable.

Here we fix the problem by
- keeping tracks of the objects that have been observed by an external call
  to the method "atom" (we call them seeds).
- remove all observer atoms that are not seeds for observers of at least
  one key deletion or key value change.

Note the fix is in some sense partial: a user could use the "atom" primitive
making the new system uncapable of avoiding a leak (see test "atom on an
object with a getter 3" where a getter is used in a weird way in an
observed object).

src/reactivity.ts

@@ -12,18 +12,29 @@ const SOURCE = Symbol("source");
 const OBSERVER = Symbol("observer");
 const KEYS = Symbol("keys");
 const ROOT = Symbol("root");
+const SEED = Symbol("Seed");
 
-export function atom<T>(source: T, observer: Observer): T {
+export function atom(source: any, observer: Observer) {
+  return _atom(source, observer, true);
+}
+
+export function _atom(source: any, observer: Observer, seed = false) {
   if (isTrackable(source) && observerSourceAtom.has(observer)) {
     source = (source as any)[SOURCE] || source;
     const oldAtom = observerSourceAtom.get(observer)!.get(source);
     if (oldAtom) {
+      if (seed) {
+        oldAtom[SEED] = true;
+      }
       return oldAtom;
     }
     if (!sourceAtoms.get(source)) {
       sourceAtoms.set(source, new Map([[ROOT, new ObserverSet()]]));
     }
     const newAtom = createAtom(source, observer);
+    if (seed) {
+      newAtom[SEED] = true;
+    }
     observerSourceAtom.get(observer)!.set(source, newAtom);
     sourceAtoms.get(source)!.get(ROOT)!.add(newAtom);
     return newAtom;
@@ -33,8 +44,13 @@ export function atom<T>(source: T, observer: Observer): T {
 
 function createAtom(source: Source, observer: Observer): Atom {
   const keys: Set<any> = new Set();
-  return new Proxy(source as any, {
-    set(target: any, key: string, value: any): boolean {
+  let seed: boolean = false;
+  const newAtom: Atom = new Proxy(source as any, {
+    set(target: any, key: any, value: any): boolean {
+      if (key === SEED) {
+        seed = value;
+        return true;
+      }
       if (!(key in target)) {
         target[key] = value;
         notify(sourceAtoms.get(source)!.get(ROOT)!);
@@ -45,20 +61,25 @@ function createAtom(source: Source, observer: Observer): Atom {
         target[key] = value;
         const observerSet = sourceAtoms.get(source)!.get(key);
         if (observerSet) {
-          notify(observerSet);
+          const clean = isTrackable(current);
+          notify(observerSet, clean);
         }
       }
       return true;
     },
     deleteProperty(target: any, key: string): boolean {
       if (key in target) {
+        const current = target[key];
         delete target[key];
         // notify source observers
-        notify(sourceAtoms.get(source)!.get(ROOT)!);
+        const clean = isTrackable(current);
+        notify(sourceAtoms.get(source)!.get(ROOT)!, clean);
         const atoms = sourceAtoms.get(source)!;
         if (atoms.has(key)) {
           // clear source-key observers
-          atoms.get(key)!.deleteKey(key);
+          for (const atom of atoms.get(key)!) {
+            atom[KEYS].delete(key);
+          }
           atoms.delete(key);
         }
       }
@@ -72,6 +93,8 @@ function createAtom(source: Source, observer: Observer): Atom {
           return source;
         case KEYS:
           return keys;
+        case SEED:
+          return seed;
         default:
           const value = target[key];
           // register observer to source-key
@@ -80,14 +103,15 @@ function createAtom(source: Source, observer: Observer): Atom {
             if (!atoms.has(key)) {
               atoms.set(key, new ObserverSet());
             }
-            atoms.get(key)!.add(proxy);
+            atoms.get(key)!.add(newAtom);
             keys.add(key);
           }
           //
-          return atom(value, observer);
+          return _atom(value, observer);
       }
     },
   });
+  return newAtom;
 }
 
 function isTrackable(value: any): boolean {
@@ -103,18 +127,25 @@ export function registerObserver(observer: Observer) {
   if (!observerSourceAtom.get(observer)) {
     observerSourceAtom.set(observer, new Map());
   }
-  return unregisterObserver.bind(null, observer);
+  return unregisterObserverAtoms.bind(null, observer, false);
 }
 
-function unregisterObserver(observer: Observer) {
-  for (const [source, atom] of observerSourceAtom.get(observer)!) {
+function unregisterObserverAtoms(observer: Observer, keepSeeds: boolean) {
+  const observerAtoms = observerSourceAtom.get(observer)!;
+  for (const [source, atom] of observerAtoms) {
+    if (keepSeeds && atom[SEED]) {
+      continue;
+    }
+    observerAtoms.delete(source);
     const atoms = sourceAtoms.get(source)!;
     atoms.get(ROOT)!.delete(atom);
     for (const key of atom[KEYS]) {
       atoms.get(key)!.delete(atom);
     }
   }
-  observerSourceAtom.delete(observer);
+  if (!keepSeeds) {
+    observerSourceAtom.delete(observer);
+  }
 }
 
 export function useState(state: any): Atom {
@@ -145,14 +176,6 @@ class ObserverSet {
       return this.callbackSet.delete(atom);
     }
   }
-  deleteKey(key: any) {
-    for (const atom of this.nodeAntichain) {
-      atom[KEYS].delete(key);
-    }
-    for (const atom of this.callbackSet) {
-      atom[KEYS].delete(key);
-    }
-  }
   union(other: ObserverSet) {
     for (const atom of other.nodeAntichain) {
       this.nodeAntichain.add(atom);
@@ -163,12 +186,22 @@ class ObserverSet {
   }
   notify() {
     for (const atom of this.nodeAntichain) {
+      // an observer cannot be found twice here: <= in add is based on observers
       atom[OBSERVER].render();
     }
+    const called = new Set();
     for (const atom of this.callbackSet) {
-      atom[OBSERVER]();
+      const observer = atom[OBSERVER];
+      if (!called.has(observer)) {
+        observer();
+      }
+      called.add(observer);
     }
   }
+  *[Symbol.iterator]() {
+    yield* this.nodeAntichain;
+    yield* this.callbackSet;
+  }
 }
 
 // Need to optimize this!
@@ -213,7 +246,13 @@ class Antichain extends Set<Atom> {
 }
 
 let toNotify: ObserverSet | null = null;
-async function notify(observers: ObserverSet) {
+let toClean: Set<Observer> = new Set();
+async function notify(observers: ObserverSet, clean = false) {
+  if (clean) {
+    for (const atom of observers) {
+      toClean.add(atom[OBSERVER]);
+    }
+  }
   if (toNotify) {
     toNotify.union(observers);
     return;
@@ -221,6 +260,10 @@ async function notify(observers: ObserverSet) {
   toNotify = new ObserverSet();
   toNotify.union(observers);
   await Promise.resolve();
+  for (const observer of toClean) {
+    unregisterObserverAtoms(observer, true);
+  }
+  toClean.clear();
   toNotify.notify();
   toNotify = null;
 }

tests/reactivity.test.ts

@@ -76,6 +76,67 @@ describe("Reactivity: atom", () => {
     expect(n).toBe(1);
   });
 
+  test("atom on an object with a getter 1", async () => {
+    let n = 0;
+    let value = 1;
+    const atom = createAtom(
+      {
+        get a() {
+          return value;
+        },
+        set a(val) {
+          value = val;
+        },
+      },
+      () => n++
+    );
+    atom.a = atom.a + 4;
+    await nextMicroTick();
+    expect(n).toBe(1);
+  });
+
+  test("atom on an object with a getter 2", async () => {
+    let n = 0;
+    let value = { b: 1 };
+    const atom = createAtom(
+      {
+        get a() {
+          return value;
+        },
+      },
+      () => n++
+    );
+    expect(atom.a.b).toBe(1);
+    atom.a.b = 2;
+    await nextMicroTick();
+    expect(n).toBe(1);
+  });
+
+  test("atom on an object with a getter 3", async () => {
+    let n = 0;
+    const values: { b: number }[] = createAtom([]);
+    function createValue() {
+      const o = { b: values.length };
+      values.push(o);
+      return o;
+    }
+    const atom = createAtom(
+      {
+        get a() {
+          return createValue();
+        },
+      },
+      () => n++
+    );
+    for (let i = 0; i < 10; i++) {
+      expect(atom.a.b).toEqual(i);
+    }
+    expect(n).toBe(0);
+    values[0].b = 3;
+    await nextMicroTick();
+    expect(n).toBe(1); // !!! atoms for each object in values are still there !!!
+  });
+
   test("atom observer is called after batch of operation", async () => {
     let n = 0;
     const atom1 = createAtom({ a: 1, b: 2 }, () => n++);
@@ -553,6 +614,7 @@ describe("Reactivity: atom", () => {
   test("properly handle already observed atom", async () => {
     let n1 = 0;
     let n2 = 0;
+
     const obj1 = createAtom({ a: 1 }, () => n1++) as any;
     const obj2 = createAtom({ b: 1 }, () => n2++) as any;
 
@@ -887,6 +949,106 @@ describe("Reactivity: atom", () => {
     await nextMicroTick();
     expect(n).toBe(0);
   });
+
+  test("atom cleaning after key deleted 1", async () => {
+    let n = 0;
+    const a = { k: {} } as any;
+    const atom = createAtom(a, () => n++);
+
+    atom.k.l = 1;
+    await nextMicroTick();
+    expect(n).toBe(1);
+
+    const kVal = atom.k;
+
+    delete atom.k;
+    await nextMicroTick();
+    expect(n).toBe(2);
+
+    kVal.l = 2;
+    await nextMicroTick();
+    expect(n).toBe(2); // kVal must no longer be observed
+  });
+
+  test("atom cleaning after key deleted 2", async () => {
+    let n = 0;
+    const b = {} as any;
+    const a = { k: b } as any;
+    const observer = () => n++;
+    const atom2 = createAtom(b, observer);
+    const atom1 = createAtom(a, observer); // atom2 should be still a seed
+
+    atom1.c = 1;
+    atom1.k.d = 2;
+    await nextMicroTick();
+    expect(n).toBe(1);
+
+    delete atom1.k;
+    await nextMicroTick();
+    expect(n).toBe(2);
+
+    atom2.e = 3;
+    await nextMicroTick();
+    expect(n).toBe(3); // b must still be observed
+  });
+
+  test("atom cleaning after key deleted 3", async () => {
+    let n = 0;
+    const b = {} as any;
+    const a = { k: b } as any;
+    const observer = () => n++;
+    const atom1 = createAtom(a, observer);
+    const atom2 = createAtom(b, observer); // atom2 should be a seed now
+
+    atom1.c = 1;
+    atom1.k.d = 2;
+    await nextMicroTick();
+    expect(n).toBe(1);
+
+    delete atom1.k;
+    await nextMicroTick();
+    expect(n).toBe(2);
+
+    atom2.e = 3;
+    await nextMicroTick();
+    expect(n).toBe(3); // b must still be observed
+  });
+
+  test("atom cleaning after key deleted 4", async () => {
+    let n = 0;
+    const a = { k: {} } as any;
+    a.k = a;
+    const atom = createAtom(a, () => n++);
+
+    atom.b = 1;
+    await nextMicroTick();
+    expect(n).toBe(1);
+
+    delete atom.k;
+    await nextMicroTick();
+    expect(n).toBe(2);
+
+    atom.c = 2;
+    await nextMicroTick();
+    expect(n).toBe(3); // a must still be observed
+  });
+
+  test("atom cleaning after trackable key value changed", async () => {
+    let n = 0;
+    const a = { k: { n: 1 } } as any;
+    const atom = createAtom(a, () => n++);
+    const kVal = atom.k; // read k
+
+    atom.k = { n: atom.k.n + 1 };
+    await nextMicroTick();
+    expect(n).toBe(1);
+    expect(atom.k).toEqual({ n: 2 });
+
+    kVal.n = 3;
+    await nextMicroTick();
+    expect(n).toBe(1);
+    expect(atom.k).toEqual({ n: 2 });
+  });
 });
 
 describe("Reactivity: useState", () => {
@@ -1324,14 +1486,11 @@ describe("Reactivity: useState", () => {
     expect([...steps]).toEqual(["list", "quantity1"]);
     steps.clear();
 
-    // secondQuantity is no longer accessible by List! But list react to changes
-    // --> this prooff that a useless atom continues to exist
     secondQuantity.quantity = 2;
     await nextMicroTick();
     await nextMicroTick();
-
     expect(fixture.innerHTML).toBe("<div><div>3</div> Total: 3 Count: 1</div>");
-    expect([...steps]).toEqual(["list"]); // should be avoided!!!
+    expect([...steps]).toEqual([]);
     steps.clear();
   });