Blox deviations

modules/signals-scheduler/src/SignalsScheduler.lua

@@ -8,31 +8,66 @@ local isContinuing = false
 
 local continuations: { work } = {}
 
-local function runWork(fn: work)
-	fn()
-	for _, work in continuations do
-		work()
+--[[
+	Counters for a profiler, supplied by the host rather than required, so nothing here
+	knows about the runtime using it. Bound as upvalues rather than read out of a table
+	per batch, so a hook that is not installed costs what a constant did.
+]]
+export type Hooks = {
+	onBatch: (() -> ())?,
+	onContinuations: ((count: number) -> ())?,
+}
+
+local function noop() end
+
+local onBatch: () -> () = noop
+local onContinuations: (count: number) -> () = noop
+
+local function setHooks(hooks: Hooks?)
+	local given: any = if hooks ~= nil then hooks else {}
+	onBatch = given.onBatch or noop
+	onContinuations = given.onContinuations or noop
+end
+
+--[[
+	Both drains below are indexed rather than iterated.
+
+	Work run during a drain can call `schedule`, which appends to the very table being
+	walked. Generalised iteration is not guaranteed to reach entries added after it
+	started, so an effect scheduled by another effect could be dropped for the rest of
+	the batch -- and `table.clear` below would then discard it entirely. Indexing by
+	position re-reads the length each time round and picks those up.
+]]
+local function runWork()
+	local i = 1
+	while i <= #continuations do
+		continuations[i]()
+		i += 1
 	end
 end
 
 local function runContinuations(): (boolean, any)
 	local firstError: any = nil
-	for _, work in continuations do
-		local ok, err: any = xpcall(work, debug.traceback)
+	local i = 1
+	while i <= #continuations do
+		local ok, err: any = xpcall(continuations[i], debug.traceback)
 		if not ok and firstError == nil then
 			firstError = err
 		end
+		i += 1
 	end
 
 	return firstError == nil, firstError
 end
 
 local function batch(fn: work)
 	if not isContinuing then
+		onBatch()
 		isContinuing = true
 
 		if SignalsSchedulerResetStateAfterErrors then
 			local ok, err: any = xpcall(fn, debug.traceback)
+			onContinuations(#continuations)
 			local continuationsOk, continuationsErr: any = runContinuations()
 
 			table.clear(continuations)
@@ -45,7 +80,9 @@ local function batch(fn: work)
 				error(continuationsErr, 0)
 			end
 		else
-			runWork(fn)
+			fn()
+			onContinuations(#continuations)
+			runWork()
 
 			table.clear(continuations)
 			isContinuing = false
@@ -67,4 +104,5 @@ return {
 	batch = batch,
 	flush = flush,
 	schedule = schedule,
+	setHooks = setHooks,
 }

modules/signals-scheduler/src/__tests__/SignalsScheduler.test.lua

@@ -136,6 +136,28 @@ it("should support batching scheduled work", function()
 	})
 end)
 
+it("should run work that is scheduled while draining", function()
+	local order = {}
+
+	batch(function()
+		schedule(function()
+			table.insert(order, "first")
+			-- Appends to the queue that is being drained right now. An effect
+			-- scheduling another effect is the ordinary case, so this has to be
+			-- picked up by the drain already in progress rather than left for a
+			-- batch that may never come.
+			schedule(function()
+				table.insert(order, "second")
+				schedule(function()
+					table.insert(order, "third")
+				end)
+			end)
+		end)
+	end)
+
+	expect(order).toEqual({ "first", "second", "third" })
+end)
+
 if SignalsSchedulerResetStateAfterErrors then
 	it("should drain scheduled work after batched work errors", function()
 		local counters = {

modules/signals-scheduler/src/init.lua

@@ -1,5 +1,6 @@
 local SignalsScheduler = require(script.SignalsScheduler)
 
 export type work = SignalsScheduler.work
+export type Hooks = SignalsScheduler.Hooks
 
 return SignalsScheduler