# signature { license=0, seed=457995617, key="628d11ffa98b65ab9a6665a7ec62a436" }
--
-- Limits resource usage while executing tasks.
-- Tasks added will be run in parallel, with a maximum of 
-- simultaneous tasks to prevent consuming all/too many resources.
-- Every task added will immediately be scheduled (if there is room)
-- using the `wait` method one can wait for completion.

local copas = require("copas")
local pack = table.pack or function(...) return {n=select('#',...),...} end
local unpack = function(t) return (table.unpack or unpack)(t, 1, t.n or #t) end

-- Add a task to the queue, returns the coroutine created
-- identical to `copas.addthread`. Can be called while the 
-- set of tasks is executing.
local function add(self, task, ...)
	local carg = pack(...)
	local coro = copas.addthread(function()
			copas.sleep(-1)							-- go to sleep until being woken
			local suc, err = pcall(task, unpack(carg)) -- start the task
			self:removethread(coroutine.running())		   -- dismiss ourselves
			if not suc then error(err) end			 -- rethrow error
		end)
	table.insert(self.queue, coro)				 -- store in list
	self:next()
	return coro
end

-- remove a task from the queue. Can be called while the 
-- set of tasks is executing. Will NOT stop the task if 
-- it is already running.
local function remove(self, coro)
	self.queue[coro] = nil
	if self.running[coro] then
		-- it is in the already running set
		self.running[coro] = nil
		self.count = self.count - 1
	else
		-- check the queue and remove if found
		for i, item in ipairs(self.queue) do
			if coro == item then 
				table.remove(self.queue, i)
				break
			end
		end	
	end  
	self:next()
end

-- schedules the next task (if any) for execution, signals completeness
local function nxt(self)
	while self.count < self.maxt do
		local coro = self.queue[1]
		if not coro then break end -- queue is empty, so nothing to add
		-- move it to running and restart the task
		table.remove(self.queue, 1)
		self.running[coro] = coro
		self.count = self.count + 1
		copas.wakeup(coro)
	end
	if self.count == 0 and next(self.waiting) then
		-- all tasks done, resume the waiting tasks so they can unblock/return
		for coro in pairs(self.waiting) do
			copas.wakeup(coro)
		end
	end
end

-- Waits for the tasks. Yields until all are finished
local function wait(self)
	if self.count == 0 then return end  -- There's nothing to do...
	local coro = coroutine.running()
	-- now store this coroutine (so we know which to wakeup) and go to sleep
	self.waiting[coro] = true
	copas.sleep(-1)
	self.waiting[coro] = nil
end

-- creats a new tasksrunner, with maximum maxt simultaneous threads
local function new(maxt)
	return {
		maxt = maxt or 99999,   -- max simultaneous tasks
		count = 0,              -- count of running tasks
		queue = {},             -- tasks waiting (list/array)
		running = {},           -- tasks currently running (indexed by coroutine)
		waiting = {},           -- coroutines, waiting for all tasks being finished (indexed by coro)
		addthread = add,
		removethread = remove,
		next = nxt,
		wait = wait,
	}
end

return { new = new }

