actively removes stale timers

src/fibers/op.lua

@@ -80,11 +80,11 @@ function Suspension:wakeup(task)
 end
 
 function Suspension:at_time(t, task)
-	self.sched:schedule_at_time(t, task)
+	return self.sched:schedule_at_time(t, task)
 end
 
 function Suspension:after(dt, task)
-	self.sched:schedule_after_sleep(dt, task)
+	return self.sched:schedule_after_sleep(dt, task)
 end
 
 function Suspension:complete(wrap, ...)

src/fibers/sched.lua

@@ -102,14 +102,14 @@ end
 ---@param t number # absolute time on the scheduler clock
 ---@param task Task
 function Scheduler:schedule_at_time(t, task)
-	self.wheel:add_absolute(t, task)
+	return self.wheel:add_absolute(t, task)
 end
 
 --- Schedule a task after a delay from the wheel's current time.
 ---@param dt number # delay in seconds
 ---@param task Task
 function Scheduler:schedule_after_sleep(dt, task)
-	self.wheel:add_delta(dt, task)
+	return self.wheel:add_delta(dt, task)
 end
 
 --- Ask all registered sources to enqueue any ready tasks.

src/fibers/sleep.lua

@@ -21,7 +21,8 @@ local function deadline_op(t)
 	---@param suspension Suspension
 	---@param wrap_fn WrapFn
 	local function block(suspension, wrap_fn)
-		suspension:at_time(t, suspension:complete_task(wrap_fn))
+		local cancel_timer = suspension:at_time(t, suspension:complete_task(wrap_fn))
+		suspension:add_cleanup(cancel_timer)
 	end
 
 	return op.new_primitive(nil, try, block)

src/fibers/timer.lua

@@ -4,8 +4,11 @@
 ---@module 'fibers.timer'
 
 ---@class TimerNode
----@field time number  # absolute due time (monotonic seconds)
----@field obj any     # scheduled payload
+---@field time number       # absolute due time (monotonic seconds)
+---@field obj any           # scheduled payload
+---@field index integer|nil # current heap index, nil when not queued
+
+---@alias TimerCancel fun(): boolean
 
 local floor, huge = math.floor, math.huge
 
@@ -21,10 +24,20 @@ local function new_heap()
 	return setmetatable({ heap = {}, size = 0 }, Heap)
 end
 
+---@param i integer
+---@param j integer
+function Heap:swap(i, j)
+	local heap = self.heap
+	heap[i], heap[j] = heap[j], heap[i]
+	heap[i].index = i
+	heap[j].index = j
+end
+
 ---@param node TimerNode
 function Heap:push(node)
 	local size = self.size + 1
 	self.size = size
+	node.index = size
 	self.heap[size] = node
 	self:heapify_up(size)
 end
@@ -38,21 +51,62 @@ function Heap:pop()
 
 	local heap = self.heap
 	local root = heap[1]
+	root.index = nil
 
 	if size == 1 then
 		heap[1] = nil
 		self.size = 0
 		return root
 	end
 
-	heap[1] = heap[size]
+	local last = heap[size]
 	heap[size] = nil
 	self.size = size - 1
+	heap[1] = last
+	last.index = 1
 	self:heapify_down(1)
 
 	return root
 end
 
+---@param node TimerNode
+---@return boolean
+function Heap:remove(node)
+	local idx = node.index
+	if type(idx) ~= 'number' or idx < 1 or idx > self.size then
+		return false
+	end
+
+	local heap = self.heap
+	if heap[idx] ~= node then
+		return false
+	end
+
+	local size = self.size
+	node.index = nil
+
+	if idx == size then
+		heap[size] = nil
+		self.size = size - 1
+		return true
+	end
+
+	local last = heap[size]
+	heap[size] = nil
+	self.size = size - 1
+	heap[idx] = last
+	last.index = idx
+
+	local parent = floor(idx / 2)
+	if idx > 1 and heap[idx].time < heap[parent].time then
+		self:heapify_up(idx)
+	else
+		self:heapify_down(idx)
+	end
+
+	return true
+end
+
 ---@param idx integer
 function Heap:heapify_up(idx)
 	local heap = self.heap
