diff --git a/src/fibers/channel.lua b/src/fibers/channel.lua index f665da3..696f114 100644 --- a/src/fibers/channel.lua +++ b/src/fibers/channel.lua @@ -1,24 +1,15 @@ --- Use of this source code is governed by the Apache 2.0 license; see COPYING. - ---- fibers.channel module +-- fibers.channel -- Provides Concurrent ML style channels for communication between fibers. --- @module fibers.channel - -local op = require 'fibers.op' -local fifo = require 'fibers.utils.fifo' +local op = require 'fibers.op' +local fifo = require 'fibers.utils.fifo' local perform = require 'fibers.performer'.perform ---- Channel class --- Represents a communication channel between fibers. --- @type Channel local Channel = {} Channel.__index = Channel ---- Create a new Channel. --- @treturn Channel The created Channel. local function new(buffer_size) - buffer_size = buffer_size or 0 -- Default to unbuffered + buffer_size = buffer_size or 0 local buffer = nil if buffer_size > 0 then @@ -26,111 +17,143 @@ local function new(buffer_size) end return setmetatable({ - buffer = buffer, + buffer = buffer, buffer_size = buffer_size, - getq = fifo.new(), -- Queue of waiting receivers - putq = fifo.new(), -- Queue of waiting senders + getq = fifo.new(), -- waiting receivers + putq = fifo.new(), -- waiting senders }, Channel) end ---- Create a put operation for the Channel. --- Make an operation that if and when it completes will rendezvous with --- a receiver fiber to send VAL over the channel. --- @param val The value to put into the Channel. --- @treturn BaseOp The created put operation. +---------------------------------------------------------------------- +-- Helpers: pop active entries (skip cancelled ones) +---------------------------------------------------------------------- + +local function pop_active(q) + while not q:empty() do + local entry = q:pop() + if not entry.cancelled then + return entry + end + end + return nil +end + +---------------------------------------------------------------------- +-- PUT side +---------------------------------------------------------------------- + function Channel:put_op(val) - local getq, putq, buffer, buffer_size = self.getq, self.putq, self.buffer, self.buffer_size + local getq, putq = self.getq, self.putq + local buffer, buffer_size = self.buffer, self.buffer_size + + -- Per-operation state for this event instance. + local entry = { + val = val, + cancelled = false, + suspension = nil, + wrap = nil, + } + local function try() - -- Case 1: If there's a waiting receiver, complete the rendezvous immediately - while not getq:empty() do - local remote = getq:pop() - if remote.suspension:waiting() then - remote.suspension:complete(remote.wrap, val) - return true - end - -- Otherwise the remote suspension is already completed, pop and continue + -- Case 1: rendezvous with a waiting receiver. + local recv = pop_active(getq) + if recv then + recv.suspension:complete(recv.wrap, val) + return true end - -- Case 2: If we have a buffer with space, add the value to the buffer + -- Case 2: buffered channel with available space. if buffer and buffer:length() < buffer_size then buffer:push(val) return true end - -- Case 3: No receivers and no buffer space + -- Case 3: no receiver and no buffer space. return false end + local function block(suspension, wrap_fn) - -- First, GC for canceled operations - while not putq:empty() and not putq:peek().suspension:waiting() do - putq:pop() - end - -- No space in buffer and no receivers, so block - putq:push({ suspension = suspension, wrap = wrap_fn, val = val }) + entry.suspension = suspension + entry.wrap = wrap_fn + entry.cancelled = false + putq:push(entry) end - return op.new_primitive(nil, try, block) + + local prim = op.new_primitive(nil, try, block) + + -- Mark this entry as cancelled if this put loses in a choice / is aborted. + return prim:on_abort(function() + entry.cancelled = true + end) end ---- Create a get operation for the Channel. --- Make an operation that if and when it completes will rendezvous with --- a sender fiber to receive one value from the channel. --- @treturn BaseOp The created get operation. +---------------------------------------------------------------------- +-- GET side +---------------------------------------------------------------------- + function Channel:get_op() - local getq, putq, buffer = self.getq, self.putq, self.buffer - -- Attempt to service one sender from putq - local function pop_from_putq() - while not putq:empty() do - local remote = putq:pop() - if remote.suspension:waiting() then - remote.suspension:complete(remote.wrap) - return remote - end + local getq, putq = self.getq, self.putq + local buffer = self.buffer + + local entry = { + cancelled = false, + suspension = nil, + wrap = nil, + } + + local function pop_sender() + local sender = pop_active(putq) + if not sender then + return nil end - return nil + -- Having chosen this sender, complete its suspension immediately. + sender.suspension:complete(sender.wrap) + return sender end + local function try() - local remote = pop_from_putq() - -- Case 1: Take from buffer if available + local remote = pop_sender() + -- Case 1: take from buffer if there is a buffered value. if buffer and buffer:length() > 0 then - local val = buffer:pop() - -- Attempt to refill buffer with one sender - if remote then buffer:push(remote.val) end - return true, val + local v = buffer:pop() + -- If there was a sender waiting, refill the buffer with its value. + if remote then + buffer:push(remote.val) + end + return true, v + end + -- Case 2: no buffered value; take directly from a sender. + if remote then + return true, remote.val end - -- Case 2: No buffer value; try to take directly from a sender - if remote then return true, remote.val end - -- Case 3: Nothing available so block + -- Case 3: nothing available. return false end + local function block(suspension, wrap_fn) - -- Clear any stale entries - while not getq:empty() and not getq:peek().suspension:waiting() do - getq:pop() - end - -- Block this receiver - getq:push({ suspension = suspension, wrap = wrap_fn }) + entry.suspension = suspension + entry.wrap = wrap_fn + entry.cancelled = false + getq:push(entry) end - return op.new_primitive(nil, try, block) + + local prim = op.new_primitive(nil, try, block) + + return prim:on_abort(function() + entry.cancelled = true + end) end ---- Put a message into the Channel. --- Send message on the channel. If there is already another fiber --- waiting to receive a message on this channel, give it our message and --- continue. Otherwise, block until a receiver becomes available. --- @tparam any message The message to put into the Channel. +---------------------------------------------------------------------- +-- Synchronous wrappers +---------------------------------------------------------------------- + function Channel:put(message) return perform(self:put_op(message)) end ---- Get a message from the Channel. --- Receive a message from the channel and return it. If there is --- already another fiber waiting to send a message on this channel, take --- its message directly. Otherwise, block until a sender becomes --- available. --- @treturn any The message retrieved from the Channel. function Channel:get() return perform(self:get_op()) end ---- @export return { - new = new + new = new, } diff --git a/src/fibers/waitgroup.lua b/src/fibers/waitgroup.lua index a0dbaa2..bd75faa 100644 --- a/src/fibers/waitgroup.lua +++ b/src/fibers/waitgroup.lua @@ -1,26 +1,41 @@ --- waitgroup.lua -local op = require 'fibers.op' - +-- fibers/waitgroup.lua +local op = require 'fibers.op' local perform = require 'fibers.performer'.perform local Waitgroup = {} Waitgroup.__index = Waitgroup local function new() - -- Use the core condition primitive directly: { wait_op = ..., signal = ... } - local wg = setmetatable({ + return setmetatable({ _counter = 0, - _cond = op.new_cond(), + _cond = nil, -- per-generation condition; nil when idle }, Waitgroup) - return wg end -function Waitgroup:add(count) - self._counter = self._counter + count - if self._counter < 0 then +function Waitgroup:add(delta) + if delta == 0 then + return + end + + local old_waiters = self._counter + local new_waiters = old_waiters + delta + + if new_waiters < 0 then error("waitgroup counter goes negative") - elseif self._counter == 0 then - self._cond.signal() + end + + self._counter = new_waiters + + if new_waiters == 0 then + -- This generation completes: wake any waiters. + if self._cond then + self._cond.signal() + end + -- _cond remains a triggered cond for this generation; a new + -- generation will allocate a fresh cond when counter rises from 0. + elseif old_waiters == 0 and new_waiters > 0 then + -- Starting a new generation: new condition for new work. + self._cond = op.new_cond() end end @@ -29,20 +44,25 @@ function Waitgroup:done() end function Waitgroup:wait_op() - -- Take the underlying cond's wait op (a BaseOp). - local cond_op = self._cond.wait_op() local function try() return self._counter == 0 end local function block(suspension, wrap_fn) + -- Re-check after try(): counter may have become zero meanwhile. if self._counter == 0 then - -- Became zero after try() but before block() ran. suspension:complete(wrap_fn) - else - -- Delegate blocking to the underlying condition's block_fn. - cond_op.block_fn(suspension, wrap_fn) + return end + + -- At this point we are in an active generation. + local cond = self._cond + if not cond then + error("waitgroup internal error: missing condition for active generation") + end + + local cond_op = cond.wait_op() + cond_op.block_fn(suspension, wrap_fn) end return op.new_primitive(nil, try, block) @@ -53,5 +73,5 @@ function Waitgroup:wait() end return { - new = new + new = new, } diff --git a/tests/test_channel.lua b/tests/test_channel.lua index bb255a5..1258e6d 100644 --- a/tests/test_channel.lua +++ b/tests/test_channel.lua @@ -62,7 +62,13 @@ local function test_unbounded() chan:get() blocked = false end) + + -- At this point, there is no value available, so get() must have blocked. assert(blocked, "get should block") + + -- Now provide a value so the spawned fibre can complete. + chan:put(123) + print("Unbounded passed") end diff --git a/tests/test_waitgroup.lua b/tests/test_waitgroup.lua index 9bea596..e61c2a8 100644 --- a/tests/test_waitgroup.lua +++ b/tests/test_waitgroup.lua @@ -29,7 +29,7 @@ local function test_simple() for _ = 1, numFibers do wg:add(1) fibers.spawn(function() - sleep.sleep(math.random()) -- Simulate some work + sleep.sleep(0.1) -- Simulate some work wg:done() end) end @@ -45,6 +45,36 @@ local function test_simple() print("Simple test: ok") end +local function test_reuse() + local wg = waitgroup.new() + -- Spawn fibers and add to the waitgroup + wg:add(1) + fibers.spawn(function() + sleep.sleep(0.1) -- Simulate some work + wg:done() + end) + + wg:wait() + + wg:add(1) + local blocked = false + fibers.spawn(function() + sleep.sleep(0.1) -- Simulate some work + wg:done() + end) + + perform( + wg:wait_op() + :or_else(function () blocked = true end) + ) + assert(blocked, "Reused Waitgroup should block.") + + wg:wait() + + print("Reuse test: ok") +end + + local function test_complex() local wg = waitgroup.new() local numFibers = 5 @@ -52,7 +82,7 @@ local function test_complex() local function one_sec_work(w) w:add(1) fibers.spawn(function() - sleep.sleep(1) -- Simulate some work + sleep.sleep(0.1) -- Simulate some work w:done() end) end @@ -76,12 +106,12 @@ local function test_complex() perform( choice( wg:wait_op():wrap(function() done = true end), - sleep.sleep_op(0.9):wrap(extra_work) + sleep.sleep_op(0.09):wrap(extra_work) ) ) end - assert(sc.monotime() - start > 1.5) + assert(sc.monotime() - start > 0.05) print("Complex test: ok") end @@ -89,6 +119,7 @@ end local function main() test_nowait() test_simple() + test_reuse() test_complex() end