From ebf67ad2087c3efbfeb64fba62350aad6c96189d Mon Sep 17 00:00:00 2001
From: Evan You <yyx990803@gmail.com>
Date: Sun, 28 Oct 2018 12:08:58 -0400
Subject: [PATCH] fix(scheduler): handle queueJob inside postFlushCbs

---
 .../scheduler/__tests__/scheduler.spec.ts     | 23 +++++++
 packages/scheduler/src/index.ts               | 61 ++++++++++++++++---
 2 files changed, 74 insertions(+), 10 deletions(-)

diff --git a/packages/scheduler/__tests__/scheduler.spec.ts b/packages/scheduler/__tests__/scheduler.spec.ts
index a0aa89fd..74c21dae 100644
--- a/packages/scheduler/__tests__/scheduler.spec.ts
+++ b/packages/scheduler/__tests__/scheduler.spec.ts
@@ -90,4 +90,27 @@ describe('scheduler', () => {
     await nextTick()
     expect(calls).toEqual(['job1', 'job2'])
   })
+
+  it('queueJob inside postFlushCb', async () => {
+    const calls: any = []
+    const job1 = () => {
+      calls.push('job1')
+    }
+    const cb1 = () => {
+      // queue another job in postFlushCb
+      calls.push('cb1')
+      queueJob(job2, cb2)
+    }
+    const job2 = () => {
+      calls.push('job2')
+    }
+    const cb2 = () => {
+      calls.push('cb2')
+    }
+
+    queueJob(job1, cb1)
+    queueJob(job2, cb2)
+    await nextTick()
+    expect(calls).toEqual(['job1', 'job2', 'cb1', 'cb2', 'job2', 'cb2'])
+  })
 })
diff --git a/packages/scheduler/src/index.ts b/packages/scheduler/src/index.ts
index 93172ffd..54c8567e 100644
--- a/packages/scheduler/src/index.ts
+++ b/packages/scheduler/src/index.ts
@@ -1,33 +1,74 @@
 const queue: Array<() => void> = []
 const postFlushCbs: Array<() => void> = []
+const postFlushCbsForNextTick: Array<() => void> = []
 const p = Promise.resolve()
 
-let hasPendingFlush = false
+let isFlushing = false
+let isFlushingPostCbs = false
 
 export function nextTick(fn?: () => void): Promise<void> {
   return p.then(fn)
 }
 
-export function queueJob(job: () => void, postFlushCb?: () => void) {
+export function queueJob(
+  job: () => void,
+  postFlushCb?: () => void,
+  onError?: (err: Error) => void
+) {
   if (queue.indexOf(job) === -1) {
     queue.push(job)
-    if (!hasPendingFlush) {
-      hasPendingFlush = true
-      nextTick(flushJobs)
+    if (!isFlushing || isFlushingPostCbs) {
+      const p = nextTick(flushJobs)
+      if (onError) p.catch(onError)
     }
   }
-  if (postFlushCb && postFlushCbs.indexOf(postFlushCb) === -1) {
-    postFlushCbs.push(postFlushCb)
+  if (postFlushCb) {
+    if (isFlushingPostCbs) {
+      // it's possible for a postFlushCb to queue another job/cb combo,
+      // e.g. triggering a state update inside the updated hook.
+      if (postFlushCbsForNextTick.indexOf(postFlushCb) === -1) {
+        postFlushCbsForNextTick.push(postFlushCb)
+      }
+    } else if (postFlushCbs.indexOf(postFlushCb) === -1) {
+      postFlushCbs.push(postFlushCb)
+    }
   }
 }
 
+const seenJobs = new Map()
+const RECURSION_LIMIT = 100
+
 function flushJobs() {
+  seenJobs.clear()
+  isFlushing = true
   let job
   while ((job = queue.shift())) {
+    if (__DEV__) {
+      if (!seenJobs.has(job)) {
+        seenJobs.set(job, 1)
+      } else {
+        const count = seenJobs.get(job)
+        if (count > RECURSION_LIMIT) {
+          throw new Error('Maximum recursive updates exceeded')
+        } else {
+          seenJobs.set(job, count + 1)
+        }
+      }
+    }
     job()
   }
-  while ((job = postFlushCbs.shift())) {
-    job()
+  isFlushingPostCbs = true
+  if (postFlushCbsForNextTick.length > 0) {
+    const postFlushCbsFromPrevTick = postFlushCbsForNextTick.slice()
+    postFlushCbsForNextTick.length = 0
+    for (let i = 0; i < postFlushCbsFromPrevTick.length; i++) {
+      postFlushCbsFromPrevTick[i]()
+    }
   }
-  hasPendingFlush = false
+  for (let i = 0; i < postFlushCbs.length; i++) {
+    postFlushCbs[i]()
+  }
+  postFlushCbs.length = 0
+  isFlushingPostCbs = false
+  isFlushing = false
 }