@@ -61,7 +115,7 @@ function Heap:heapify_up(idx)
 		if heap[parent].time <= heap[idx].time then
 			break
 		end
-		heap[parent], heap[idx] = heap[idx], heap[parent]
+		self:swap(parent, idx)
 		idx = parent
 	end
 end
@@ -87,7 +141,7 @@ function Heap:heapify_down(idx)
 			break
 		end
 
-		heap[idx], heap[smallest] = heap[smallest], heap[idx]
+		self:swap(idx, smallest)
 		idx = smallest
 	end
 end
@@ -108,15 +162,30 @@ end
 --- Schedule an object at absolute time t.
 ---@param t number   # absolute due time
 ---@param obj any    # payload to pass to the scheduler
+---@return TimerCancel cancel # idempotent cancellation handle
 function Timer:add_absolute(t, obj)
-	self.heap:push { time = t, obj = obj }
+	local node = { time = t, obj = obj, index = nil }
+	self.heap:push(node)
+
+	local cancelled = false
+	return function ()
+		if cancelled then
+			return false
+		end
+
+		cancelled = true
+		local removed = self.heap:remove(node)
+		node.obj = nil
+		return removed
+	end
 end
 
 --- Schedule an object after a delay from the current timer time.
 ---@param dt number  # delay in seconds from self.now
 ---@param obj any    # payload to pass to the scheduler
+---@return TimerCancel cancel # idempotent cancellation handle
 function Timer:add_delta(dt, obj)
-	self:add_absolute(self.now + dt, obj)
+	return self:add_absolute(self.now + dt, obj)
 end
 
 --- Get the time of the next scheduled entry, or math.huge if none exist.
@@ -140,8 +209,10 @@ function Timer:advance(t, sched)
 
 	while heap.size > 0 and t >= heap.heap[1].time do
 		local node = assert(heap:pop()) -- non-nil since size>0
+		local obj = node.obj
+		node.obj = nil
 		self.now = node.time
-		sched:schedule(node.obj)
+		sched:schedule(obj)
 	end
 
 	self.now = t

tests/test.lua

@@ -22,6 +22,7 @@ local modules = {
 	{ 'pulse' },
 	{ 'cond' },
 	{ 'sleep' },
+	{ 'sleep', 'timer_cancel' },
 	{ 'waitgroup' },
 	{ 'scope' },
 }

tests/test_sleep-timer_cancel.lua

@@ -0,0 +1,57 @@
+--- Regression test for losing sleep arms in choices.
+---
+--- A sleep operation that loses a choice must cancel and remove its scheduled
+--- timer task immediately. Otherwise the timer heap retains the CompleteTask,
+--- which retains the Suspension and its captured values until the deadline.
+print('testing: fibers.sleep timer cancellation')
+
+package.path = '../src/?.lua;' .. package.path
+
+local fibers  = require 'fibers'
+local op      = require 'fibers.op'
+local sleep   = require 'fibers.sleep'
+local runtime = require 'fibers.runtime'
+
+local count = 2000
+
+local function next_turn_op(value)
+	return op.new_primitive(nil,
+		function ()
+			return false
+		end,
+		function (suspension, wrap_fn)
+			suspension:wakeup(suspension:complete_task(wrap_fn, value))
+		end
+	)
+end
+
+local completed = 0
+
+fibers.run(function ()
+	local wheel = runtime.current_scheduler.wheel
+	assert(wheel.heap.size == 0, 'timer heap unexpectedly non-empty at start')
+
+	for i = 1, count do
+		local value = fibers.perform(fibers.choice(
+			next_turn_op(i),
+			sleep.sleep_op(1e6)
+		))
+
+		assert(value == i)
+		completed = completed + 1
+
+		assert(wheel.heap.size == 0,
+			('losing sleep_op left %d timer entries after iteration %d')
+			:format(wheel.heap.size, i))
+
+		if i % 250 == 0 then
+			collectgarbage('collect')
+			assert(wheel.heap.size == 0,
+				('timer heap retained entries after GC at iteration %d'):format(i))
+		end
+	end
+end)
+
+assert(completed == count)
+
+print('test: ok')