feat(servicebag): Add SetInitFrameBudget to spread Init and Start over frames

src/servicebag/src/Shared/ServiceBag.lua

@@ -67,6 +67,12 @@ export type ServiceBag =
 			_serviceStartLogger: ServiceInitLogger.ServiceInitLogger,
 
 			_destroyingSignal: Signal.Signal<()>,
+
+			_initFrameBudget: number?,
+			_frameStartClock: number,
+			_budgetPauses: number,
+			_budgetLongestStretch: number,
+			_budgetWaitingThreads: { [thread]: true },
 		},
 		{} :: typeof({ __index = ServiceBag })
 	))
@@ -92,6 +98,12 @@ function ServiceBag.new(parentProvider: ServiceBag?): ServiceBag
 	self._initRunAllowed = false
 	self._destructing = false
 
+	self._initFrameBudget = nil
+	self._frameStartClock = 0
+	self._budgetPauses = 0
+	self._budgetLongestStretch = 0
+	self._budgetWaitingThreads = setmetatable({}, { __mode = "k" }) :: any
+
 	self._serviceInitLogger = ServiceInitLogger.new("initialized")
 	self._serviceStartLogger = ServiceInitLogger.new("started")
 
@@ -174,6 +186,62 @@ function ServiceBag.HasService<T>(self: ServiceBag, serviceType: T): boolean
 	end
 end
 
