Packages

An asynchronous task pool with limited concurrency

Current section

Files

Jump to
crew src crew.gleam
Raw

src/crew.gleam

//// A worker pool for Gleam that distributes work across a limited number of
//// worker processes. Workers are pooled to avoid the overhead of
//// spawning and killing processes for each task.
////
//// The pool manages a queue of work items and distributes them to idle workers.
//// When no workers are available, work is queued until a worker becomes free.
//// The pool handles worker crashes gracefully and automatically manages the
//// worker lifecycle.
////
//// ## Example
////
//// ```gleam
//// import crew
//// import gleam/otp/static_supervisor as supervisor
////
//// pub fn main() {
//// // Create a pool name
//// let pool_name = process.new_name("my_crew")
////
//// // Start an unsupervised pool
//// let assert Ok(_) =
//// crew.new(pool_name)
//// |> crew.fixed_size(4)
//// |> crew.start
////
//// // Execute work on the pool
//// let result = crew.work(pool_name, 5000, fn() {
//// // Some expensive computation
//// expensive_computation()
//// })
//// }
//// ```
// -- IMPORTS -----------------------------------------------------------------
import gleam/bool
import gleam/deque.{type Deque}
import gleam/dict.{type Dict}
import gleam/erlang/process.{
type Monitor, type Name, type Pid, type Selector, type Subject,
}
import gleam/int
import gleam/list
import gleam/otp/actor.{type Next, type StartError}
import gleam/otp/static_supervisor.{type Supervisor}
import gleam/otp/supervision.{type ChildSpecification}
import gleam/result
import gleam/set.{type Set}
import gleam/string
// TODO: - min/max dynamic scaling? what (if any) is the load metric?
// TODO: fairer scheduler? group work by caller and round-robin maybe?
// poolboy: https://github.com/devinus/poolboy
// lifeguard: https://github.com/Pevensie/lifeguard
// gen_stage: https://github.com/elixir-lang/gen_stage
// -- BUILDER -----------------------------------------------------------------
/// A builder for configuring a worker pool before starting it.
pub opaque type Builder {
Builder(name: Name(PoolMsg), size: Int)
}
/// Create a new worker pool builder with the given name.
///
/// The name is used to register the pool so that work can be sent to it.
///
/// By default, the pool will have a number of workers equal to the number of
/// scheduler threads available on the system (typically the number of CPU cores).
///
/// ## Example
///
/// ```gleam
/// let pool_name = process.new_name("pool")
/// let builder = crew.new(pool_name)
/// ```
pub fn new(name: Name(PoolMsg)) -> Builder {
Builder(name:, size: scheduler_count())
}
/// Set the number of worker processes in the pool to a fixed number.
///
/// When set to less than 1 starting the pool will fail.
///
/// ## Example
///
/// ```gleam
/// crew.new(pool_name)
/// |> crew.fixed_size(8) // Use 8 workers regardless of CPU count
/// ```
pub fn fixed_size(builder: Builder, size: Int) -> Builder {
Builder(..builder, size:)
}
// -- START -------------------------------------------------------------------
/// Start an unsupervised worker pool from the given builder.
///
/// Returns a supervisor that manages the pool and its workers. In most cases,
/// you should use `supervised` instead to get a child specification that can
/// be added to your application's supervision tree.
///
/// ## Panics
/// This function will exit the process if any workers fail to start, similar
/// to `static_supervisor.start`.
///
/// ## Example
///
/// ```gleam
/// let assert Ok(pool_supervisor) =
/// crew.new(pool_name)
/// |> crew.start
/// ```
pub fn start(builder: Builder) -> Result(Supervisor, StartError) {
let Builder(name:, size:) = builder
use result <- result.map(start_tree(name, size))
result.data
}
/// Create a child specification for a supervised worker pool.
///
/// This is the recommended way to start a worker pool as part of your
/// application's supervision tree. The returned child specification can be
/// added to a supervisor using `static_supervisor.add`.
///
/// ## Example
///
/// ```gleam
/// let pool_spec =
/// crew.new(pool_name)
/// |> crew.fixed_size(4)
/// |> crew.supervised
///
/// let assert Ok(_) =
/// supervisor.new(supervisor.OneForOne)
/// |> supervisor.add(pool_spec)
/// |> supervisor.start
/// ```
pub fn supervised(builder: Builder) -> ChildSpecification(_) {
let Builder(name:, size:) = builder
use <- supervision.supervisor
start_tree(name, size)
}
fn start_tree(
name: Name(PoolMsg),
size: Int,
) -> Result(actor.Started(_), StartError) {
use <- bool.guard(
when: size <= 0,
return: Error(actor.InitFailed("pool size must be greater than zero")),
)
let main_supervisor = static_supervisor.new(static_supervisor.RestForOne)
let worker_supervisor = static_supervisor.new(static_supervisor.OneForOne)
let pool_spec = {
use <- supervision.worker
actor.new_with_initialiser(1000, init_pool)
|> actor.named(name)
|> actor.on_message(pool)
|> actor.start
}
let worker_spec = {
use <- supervision.worker
let pid = process.spawn(fn() { worker(name) })
Ok(actor.Started(pid, Nil))
}
let worker_supervisor_spec = {
use <- supervision.supervisor
worker_supervisor
|> repeat(times: size, with: static_supervisor.add(_, worker_spec))
|> static_supervisor.start
}
main_supervisor
|> static_supervisor.add(pool_spec)
|> static_supervisor.add(worker_supervisor_spec)
|> static_supervisor.start()
}
// -- WORK --------------------------------------------------------------------
/// Execute a single piece of work on the pool and wait for the result.
///
/// This function blocks until the work is completed or the timeout is reached.
/// The work function is executed on one of the worker processes in the pool.
///
/// ## Parameters
/// - `pool` - The name of the pool to execute work on
/// - `timeout` - Maximum time to wait for completion in milliseconds
/// - `work` - A function containing the work to be executed
///
/// ## Panics
/// - If the pool does not complete the work within the specified timeout
/// - If the pool is not running
/// - If the worker crashes while executing the work
///
/// ## Example
///
/// ```gleam
/// let response = crew.work(pool_name, 5000, fn() {
/// httpc.send(...)
/// })
/// ```
pub fn work(
in pool: Name(PoolMsg),
timeout timeout: Int,
do work: fn() -> any,
) -> any {
let assert [result] = work_many(pool, timeout, [work])
result
}
/// Execute multiple pieces of work concurrently on the pool, without
/// ordering guarantees.
///
/// Work is distributed among the available workers and executed concurrently.
/// Results are returned in the order they complete, not the order they were
/// submitted.
///
/// ## Parameters
/// - `pool` - The name of the pool to execute work on
/// - `timeout` - Maximum time to wait for all work to complete in milliseconds
/// - `work` - A list of functions containing work to be executed
///
/// ## Panics
/// - If the pool does not complete all work within the specified timeout
/// - If the pool is not running
/// - If any worker crashes while executing work
///
/// ## Example
///
/// ```gleam
/// let users = crew.work_many(pool_name, 10000, [
/// fn() { fetch_user_data(user1) },
/// fn() { fetch_user_data(user2) },
/// fn() { fetch_user_data(user3) },
/// ])
/// ```
pub fn work_many(
in pool: Name(PoolMsg),
timeout timeout: Int,
do work: List(fn() -> any),
) -> List(any) {
list.reverse(do_work(pool, timeout, work))
}
/// Execute multiple pieces of work concurrently on the pool and return results
/// in submission order.
///
/// Work functions are distributed across available workers and executed
/// concurrently, but results are reordered to match the original submission
/// order before being returned.
///
/// ## Parameters
/// - `pool` - The name of the pool to execute work on
/// - `timeout` - Maximum time to wait for all work to complete in milliseconds
/// - `work` - A list of functions containing work to be executed
///
/// ## Panics
/// - If the pool does not complete all work within the specified timeout
/// - If the pool is not running
/// - If any worker crashes while executing work
///
/// ## Example
///
/// ```gleam
/// let assert [user1, user2, user3] = crew.work_ordered(pool_name, 10000, [
/// fn() { fetch_user_data(user1) },
/// fn() { fetch_user_data(user2) },
/// fn() { fetch_user_data(user3) },
/// ])
/// ```
pub fn work_ordered(
in pool: Name(PoolMsg),
timeout timeout: Int,
do work: List(fn() -> any),
) -> List(any) {
parallel_map(work, pool, timeout, fn(f) { f() })
}
/// Apply a function to each element of a list concurrently using the worker pool.
///
/// This is similar to `list.map` but executes the mapping function
/// concurrently across all workers. Results are returned in the same
/// order as the input list.
///
/// There's a bunch of extra overhead involved with spawning a work item
/// per list element and making sure the order matches. Depending on your
/// workload it might make sense to split your list into chunks first to reduce
/// work queue pressure.
///
/// ## Parameters
/// - `list` - The list of items to map over
/// - `pool` - The name of the pool to execute work on
/// - `timeout` - Maximum time to wait for all mappings to complete in milliseconds
/// - `fun` - The function to apply to each item
///
/// ## Panics
/// - If the pool does not complete all mappings within the specified timeout
/// - If the pool is not running
/// - If any worker crashes while executing a mapping
///
/// ## Example
///
/// ```gleam
/// let user_ids = [1, 2, 3, 4, 5]
/// let users = crew.parallel_map(user_ids, pool_name, 5000, fetch_user_data)
/// ```
pub fn parallel_map(
over list: List(a),
in pool: Name(PoolMsg),
timeout timeout: Int,
with fun: fn(a) -> b,
) -> List(b) {
// we don't want to do the multiple subject trick here since that is
// potentially O(n^2)
let unordered: List(#(Int, b)) =
list
|> list.index_map(fn(item, index) { fn() { #(index, fun(item)) } })
|> do_work(pool, timeout, _)
unordered
|> list.sort(fn(a, b) { int.compare(a.0, b.0) })
|> list.map(fn(x) { x.1 })
}
fn do_work(
pool: Name(PoolMsg),
timeout: Int,
work: List(fn() -> any),
) -> List(any) {
let channel = subscribe(pool)
enqueue_many(channel, work)
let selector =
process.new_selector()
|> select_map(channel, fn(x) { x })
let timeout_end = system_time() + timeout
let result = receive_loop(work, selector, timeout_end, [])
// sending unsubscribe unconditionally in an attempt to handle
// cases more gracefully where this crash is rescued by the caller.
unsubscribe(channel, selector)
case result {
Ok(value) -> value
Error(_) -> panic as "Pool did not complete work in the alloted time"
}
}
fn receive_loop(
work: List(fn() -> any),
selector: Selector(any),
timeout_end: Int,
state: List(any),
) -> Result(List(any), Nil) {
case work {
[] -> Ok(state)
[_, ..work] -> {
let timeout = timeout_end - system_time()
case process.selector_receive(selector, timeout) {
Ok(result) ->
receive_loop(work, selector, timeout_end, [result, ..state])
Error(Nil) -> Error(Nil)
}
}
}
}
// -- SUBSCRIBE / UNSUBSCRIBE -------------------------------------------------
/// A channel represents a typed subscription to a worker pool for receiving
/// work results. This allows you to send work and receive results asynchronously.
pub opaque type Channel(a) {
Channel(
pool: Subject(PoolMsg),
receive: Subject(WorkResult),
monitor: Monitor,
)
}
/// Create a typed subscription channel to the worker pool for work submission.
///
/// This function provides a lower-level interface for submitting work to the
/// pool. In many cases calling the `work*` function from a separate process
/// will be easier.
///
/// Subscribing to the pool creates a reference to the current process in the pool
/// that must be cleaned up using `unsubscribe`. The returned channel can be used
/// with `select_map` and `enqueue` for custom message handling.
///
/// ## Panics
/// - If the pool is not running
///
/// ## Example
///
/// ```gleam
/// let channel = crew.subscribe(pool_name)
/// crew.enqueue(channel, fn() { some_work() })
/// // Handle results with selector...
/// crew.unsubscribe(channel, selector)
/// ```
pub fn subscribe(pool: Name(PoolMsg)) -> Channel(a) {
let pool_subject = process.named_subject(pool)
let assert Ok(pool_pid) = process.subject_owner(pool_subject)
as "Pool is not running"
let monitor = process.monitor(pool_pid)
let receive: Subject(WorkResult) = process.new_subject()
actor.send(pool_subject, Subscribe(receive))
Channel(pool: pool_subject, receive:, monitor:)
}
/// Add a worker pool channel to a selector for receiving work results.
///
/// This allows you to receive work results as part of the larger message-handling
/// loop. Used in conjunction with `subscribe` and `enqueue` for lower-level
/// pool usage.
///
/// Work results arrive in completion order. It is your responsibility to handle
/// the lifecycle and ordering.
///
/// This adds a specific monitor to the selector to handle the case where the pool
/// goes down, so it's better to prefer few channels over many.
///
/// ## Panics
/// - If a worker or the pool crashes while executing work, the selector will
/// panic with details about the crash.
pub fn select_map(
selector: Selector(msg),
channel: Channel(a),
tagger: fn(a) -> msg,
) -> Selector(msg) {
selector
|> process.select_specific_monitor(channel.monitor, fn(down) {
let msg =
"Pool exited while waiting for work to complete: " <> string.inspect(down)
panic as msg
})
|> process.select_map(channel.receive, fn(result) {
case result {
Done(value) -> tagger(cast(value))
WorkerExited(reason) -> {
let msg = "Worker exited: " <> string.inspect(reason)
panic as msg
}
}
})
}
/// Remove a subscription channel from the pool and clean up resources.
///
/// This function should be called when you're done using a channel obtained
/// from `subscribe`. It stops any in-progress work for this channel and
/// removes the channel's handlers from the selector.
///
/// Any finished work that is still coming in is dropped.
pub fn unsubscribe(
channel: Channel(a),
selector: Selector(msg),
) -> Selector(msg) {
let unsubscribed = process.new_subject()
actor.send(channel.pool, Unsubscribe(channel.receive, unsubscribed:))
wait_for_unsubscribe(channel, unsubscribed)
process.demonitor_process(channel.monitor)
selector
|> process.deselect_specific_monitor(channel.monitor)
|> process.deselect(channel.receive)
}
fn wait_for_unsubscribe(channel: Channel(a), unsubscribed: Subject(b)) -> Nil {
let unsubscribe_selector =
process.new_selector()
|> process.select_map(unsubscribed, fn(_) { True })
|> process.select_specific_monitor(channel.monitor, fn(_) { True })
|> process.select_map(channel.receive, fn(_) { False })
drop_messages_loop(unsubscribe_selector)
}
fn drop_messages_loop(selector: Selector(Bool)) -> Nil {
case process.selector_receive_forever(selector) {
True -> Nil
False -> drop_messages_loop(selector)
}
}
/// Submit a single piece of work to the pool using a subscription channel.
///
/// This is a lower-level function for submitting work. Results must be handled
/// using a selector with `select_map`. It is the callers responsibility to handle
/// timeouts and submission order. Most users should prefer the `work*` functions.
pub fn enqueue(channel: Channel(a), work: fn() -> a) -> Nil {
enqueue_many(channel, [work])
}
/// Submit multiple pieces of work to the pool using a subscription channel.
///
/// This is a lower-level function for submitting work. Results must be handled
/// using a selector with `select_map`. It is the callers responsibility to handle
/// timeouts and submission order. Most users should prefer the `work*` functions.
pub fn enqueue_many(channel: Channel(a), work: List(fn() -> a)) -> Nil {
actor.send(channel.pool, Enqueue(channel.receive, cast(work)))
}
// -- POOL --------------------------------------------------------------------
pub opaque type PoolMsg {
WorkerStarted(pid: Pid, send: Subject(Work))
WorkerIdle(pid: Pid)
MonitoredProcessExited(reason: process.Down)
//
Subscribe(receive: Subject(WorkResult))
Unsubscribe(receive: Subject(WorkResult), unsubscribed: Subject(Nil))
Enqueue(receive: Subject(WorkResult), work: List(fn() -> Any))
}
type State {
State(
idle_workers: List(Worker),
active_workers: Dict(Pid, ActiveWorker),
callers: Dict(Pid, Caller),
requests: Dict(Subject(WorkResult), Request),
queue: Deque(QueueItem),
)
}
type Worker {
Worker(pid: Pid, send: Subject(Work), monitor: Monitor)
}
type QueueItem {
QueueItem(work: fn() -> Any, caller: Pid, receive: Subject(WorkResult))
}
type ActiveWorker {
ActiveWorker(worker: Worker, caller: Pid, receive: Subject(WorkResult))
}
type Caller {
Caller(pid: Pid, monitor: Monitor, requests: Set(Subject(WorkResult)))
}
type Request {
Request(from: Pid, workers: Set(Pid), receive: Subject(WorkResult))
}
fn init_pool(self: Subject(PoolMsg)) {
let selector =
process.new_selector()
|> process.select(self)
|> process.select_monitors(MonitoredProcessExited)
let state =
State(
idle_workers: [],
active_workers: dict.new(),
callers: dict.new(),
requests: dict.new(),
queue: deque.new(),
)
actor.initialised(state)
|> actor.selecting(selector)
|> actor.returning(self)
|> Ok
}
fn pool(state: State, msg: PoolMsg) -> Next(State, PoolMsg) {
let next_state = case msg {
WorkerStarted(pid:, send:) -> {
let monitor = process.monitor(pid)
let worker = Worker(pid:, send:, monitor:)
try_dequeue_work(state, worker)
}
MonitoredProcessExited(process.ProcessDown(pid:, monitor: _, reason:)) -> {
case dict.get(state.callers, pid), dict.get(state.active_workers, pid) {
// Caller down -
// We will get a second Down message for when the process is caller and worker.
Ok(_), _ -> {
// this sends a message back to the caller, which we already
// know is down. This should not matter.
abort_caller(state, pid, reason)
}
// Active worker down
Error(_), Ok(ActiveWorker(caller:, ..)) -> {
// after one worker exits, we want to stop all workers for this caller
// and remove it as an active caller. The caller will no longer receive messages.
let active_workers = dict.delete(state.active_workers, pid)
State(..state, active_workers:)
|> abort_caller(caller, reason)
}
// idle worker down
Error(_), Error(_) -> {
let idle_workers =
list.filter(state.idle_workers, fn(worker) { worker.pid != pid })
State(..state, idle_workers:)
}
}
}
// we only monitor processes and can ignore all other Down messages
MonitoredProcessExited(_) -> state
WorkerIdle(pid:) -> {
// got an idle message from a worker that is not active
use ActiveWorker(worker:, receive:, caller: _) <- try_(
dict.get(state.active_workers, pid),
state,
)
let active_workers = dict.delete(state.active_workers, pid)
case dict.get(state.requests, receive) {
Ok(request) -> {
let request =
Request(..request, workers: set.delete(request.workers, pid))
let requests = dict.insert(state.requests, receive, request)
State(..state, requests:, active_workers:)
|> try_dequeue_work(worker)
}
// the requests unsubscribed while we were still working.
// we sent a kill request while unsubscribing, so we drop the worker here.
Error(_) -> State(..state, active_workers:)
}
}
Enqueue(work: [], ..) -> state
Enqueue(receive:, work:) -> {
use request <- try_(dict.get(state.requests, receive), state)
enqueue_loop(state, request, work)
}
Subscribe(receive:) -> {
use pid <- try_(process.subject_owner(receive), state)
let request =
dict.get(state.requests, receive)
|> result.unwrap(Request(from: pid, workers: set.new(), receive:))
let caller = case dict.get(state.callers, pid) {
Ok(caller) ->
Caller(..caller, requests: set.insert(caller.requests, receive))
Error(_) -> {
let monitor = process.monitor(pid)
Caller(pid: pid, monitor:, requests: set.new() |> set.insert(receive))
}
}
let callers = dict.insert(state.callers, pid, caller)
let requests = dict.insert(state.requests, receive, request)
State(..state, callers:, requests:)
}
Unsubscribe(receive:, unsubscribed:) -> {
use request <- try_(dict.get(state.requests, receive), state)
let requests = dict.delete(state.requests, receive)
set.each(request.workers, process.kill)
use caller <- try_(
dict.get(state.callers, request.from),
State(..state, requests:),
)
let caller =
Caller(..caller, requests: set.delete(caller.requests, receive))
let callers = case set.is_empty(caller.requests) {
True -> {
process.demonitor_process(caller.monitor)
dict.delete(state.callers, caller.pid)
}
False -> {
dict.insert(state.callers, caller.pid, caller)
}
}
process.send(unsubscribed, Nil)
State(..state, requests:, callers:)
}
}
actor.continue(next_state)
}
fn enqueue_loop(
state: State,
request: Request,
work: List(fn() -> Any),
) -> State {
let Request(from: caller, receive:, workers:) = request
case work, state.idle_workers {
[work, ..rest], [worker, ..idle_workers] -> {
// got a work item and a worker - it's a match!
process.send(worker.send, Work(work:, receive:))
let active_worker = ActiveWorker(worker:, caller:, receive:)
let active_workers =
dict.insert(state.active_workers, worker.pid, active_worker)
let request = Request(..request, workers: set.insert(workers, worker.pid))
State(..state, active_workers:, idle_workers:)
|> enqueue_loop(request, rest)
}
_, _ -> {
let queue =
list.fold(work, state.queue, fn(queue, work) {
deque.push_back(queue, QueueItem(work:, caller:, receive:))
})
let requests = dict.insert(state.requests, receive, request)
State(..state, queue:, requests:)
}
}
}
fn try_dequeue_work(state: State, worker: Worker) -> State {
use #(QueueItem(work:, caller:, receive:), queue) <- try(
deque.pop_front(state.queue),
// the queue is empty, this worker becomes idle.
fn(_) { State(..state, idle_workers: [worker, ..state.idle_workers]) },
)
use request <- try(dict.get(state.requests, receive), fn(_) {
// this request unsubscribed
try_dequeue_work(State(..state, queue:), worker)
})
process.send(worker.send, Work(work:, receive:))
let active_worker = ActiveWorker(worker:, caller:, receive:)
let active_workers =
dict.insert(state.active_workers, worker.pid, active_worker)
let request =
Request(..request, workers: set.insert(request.workers, worker.pid))
let requests = dict.insert(state.requests, receive, request)
State(..state, active_workers:, requests:, queue:)
}
fn abort_caller(state: State, caller: Pid, reason: process.ExitReason) -> State {
// we do not have an active caller, so this worker was stopped in
// response to a caller exiting.
use caller <- try_(dict.get(state.callers, caller), state)
process.demonitor_process(caller.monitor)
let requests =
set.fold(caller.requests, state.requests, fn(requests, receive) {
use request <- try_(dict.get(state.requests, receive), requests)
// we kill all workers currently working on things for this caller.
// the WorkerStopped messages will then clean up the workers.
set.each(request.workers, process.kill)
// notify the caller that we crashed. this is the last message we send.
process.send(receive, WorkerExited(reason))
dict.delete(requests, receive)
})
// the queue is cleaned up while consuming it if the caller is missing.
let callers = dict.delete(state.callers, caller.pid)
State(..state, callers:, requests:)
}
// -- WORKER ------------------------------------------------------------------
type Any
type Work {
Work(work: fn() -> Any, receive: Subject(WorkResult))
}
type WorkResult {
Done(Any)
WorkerExited(process.ExitReason)
}
fn worker(pool: Name(PoolMsg)) -> Nil {
let self = process.self()
let subject = process.new_subject()
let pool_subject = process.named_subject(pool)
process.send(pool_subject, WorkerStarted(self, subject))
worker_loop(pool_subject, self, subject)
}
fn worker_loop(pool: Subject(PoolMsg), self: Pid, subject: Subject(Work)) {
let Work(work:, receive:) = process.receive_forever(subject)
let result = work()
process.send(pool, WorkerIdle(self))
process.send(receive, Done(result))
worker_loop(pool, self, subject)
}
// -- HELPERS -----------------------------------------------------------------
fn repeat(
times times: Int,
from state: state,
with f: fn(state) -> state,
) -> state {
case times > 0 {
True -> repeat(times - 1, f(state), f)
False -> state
}
}
fn try(result: Result(a, x), or: fn(x) -> b, then: fn(a) -> b) {
case result {
Ok(x) -> then(x)
Error(x) -> or(x)
}
}
fn try_(result: Result(a, _), or: b, then: fn(a) -> b) {
case result {
Ok(x) -> then(x)
Error(_) -> or
}
}
@external(erlang, "crew_ffi", "scheduler_count")
fn scheduler_count() -> Int
@external(erlang, "crew_ffi", "system_time")
fn system_time() -> Int
@external(erlang, "crew_ffi", "identity")
fn cast(value: a) -> b