From cef64024e6ddb2b53f9dce952fb0dac606e72306 Mon Sep 17 00:00:00 2001
From: Eryn Lynn <eryn@eryn.io>
Date: Sun, 29 Sep 2019 00:07:32 -0400
Subject: [PATCH] Add Promise.delay, :timeout

Closes #8
Closes #7
---
 CHANGELOG.md  |  1 +
 lib/README.md | 29 +++++++++++++++++++-
 lib/init.lua  | 73 +++++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 102 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 4e0ab4a..b82569d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -10,6 +10,7 @@
 - Add Promise.try
 - Add `done`, `doneCall`, `doneReturn`
 - Add `andThenReturn`, `finallyReturn`
+- Add `Promise.delay`, `promise:timeout`
 
 # 2.4.0
 
diff --git a/lib/README.md b/lib/README.md
index 6538432..6996845 100644
--- a/lib/README.md
+++ b/lib/README.md
@@ -220,6 +220,18 @@ docs:
       static: true
       params: "promises: array<Promise<T>>"
       returns: Promise<T>
+    - name: delay
+      desc: |
+        Returns a Promise that resolves after `seconds` seconds have passed. The Promise resolves with the actual amount of time that was waited.
+
+        This function is **not** a wrapper around `wait`. `Promise.delay` uses a custom scheduler which provides more accurate timing. As an optimization, cancelling this Promise instantly removes the task from the scheduler.
+
+        ::: warning
+          Passing `NaN`, infinity, or a number less than 1/60 is equivalent to passing 1/60.
+        :::
+      params: "seconds: number"
+      returns: Promise<number>
+      static: true
     - name: is
       desc: Returns whether the given object is a Promise. This only checks if the object is a table and has an `andThen` method.
       static: true
@@ -473,12 +485,27 @@ docs:
           desc: Values to return from the function.
       returns: Promise<...any?>
 
+    - name: timeout
+      params: "seconds: number, rejectionValue: any?"
+      desc: |
+        Returns a new Promise that resolves if the chained Promise resolves within `seconds` seconds, or rejects if execution time exceeds `seconds`. The chained Promise will be cancelled if the timeout is reached.
+
+        Sugar for:
+
+        ```lua
+        Promise.race({
+          Promise.delay(seconds):andThen(function()
+            return Promise.reject(rejectionValue == nil and "Timed out" or rejectionValue)
+          end),
+          promise
+        })
+        ```
 
     - name: cancel
       desc: |
         Cancels this promise, preventing the promise from resolving or rejecting. Does not do anything if the promise is already settled.
 
-        Cancellations will propagate upwards through chained promises.
+        Cancellations will propagate upwards and downwards through chained promises.
 
         Promises will only be cancelled if all of their consumers are also cancelled. This is to say that if you call `andThen` twice on the same promise, and you cancel only one of the child promises, it will not cancel the parent promise until the other child promise is also cancelled.
 
diff --git a/lib/init.lua b/lib/init.lua
index 0be1852..eabe803 100644
--- a/lib/init.lua
+++ b/lib/init.lua
@@ -392,6 +392,79 @@ function Promise.promisify(callback)
 	end
 end
 
+--[[
+	Creates a Promise that resolves after given number of seconds.
+]]
+do
+	local connection
+	local queue = {}
+
+	local function enqueue(callback, seconds)
+		table.insert(queue, {
+			callback = callback,
+			startTime = tick(),
+			endTime = tick() + math.max(seconds, 1/60)
+		})
+
+		table.sort(queue, function(a, b)
+			return a.endTime < b.endTime
+		end)
+
+		if not connection then
+			connection = RunService.Heartbeat:Connect(function()
+				while #queue > 0 and queue[1].endTime <= tick() do
+					local item = table.remove(queue, 1)
+
+					item.callback(tick() - item.startTime)
+				end
+
+				if #queue == 0 then
+					connection:Disconnect()
+					connection = nil
+				end
+			end)
+		end
+	end
+
+	local function dequeue(callback)
+		for i, item in ipairs(queue) do
+			if item.callback == callback then
+				table.remove(queue, i)
+				break
+			end
+		end
+	end
+
+	function Promise.delay(seconds)
+		assert(type(seconds) == "number", "Bad argument #1 to Promise.delay, must be a number.")
+		-- If seconds is -INF, INF, or NaN, assume seconds is 0.
+		-- This mirrors the behavior of wait()
+		if seconds < 0 or seconds == math.huge or seconds ~= seconds then
+			seconds = 0
+		end
+
+		return Promise.new(function(resolve, _, onCancel)
+			enqueue(resolve, seconds)
+
+			onCancel(function()
+				dequeue(resolve)
+			end)
+		end)
+	end
+end
+
+--[[
+	Rejects the promise after `seconds` seconds.
+]]
+function Promise.prototype:timeout(seconds, timeoutValue)
+	return Promise.race({
+		Promise.delay(seconds):andThen(function()
+			return Promise.reject(timeoutValue == nil and "Timed out" or timeoutValue)
+		end),
+		self
+	})
+end
+
 function Promise.prototype:getStatus()
 	return self._status
 end