diff --git a/src/fibers/mailbox.lua b/src/fibers/mailbox.lua new file mode 100644 index 0000000..48329a4 --- /dev/null +++ b/src/fibers/mailbox.lua @@ -0,0 +1,304 @@ +-- fibers/mailbox.lua +-- +-- Mailbox: closeable, drainable queue for fibers. +-- +-- Conventions +-- * nil payloads are forbidden; nil is reserved for end-of-stream. +-- * rx:recv() returns: +-- - a non-nil message, or +-- - nil when the mailbox is closed and drained. +-- rx:why() yields the close reason (if any). +-- * tx:send(v) returns: +-- - true if accepted/delivered, +-- - nil if the mailbox is closed (send rejected). +-- tx:why() yields the close reason (if any). +-- * Multi-producer: +-- - tx:clone() creates a new counted sender handle. +-- - each counted handle should be closed once finished. +-- - mailbox closes-for-send when the last counted handle closes. + +---@module 'fibers.mailbox' + +local op = require 'fibers.op' +local fifo = require 'fibers.utils.fifo' +local perform = require 'fibers.performer'.perform + +---@alias MailboxWant nil -- reserved for future extensions + +---@class MailboxState +---@field cap integer +---@field buf any|nil -- FIFO buffer when cap>0; nil for rendezvous +---@field getq any -- FIFO of waiting receivers +---@field putq any -- FIFO of waiting senders +---@field closed boolean +---@field reason any|nil +---@field senders integer -- counted sender handles still open + +---@class MailboxTx +---@field _st MailboxState +---@field _closed boolean -- this handle closed (idempotent) +---@field _counted boolean -- whether this handle contributes to st.senders +local Tx = {} +Tx.__index = Tx + +---@class MailboxRx +---@field _st MailboxState +local Rx = {} +Rx.__index = Rx + +---------------------------------------------------------------------- +-- Internal helpers +---------------------------------------------------------------------- + +--- Pop the next entry whose suspension is still waiting, if any. +---@param q any +---@return table|nil +local function pop_active(q) + while not q:empty() do + local e = q:pop() + local s = e.suspension + if not s or s:waiting() then + return e + end + end +end + +---@param st MailboxState +---@param reason any|nil +local function record_reason(st, reason) + if st.reason == nil and reason ~= nil then + st.reason = reason + end +end + +--- Close the mailbox state (idempotent), record reason, and wake blocked parties. +--- Receivers drain buffered values (if any), then receive nil. +--- Waiting senders are rejected (nil). +---@param st MailboxState +---@param reason any|nil +local function close_state(st, reason) + if st.closed then + record_reason(st, reason) + return + end + + st.closed = true + record_reason(st, reason) + + -- Wake receivers: deliver buffered values first, then nil when buffer empty. + while true do + local recv = pop_active(st.getq) + if not recv then break end + + local v + if st.buf and st.buf:length() > 0 then + v = st.buf:pop() + end + recv.suspension:complete(recv.wrap, v) + end + + -- Reject senders. + while true do + local snd = pop_active(st.putq) + if not snd then break end + snd.suspension:complete(snd.wrap, nil) + end +end + +---------------------------------------------------------------------- +-- Construction +---------------------------------------------------------------------- + +--- Create a mailbox. Returns (tx, rx). +---@param capacity? integer # 0 or nil -> rendezvous; >0 -> buffered capacity +---@return MailboxTx tx, MailboxRx rx +local function new(capacity) + capacity = capacity or 0 + assert(type(capacity) == 'number' and capacity >= 0, 'mailbox.new: capacity must be >= 0') + + ---@type MailboxState + local st = { + cap = capacity, + buf = (capacity > 0) and fifo.new() or nil, + getq = fifo.new(), + putq = fifo.new(), + closed = false, + reason = nil, + senders = 1, + } + + local tx = setmetatable({ _st = st, _closed = false, _counted = true }, Tx) + local rx = setmetatable({ _st = st }, Rx) + return tx, rx +end + +---------------------------------------------------------------------- +-- Tx (sender) +---------------------------------------------------------------------- + +--- Return the mailbox close reason (if any). +---@return any|nil +function Tx:why() + return self._st.reason +end + +--- Clone this sender handle (multi-producer). +--- If the mailbox or this handle is closed, returns an inert, uncounted handle. +---@return MailboxTx +function Tx:clone() + local st = self._st + if self._closed or st.closed then + return setmetatable({ _st = st, _closed = true, _counted = false }, Tx) + end + st.senders = st.senders + 1 + return setmetatable({ _st = st, _closed = false, _counted = true }, Tx) +end + +--- Close this sender handle (idempotent). +--- When the last counted sender closes, the mailbox closes-for-send. +---@param reason any|nil +---@return boolean ok +function Tx:close(reason) + local st = self._st + record_reason(st, reason) + + if self._closed then return true end + + self._closed = true + + -- If already uncounted, or mailbox already closed, nothing to do. + if not self._counted or st.closed then + self._counted = false + return true + end + + self._counted = false + st.senders = st.senders - 1 + if st.senders <= 0 then + st.senders = 0 + close_state(st, reason) + end + + return true +end + +--- Op that sends a message. +--- When performed: true on success, nil when closed (send rejected). +---@param v any # MUST NOT be nil +---@return Op +function Tx:send_op(v) + assert(v ~= nil, 'mailbox.send: nil payload is not permitted') + + local st = self._st + local getq, putq, buf, cap = st.getq, st.putq, st.buf, st.cap + + local function try() + if st.closed or self._closed then return true, nil end + + -- Rendezvous with a waiting receiver. + local recv = pop_active(getq) + if recv then + recv.suspension:complete(recv.wrap, v) + return true, true + end + + -- Buffered enqueue. + if buf and buf:length() < cap then + buf:push(v) + return true, true + end + + return false + end + + ---@param suspension Suspension + ---@param wrap_fn WrapFn + local function block(suspension, wrap_fn) + if st.closed or self._closed then + return suspension:complete(wrap_fn, nil) + end + putq:push { val = v, suspension = suspension, wrap = wrap_fn } + end + + return op.new_primitive(nil, try, block) +end + +--- Synchronously send a message. +---@param v any +---@return boolean|nil ok +function Tx:send(v) + return perform(self:send_op(v)) +end + +---------------------------------------------------------------------- +-- Rx (receiver) +---------------------------------------------------------------------- + +--- Return the mailbox close reason (if any). +---@return any|nil +function Rx:why() + return self._st.reason +end + +--- Op that receives the next message. +--- When performed: a non-nil value, or nil when closed and drained. +---@return Op +function Rx:recv_op() + local st = self._st + local getq, putq, buf = st.getq, st.putq, st.buf + + local function try() + -- Prefer unblocking a waiting sender (if present); we may still return + -- a buffered value first. + local snd = pop_active(putq) + if snd then snd.suspension:complete(snd.wrap, true) end + + if buf and buf:length() > 0 then + local v = buf:pop() + if snd then buf:push(snd.val) end + return true, v + end + + if snd then return true, snd.val end + + if st.closed then return true, nil end + + return false + end + + ---@param suspension Suspension + ---@param wrap_fn WrapFn + local function block(suspension, wrap_fn) + if st.closed then + return suspension:complete(wrap_fn, nil) + end + getq:push { suspension = suspension, wrap = wrap_fn } + end + + return op.new_primitive(nil, try, block) +end + +--- Synchronously receive the next message. +---@return any|nil v +function Rx:recv() + return perform(self:recv_op()) +end + +--- Iterator over received messages, ending at nil (closed and drained). +---@return fun(): any|nil +function Rx:iter() + return function () + return self:recv() + end +end + +---------------------------------------------------------------------- +-- Public API +---------------------------------------------------------------------- + +return { + new = new, + + Tx = Tx, + Rx = Rx, +} diff --git a/tests/test.lua b/tests/test.lua index 0e4dcf8..ddaf27a 100644 --- a/tests/test.lua +++ b/tests/test.lua @@ -17,6 +17,7 @@ local modules = { { 'sched' }, { 'runtime' }, { 'channel' }, + { 'mailbox' }, { 'cond' }, { 'sleep' }, { 'waitgroup' }, diff --git a/tests/test_mailbox.lua b/tests/test_mailbox.lua new file mode 100644 index 0000000..65b049a --- /dev/null +++ b/tests/test_mailbox.lua @@ -0,0 +1,319 @@ +print('testing: fibers.mailbox') + +package.path = '../src/?.lua;' .. package.path + +local fibers = require 'fibers' +local mailbox = require 'fibers.mailbox' +local sleep = require 'fibers.sleep' +local op = require 'fibers.op' +local wg_mod = require 'fibers.waitgroup' + +---------------------------------------------------------------------- +-- helpers +---------------------------------------------------------------------- + +local function assert_eq(a, b, msg) + if a ~= b then + error((msg or 'assert_eq failed') .. (': expected ' .. tostring(b) .. ', got ' .. tostring(a)), 2) + end +end + +local function assert_truthy(v, msg) + if not v then + error(msg or 'assert_truthy failed', 2) + end +end + +local function assert_falsy(v, msg) + if v then + error(msg or 'assert_falsy failed', 2) + end +end + +local function collect_iter(rx) + local out = {} + for v in rx:iter() do + out[#out + 1] = v + end + return out +end + +---------------------------------------------------------------------- +-- tests +---------------------------------------------------------------------- + +local function test_rendezvous_basic() + local tx, rx = mailbox.new(0) + + local wg = wg_mod.new() + wg:add(2) + + local got = {} + + fibers.spawn(function () + -- receiver: read one, then see end + got[1] = rx:recv() + got[2] = rx:recv() + wg:done() + end) + + fibers.spawn(function () + sleep.sleep(0.02) -- ensure receiver blocks first + local ok = tx:send('hello') + assert_eq(ok, true, 'send should succeed') + tx:close('done') + wg:done() + end) + + wg:wait() + + assert_eq(got[1], 'hello', 'receiver should get message') + assert_eq(got[2], nil, 'receiver should see end after close+drain') + assert_eq(rx:why(), 'done', 'receiver should see close reason') +end + +local function test_buffered_fifo_order_and_drain() + local tx, rx = mailbox.new(4) + + local ok1 = tx:send(1) + local ok2 = tx:send(2) + local ok3 = tx:send(3) + assert_eq(ok1, true) + assert_eq(ok2, true) + assert_eq(ok3, true) + + tx:close('finished') + + local xs = collect_iter(rx) + assert_eq(#xs, 3) + assert_eq(xs[1], 1) + assert_eq(xs[2], 2) + assert_eq(xs[3], 3) + + assert_eq(rx:why(), 'finished') +end + +local function test_send_after_close_is_rejected() + local tx, rx = mailbox.new(0) + + tx:close('no more') + + local ok = tx:send('x') + assert_eq(ok, nil, 'send after close should be rejected') + + local v = rx:recv() + assert_eq(v, nil, 'recv on closed+empty should return nil') + assert_eq(rx:why(), 'no more') +end + +local function test_nil_payload_is_error() + local tx, rx = mailbox.new(0) + + local ok, err = pcall(function () + tx:send(nil) + end) + + assert_falsy(ok, 'sending nil should error') + assert_truthy(err ~= nil, 'expected an error object/message') + + -- close so receiver is not left blocking if someone uses it later + tx:close('end') + assert_eq(rx:recv(), nil) +end + +local function test_multi_producer_clone_and_reason_first_non_nil_wins() + local tx, rx = mailbox.new(8) + local tx2 = tx:clone() + + local wg = wg_mod.new() + wg:add(2) + + fibers.spawn(function () + for i = 1, 3 do + assert_eq(tx:send('a' .. i), true) + end + -- close with a reason first + tx:close('r1') + wg:done() + end) + + fibers.spawn(function () + for i = 1, 2 do + assert_eq(tx2:send('b' .. i), true) + end + -- later close with another reason; should not override + sleep.sleep(0.02) + tx2:close('r2') + wg:done() + end) + + wg:wait() + + local got = collect_iter(rx) + assert_eq(#got, 5) + -- Do not assert interleaving between producers, but do assert membership. + local seen = {} + for _, v in ipairs(got) do seen[v] = (seen[v] or 0) + 1 end + assert_eq(seen['a1'], 1) + assert_eq(seen['a2'], 1) + assert_eq(seen['a3'], 1) + assert_eq(seen['b1'], 1) + assert_eq(seen['b2'], 1) + + assert_eq(rx:why(), 'r1', 'first non-nil close reason should win') +end + +local function test_choice_timeout_then_message() + local tx, rx = mailbox.new(0) + + -- First, no message available: should timeout. + local tag = fibers.perform( + op.choice( + rx:recv_op():wrap(function (v) return 'msg', v end), + sleep.sleep_op(0.03):wrap(function () return 'timeout' end) + ) + ) + assert_eq(tag, 'timeout', 'expected timeout when mailbox empty') + + -- Now arrange a message, then race again. + local wg = wg_mod.new() + wg:add(1) + fibers.spawn(function () + sleep.sleep(0.01) + tx:send('ok') + tx:close('done') + wg:done() + end) + + local tag2, v2 = fibers.perform( + op.choice( + rx:recv_op():wrap(function (v) return 'msg', v end), + sleep.sleep_op(0.20):wrap(function () return 'timeout', nil end) + ) + ) + assert_eq(tag2, 'msg') + assert_eq(v2, 'ok') + + -- Drain end-of-stream. + assert_eq(rx:recv(), nil) + assert_eq(rx:why(), 'done') + wg:wait() +end + +local function test_choice_cancels_blocked_send_and_does_not_deliver() + local tx, rx = mailbox.new(0) + + -- Start a choice that would block on send, but we time it out. + local tag = fibers.perform( + op.choice( + tx:send_op('BAD'):wrap(function (ok) return 'sent', ok end), + sleep.sleep_op(0.03):wrap(function () return 'timeout' end) + ) + ) + assert_eq(tag, 'timeout') + + -- Now send a real message with a receiver; must not receive 'BAD'. + local wg = wg_mod.new() + wg:add(2) + + local got + + fibers.spawn(function () + got = rx:recv() + wg:done() + end) + + fibers.spawn(function () + sleep.sleep(0.02) + assert_eq(tx:send('GOOD'), true) + tx:close('end') + wg:done() + end) + + wg:wait() + + assert_eq(got, 'GOOD', 'cancelled send must not be delivered later') + assert_eq(rx:recv(), nil) +end + +local function test_close_wakes_blocked_receiver() + local tx, rx = mailbox.new(0) + + local wg = wg_mod.new() + wg:add(1) + + local got + + fibers.spawn(function () + got = rx:recv() -- should block until close + wg:done() + end) + + sleep.sleep(0.02) + tx:close('closed without messages') + + wg:wait() + + assert_eq(got, nil, 'blocked receiver should wake and see end-of-stream') + assert_eq(rx:why(), 'closed without messages') +end + +local function test_close_wakes_blocked_sender() + local tx, rx = mailbox.new(0) + local tx2 = tx:clone() + + local wg = wg_mod.new() + wg:add(1) + + local send_ok + + fibers.spawn(function () + -- No receiver: this blocks until mailbox closes. + send_ok = tx2:send('will-not-deliver') + wg:done() + end) + + -- Ensure sender is blocked. + sleep.sleep(0.03) + + -- Closing both handles causes mailbox closure; blocked send should return nil. + tx:close('shutdown') + tx2:close('shutdown') + + wg:wait() + + assert_eq(send_ok, nil, 'blocked send should be rejected when mailbox closes') + assert_eq(rx:recv(), nil, 'receiver should see end-of-stream') + assert_eq(rx:why(), 'shutdown') +end + +local function test_clone_after_close_is_inert() + local tx, rx = mailbox.new(0) + tx:close('done') + + local tx2 = tx:clone() + assert_eq(tx2:send('x'), nil, 'clone after close should not send') + assert_eq(rx:recv(), nil) + assert_eq(rx:why(), 'done') +end + +---------------------------------------------------------------------- +-- main +---------------------------------------------------------------------- + +local function main() + test_rendezvous_basic() + test_buffered_fifo_order_and_drain() + test_send_after_close_is_rejected() + test_nil_payload_is_error() + test_multi_producer_clone_and_reason_first_non_nil_wins() + test_choice_timeout_then_message() + test_choice_cancels_blocked_send_and_does_not_deliver() + test_close_wakes_blocked_receiver() + test_close_wakes_blocked_sender() + test_clone_after_close_is_inert() + + print('All mailbox tests passed!') +end + +fibers.run(main)