From 4a037040e5d9d806a9d7a9c8fd3fc3b243e4617f Mon Sep 17 00:00:00 2001
From: Michael Peyper <mpeyper7@gmail.com>
Date: Tue, 31 Aug 2021 21:25:26 +1000
Subject: [PATCH] fix(async-utils): prevent timeout and interval checks in wait
 from leaving open handles (#682)

* fix(async-utils): prevent timeout and interval checks in wait from leaving open handles

* refactor(async-utils): rename timeoutSignal to timeoutController
---
 src/core/asyncUtils.ts                 | 33 ++++++++--------------
 src/helpers/createTimeoutController.ts | 39 ++++++++++++++++++++++++++
 src/helpers/promises.ts                | 10 -------
 3 files changed, 50 insertions(+), 32 deletions(-)
 create mode 100644 src/helpers/createTimeoutController.ts
 delete mode 100644 src/helpers/promises.ts

diff --git a/src/core/asyncUtils.ts b/src/core/asyncUtils.ts
index fe44c715..a7424036 100644
--- a/src/core/asyncUtils.ts
+++ b/src/core/asyncUtils.ts
@@ -7,7 +7,7 @@ import {
   AsyncUtils
 } from '../types'
 
-import { resolveAfter, callAfter } from '../helpers/promises'
+import { createTimeoutController } from '../helpers/createTimeoutController'
 import { TimeoutError } from '../helpers/error'
 
 const DEFAULT_INTERVAL = 50
@@ -20,37 +20,26 @@ function asyncUtils(act: Act, addResolver: (callback: () => void) => void): Asyn
       return callbackResult ?? callbackResult === undefined
     }
 
+    const timeoutSignal = createTimeoutController(timeout)
+
     const waitForResult = async () => {
       while (true) {
-        await Promise.race(
-          [
-            new Promise<void>((resolve) => addResolver(resolve)),
-            interval && resolveAfter(interval)
-          ].filter(Boolean)
-        )
-
-        if (checkResult()) {
+        const intervalSignal = createTimeoutController(interval)
+        timeoutSignal.onTimeout(() => intervalSignal.cancel())
+
+        await intervalSignal.wrap(new Promise<void>(addResolver))
+
+        if (checkResult() || timeoutSignal.timedOut) {
           return
         }
       }
     }
 
-    let timedOut = false
-
     if (!checkResult()) {
-      if (timeout) {
-        const timeoutPromise = () =>
-          callAfter(() => {
-            timedOut = true
-          }, timeout)
-
-        await act(() => Promise.race([waitForResult(), timeoutPromise()]))
-      } else {
-        await act(waitForResult)
-      }
+      await act(() => timeoutSignal.wrap(waitForResult()))
     }
 
-    return !timedOut
+    return !timeoutSignal.timedOut
   }
 
   const waitFor = async (
diff --git a/src/helpers/createTimeoutController.ts b/src/helpers/createTimeoutController.ts
new file mode 100644
index 00000000..643d3768
--- /dev/null
+++ b/src/helpers/createTimeoutController.ts
@@ -0,0 +1,39 @@
+import { WaitOptions } from '../types'
+
+function createTimeoutController(timeout: WaitOptions['timeout']) {
+  let timeoutId: NodeJS.Timeout
+  const timeoutCallbacks: Array<() => void> = []
+
+  const timeoutController = {
+    onTimeout(callback: () => void) {
+      timeoutCallbacks.push(callback)
+    },
+    wrap(promise: Promise<void>) {
+      return new Promise<void>((resolve, reject) => {
+        timeoutController.timedOut = false
+        timeoutController.onTimeout(resolve)
+
+        if (timeout) {
+          timeoutId = setTimeout(() => {
+            timeoutController.timedOut = true
+            timeoutCallbacks.forEach((callback) => callback())
+            resolve()
+          }, timeout)
+        }
+
+        promise
+          .then(resolve)
+          .catch(reject)
+          .finally(() => timeoutController.cancel())
+      })
+    },
+    cancel() {
+      clearTimeout(timeoutId)
+    },
+    timedOut: false
+  }
+
+  return timeoutController
+}
+
+export { createTimeoutController }
diff --git a/src/helpers/promises.ts b/src/helpers/promises.ts
deleted file mode 100644
index 2fa89e5f..00000000
--- a/src/helpers/promises.ts
+++ /dev/null
@@ -1,10 +0,0 @@
-function resolveAfter(ms: number) {
-  return new Promise<void>((resolve) => setTimeout(resolve, ms))
-}
-
-async function callAfter(callback: () => void, ms: number) {
-  await resolveAfter(ms)
-  callback()
-}
-
-export { resolveAfter, callAfter }