+--[=[
+	Spreads [ServiceBag.Init] and [ServiceBag.Start] over several frames. Before each service
+	initializes or starts, if the current frame has already spent `seconds`, the bag waits a frame
+	first. Services still initialize and start in the same order; this only stops a large bag from
+	doing all of its work in one long frame, such as a client joining with a hundred or more
+	services.
+
+	A service that yields in its own `Init` still errors. With a budget set, [ServiceBag.Init] and
+	[ServiceBag.Start] yield, so call them from a thread that can yield.
+
+	```lua
+	local serviceBag = ServiceBag.new()
+	serviceBag:SetInitFrameBudget(0.012)
+	serviceBag:GetService(require("MyServiceClient"))
+	serviceBag:Init()
+	serviceBag:Start()
+	```
+
+	@param seconds number? -- Work allowed per frame. `nil` (the default) does everything at once.
+]=]
+function ServiceBag.SetInitFrameBudget(self: ServiceBag, seconds: number?): ()
+	assert(seconds == nil or (type(seconds) == "number" and seconds > 0), "Bad seconds")
+	assert(not self._initRunAllowed and self._serviceTypesToInitializeSet, "Set the budget before Init")
+
+	self._initFrameBudget = seconds
+end
+
+--[=[
+	How many frames [ServiceBag.Init] and [ServiceBag.Start] waited because of the frame budget,
+	and the longest stretch of work between two waits, in seconds.
+
+	@return (number, number)
+]=]
+function ServiceBag.GetFrameBudgetStats(self: ServiceBag): (number, number)
+	if self._frameStartClock == 0 then
+		return self._budgetPauses, self._budgetLongestStretch
+	end
+
+	return self._budgetPauses, math.max(self._budgetLongestStretch, os.clock() - self._frameStartClock)
+end
+
+function ServiceBag._pauseIfOverBudget(self: ServiceBag): ()
+	local budget = self._initFrameBudget
+	if not budget or os.clock() - self._frameStartClock < budget then
+		return
+	end
+
+	local thread = coroutine.running()
+	self._budgetPauses += 1
+	self._budgetLongestStretch = math.max(self._budgetLongestStretch, os.clock() - self._frameStartClock)
+	self._budgetWaitingThreads[thread] = true
+	task.wait()
+	self._budgetWaitingThreads[thread] = nil
+	self._frameStartClock = os.clock()
+end
+
 --[=[
 	Initializes the service bag and ensures recursive initialization
 	can occur
@@ -182,6 +250,7 @@ function ServiceBag.Init(self: ServiceBag)
 	assert(not self._initRunAllowed, "Already initializing")
 	assert(self._serviceTypesToInitializeSet, "Already initialized")
 	self._initRunAllowed = true
+	self._frameStartClock = os.clock()
 
 	while next(self._serviceTypesToInitializeSet) do
 		local serviceType = next(self._serviceTypesToInitializeSet)
@@ -201,8 +270,11 @@ function ServiceBag.Start(self: ServiceBag)
 	assert(not self._serviceTypesToInitializeSet, "Not initialized yet. Call serviceBag:Init() first.")
 
 	self._initRunAllowed = false
+	self._frameStartClock = os.clock()
 
 	while next(self._serviceTypesToStartSet) do
+		-- With a frame budget, :Start() is spread over frames too (between services only).
+		self:_pauseIfOverBudget()
 		local serviceType = table.remove(self._serviceTypesToStartSet)
 		local service = assert(self._services[serviceType], "No service")
 		local serviceName = self:_getServiceName(serviceType)
@@ -313,6 +385,7 @@ function ServiceBag._ensureInitialization<T>(self: ServiceBag, serviceType: T)
 		end
 
 		self._initializedServiceTypeSet[serviceType] = true
+		self:_pauseIfOverBudget()
 		self:_initService(serviceType)
 	elseif self._serviceTypesToInitializeSet then
 		self._serviceTypesToInitializeSet[serviceType] = true
@@ -330,20 +403,45 @@ function ServiceBag._initService(self: ServiceBag, serviceType)
 
 	if service.Init then
 		local current
+		local done = false
+		local waiter: thread? = nil
 		task.spawn(function()
 			debug.setmemorycategory(serviceName)
 
 			local stopClock = self._serviceInitLogger:StartInitClock(serviceName)
 
 			current = coroutine.running()
-			service:Init(self)
 
-			stopClock()
+			-- Caught so that an Init that errors still counts as finished, as it does without a
+			-- budget, and a thread waiting on it after a budget pause is resumed instead of hanging.
+			local ok, err = xpcall(service.Init, debug.traceback, service, self)
+			if ok then
+				stopClock()
+			end
+
+			done = true
+			if waiter then
+				task.spawn(waiter)
+			end
+
+			if not ok then
+				error(err, 0)
+			end
 		end)
 
-		local isDead = coroutine.status(current) == "dead"
-		if not isDead then
-			error(debug.traceback(current, string.format("Initializing service %q yielded", serviceName)))
+		if not done then
+			local waitingThreads = self._budgetWaitingThreads
+			if waitingThreads[current] then
+				-- Paused by the frame budget, or waiting on a service that was. Wait for it to finish,
+				-- marked the same way, so that anything waiting on this thread waits too.
+				local thread = coroutine.running()
+				waiter = thread
+				waitingThreads[thread] = true
+				coroutine.yield()
+				waitingThreads[thread] = nil
+			else
+				error(debug.traceback(current, string.format("Initializing service %q yielded", serviceName)))
+			end
 		end
 	end
 

src/servicebag/src/Shared/ServiceBag.spec.lua

@@ -294,3 +294,128 @@ describe("ServiceBag destruction order", function()
 		end
 	end)
 end)
+
+describe("ServiceBag frame budget", function()
+	local function spin(seconds: number): number
+		local start = os.clock()
+		local iterations = 0
+		while os.clock() - start < seconds do
+			iterations += 1
+		end
+		return iterations
+	end
+
+	local function makeSlowService(name: string, events: { string }, cost: number, initDependencies: { any }?)
+		local serviceDefinition: any = makeService(name, events, initDependencies)
+		local init = serviceDefinition.Init
+		function serviceDefinition:Init(serviceBag)
+			spin(cost)
+			init(self, serviceBag)
+		end
+		return serviceDefinition
+	end
+
+	it("spreads Init over frames, keeping dependencies first", function()
+		local events = {}
+		local serviceBag = ServiceBag.new()
+		local leaf = makeSlowService("Leaf", events, 0.003)
+		local middle = makeSlowService("Middle", events, 0.003, { leaf })
+		serviceBag:GetService(makeSlowService("Root", events, 0.003, { middle }))
+		for index = 1, 4 do
+			serviceBag:GetService(makeSlowService("Other" .. index, events, 0.003))
+		end
+
+		serviceBag:SetInitFrameBudget(0.004)
+		serviceBag:Init()
+		serviceBag:Start()
+
+		local pauses = serviceBag:GetFrameBudgetStats()
+		local inits = filterSuffix(events, ".init")
+		local starts = filterSuffix(events, ".start")
+
+		expect(pauses > 0).toBe(true)
+		expect(#inits).toBe(7)
+		expect(#starts).toBe(7)
+		expect(table.find(inits, "Leaf") < table.find(inits, "Middle")).toBe(true)
+		expect(table.find(inits, "Middle") < table.find(inits, "Root")).toBe(true)
+
+		serviceBag:Destroy()
+	end)
+
+	it("doesn't pause without a budget", function()
+		local events = {}
+		local serviceBag = ServiceBag.new()
+		for index = 1, 4 do
+			serviceBag:GetService(makeSlowService("Service" .. index, events, 0.003))
+		end
+
+		serviceBag:Init()
+		serviceBag:Start()
+
+		expect((serviceBag:GetFrameBudgetStats())).toBe(0)
+		expect(#filterSuffix(events, ".init")).toBe(4)
+
+		serviceBag:Destroy()
+	end)
+
+	it("still errors when a service yields in Init", function()
+		local serviceBag = ServiceBag.new()
+		local yielding = {}
+		yielding.ServiceName = "Yielding"
+		function yielding:Init(_serviceBag)
+			task.wait()
+		end
+		serviceBag:GetService(yielding)
+		serviceBag:SetInitFrameBudget(0.004)
+
+		expect(function()
+			serviceBag:Init()
+		end).toThrow("yielded")
+	end)
+
+	it("carries on when a service errors in Init after a budget pause", function()
+		local events = {}
+		local serviceBag = ServiceBag.new()
+		local slow = makeSlowService("Slow", events, 0.006)
+		local fast = makeService("Fast", events)
+		local erroring = {}
+		erroring.ServiceName = "Erroring"
+		function erroring:Init(bag)
+			bag:GetService(slow)
+			bag:GetService(fast)
+			error("Init failed")
+		end
+		serviceBag:GetService(erroring)
+		serviceBag:GetService(makeSlowService("After", events, 0.003))
+		serviceBag:SetInitFrameBudget(0.004)
+
+		local finished = false
+		task.spawn(function()
+			serviceBag:Init()
+			serviceBag:Start()
+			finished = true
+		end)
+		for _ = 1, 60 do
+			if finished then
+				break
+			end
+			task.wait()
+		end
+
+		expect(finished).toBe(true)
+		expect(table.find(filterSuffix(events, ".start"), "After") ~= nil).toBe(true)
+
+		serviceBag:Destroy()
+	end)
+
+	it("can't be set once Init has run", function()
+		local serviceBag = ServiceBag.new()
+		serviceBag:Init()
+
+		expect(function()
+			serviceBag:SetInitFrameBudget(0.004)
+		end).toThrow("Set the budget before Init")
+
+		serviceBag:Destroy()
+	end)
+end)