From f8a29022c93c63f4d344588443c5c70781b7532e Mon Sep 17 00:00:00 2001 From: Rich Thanki Date: Mon, 1 Jun 2026 08:04:02 +0000 Subject: [PATCH] removes stale timers --- src/fibers/op.lua | 4 +- src/fibers/sched.lua | 4 +- src/fibers/sleep.lua | 3 +- src/fibers/timer.lua | 87 ++++++++++++++++++++++++++++--- tests/test.lua | 1 + tests/test_sleep-timer_cancel.lua | 57 ++++++++++++++++++++ 6 files changed, 143 insertions(+), 13 deletions(-) create mode 100644 tests/test_sleep-timer_cancel.lua diff --git a/src/fibers/op.lua b/src/fibers/op.lua index fe811eb..e52812f 100644 --- a/src/fibers/op.lua +++ b/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, ...) diff --git a/src/fibers/sched.lua b/src/fibers/sched.lua index 8700e4c..7c80dd6 100644 --- a/src/fibers/sched.lua +++ b/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. diff --git a/src/fibers/sleep.lua b/src/fibers/sleep.lua index 5587b77..de84eb0 100644 --- a/src/fibers/sleep.lua +++ b/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) diff --git a/src/fibers/timer.lua b/src/fibers/timer.lua index d49ce8f..e4e3b90 100644 --- a/src/fibers/timer.lua +++ b/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,6 +51,7 @@ function Heap:pop() local heap = self.heap local root = heap[1] + root.index = nil if size == 1 then heap[1] = nil @@ -45,14 +59,54 @@ function Heap:pop() 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 diff --git a/tests/test.lua b/tests/test.lua index 0ea7ac7..01418d5 100644 --- a/tests/test.lua +++ b/tests/test.lua @@ -22,6 +22,7 @@ local modules = { { 'pulse' }, { 'cond' }, { 'sleep' }, + { 'sleep', 'timer_cancel' }, { 'waitgroup' }, { 'scope' }, } diff --git a/tests/test_sleep-timer_cancel.lua b/tests/test_sleep-timer_cancel.lua new file mode 100644 index 0000000..8f4eb4a --- /dev/null +++ b/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')