From d80fad760cf8094d270b181a0d78b8666f8d4a0b Mon Sep 17 00:00:00 2001
From: Samuel Degueldre <sad@odoo.com>
Date: Tue, 29 Mar 2022 08:17:47 +0200
Subject: [PATCH] [FIX] utils: fix calls to batched callback from within the
 callback

Previously, calling the batched function from within the callback being
batched would fail as it would be treated as part of the same batch.
This commit fixes that by scheduling the reset of the "called" flag
before calling the callback. This means that all microtasks that were
already in the microtask queue when a batch is about to run are treated
as part of the batch, and all microtasks that will be added by the
callback are not.
---
 src/utils.ts             | 10 ++++++----
 tests/reactivity.test.ts | 40 ++++++++++++++++++++++++++++++++++++++++
 tests/utils.test.ts      | 20 ++++++++++++++++++++
 3 files changed, 66 insertions(+), 4 deletions(-)

diff --git a/src/utils.ts b/src/utils.ts
index 4b3cda783..78d328763 100644
--- a/src/utils.ts
+++ b/src/utils.ts
@@ -15,11 +15,13 @@ export function batched(callback: Callback): Callback {
     await Promise.resolve();
     if (!called) {
       called = true;
-      callback();
       // wait for all calls in this microtick to fall through before resetting "called"
-      // so that only the first call to the batched function calls the original callback
-      await Promise.resolve();
-      called = false;
+      // so that only the first call to the batched function calls the original callback.
+      // Schedule this before calling the callback so that calls to the batched function
+      // within the callback will proceed only after resetting called to false, and have
+      // a chance to execute the callback again
+      Promise.resolve().then(() => (called = false));
+      callback();
     }
   };
 }
diff --git a/tests/reactivity.test.ts b/tests/reactivity.test.ts
index 8a832024b..8d36d9258 100644
--- a/tests/reactivity.test.ts
+++ b/tests/reactivity.test.ts
@@ -208,6 +208,46 @@ describe("Reactivity", () => {
     expect(n).toBe(1); // two operations but only one notification
   });
 
+  test("batched: modifying the reactive in the callback doesn't break reactivity", async () => {
+    let n = 0;
+    let obj = { a: 1 };
+    const state = createReactive(
+      obj,
+      batched(() => {
+        state.a; // subscribe to a
+        state.a = 2;
+        n++;
+      })
+    );
+    expect(n).toBe(0);
+    state.a = 2;
+    expect(n).toBe(0);
+    await nextMicroTick();
+    expect(n).toBe(0); // key has not be read yet
+    state.a = state.a + 5; // key is read and then modified
+    expect(n).toBe(0);
+    await nextMicroTick();
+    expect(n).toBe(1);
+    // the write a = 2 inside the batched callback triggered another notification, wait for it
+    await nextMicroTick();
+    expect(n).toBe(2);
+    // Should now be stable as we're writing the same value again
+    await nextMicroTick();
+    expect(n).toBe(2);
+
+    // Do it again to check it's not broken
+    state.a = state.a + 5; // key is read and then modified
+    expect(n).toBe(2);
+    await nextMicroTick();
+    expect(n).toBe(3);
+    // the write a = 2 inside the batched callback triggered another notification, wait for it
+    await nextMicroTick();
+    expect(n).toBe(4);
+    // Should now be stable as we're writing the same value again
+    await nextMicroTick();
+    expect(n).toBe(4);
+  });
+
   test("setting property to same value does not trigger callback", async () => {
     let n = 0;
     const state = createReactive({ a: 1 }, () => n++);
diff --git a/tests/utils.test.ts b/tests/utils.test.ts
index 19081358c..e2a795006 100644
--- a/tests/utils.test.ts
+++ b/tests/utils.test.ts
@@ -50,4 +50,24 @@ describe("batched", () => {
     await nextMicroTick();
     expect(n).toBe(1);
   });
+
+  test("calling batched function from within the callback is not treated as part of the original batch", async () => {
+    let n = 0;
+    let fn = batched(() => {
+      n++;
+      if (n === 1) {
+        fn();
+      }
+    });
+
+    expect(n).toBe(0);
+    fn();
+    expect(n).toBe(0);
+    await nextMicroTick(); // First batch
+    expect(n).toBe(1);
+    await nextMicroTick(); // Second batch initiated from within the callback
+    expect(n).toBe(2);
+    await nextMicroTick();
+    expect(n).toBe(2);
+  });
 });