//// 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/erlang/process //// //// pub fn main() { //// // Create a pool name //// let pool_name = process.new_name("image_downloader") //// //// // Start an unsupervised pool //// let assert Ok(_) = //// crew.new(pool_name, download_image) //// |> crew.fixed_size(4) //// |> crew.start //// //// // Execute work on the pool //// let result = crew.run(pool_name, 5000, "tasteful-ramen.jpeg") //// } //// ``` // -- 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/list import gleam/option.{type Option, None, Some} 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(state, work, result) { Builder( name: Name(PoolMsg(work, result)), size: Int, max_queue_length: Option(Int), init: fn() -> state, init_timeout: Int, work: fn(state, work) -> result, ) } /// 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, worker) /// ``` pub fn new( name: Name(PoolMsg(work, result)), work: fn(work) -> result, ) -> Builder(Nil, work, result) { new_with_state(name, Nil, fn(_state, input) { work(input) }) } /// Create a new worker pool builder with the given name and state. /// /// 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). pub fn new_with_state( name: Name(PoolMsg(work, result)), state: state, work: fn(state, work) -> result, ) -> Builder(state, work, result) { new_with_initialiser(name, 1000, fn() { state }, work) } /// Create a new worker pool builder with the given name and initialiser. /// /// The name is used to register the pool so that work can be sent to it. /// The initialiser will run on the worker process before it registers itself /// as a worker. /// /// 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). pub fn new_with_initialiser( name: Name(PoolMsg(work, result)), timeout init_timeout: Int, init init: fn() -> state, run work: fn(state, work) -> result, ) -> Builder(state, work, result) { let size = scheduler_count() Builder(name:, size:, max_queue_length: None, init_timeout:, init:, work:) } /// 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(state, work, result), size: Int, ) -> Builder(state, work, result) { Builder(..builder, size:) } /// To avoid overloading the system, the internal work queue is limited. /// /// If this limit is reached, callers of `enqueue` have to wait until enough /// work has been done before continuing. /// /// By default, the `max_queue_size` depends on the number of workers in the pool. /// /// ## Example /// /// ```gleam /// crew.new(pool_name, worder) /// |> crew.max_queue_length(100) /// ``` pub fn max_queue_length( builder: Builder(state, work, result), max_queue_length: Int, ) -> Builder(state, work, result) { Builder(..builder, max_queue_length: Some(max_queue_length)) } // -- 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(state, work, result), ) -> Result(Supervisor, StartError) { use result <- result.map(start_tree(builder)) 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(state, work, result), ) -> ChildSpecification(Supervisor) { use <- supervision.supervisor start_tree(builder) } fn start_tree( builder: Builder(state, work, result), ) -> Result(actor.Started(_), StartError) { use <- bool.guard( when: builder.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(builder, _)) |> actor.named(builder.name) |> actor.on_message(pool) |> actor.start } let worker_spec = { use <- supervision.worker start_worker(builder) } let worker_supervisor_spec = { use <- supervision.supervisor worker_supervisor |> repeat(times: builder.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() } // -- CALL -------------------------------------------------------------------- /// Send a single piece of work to one of the workers and wait for the result. /// /// This function blocks until the work is completed or the timeout is reached. /// If no worker is available, the message will be queued and sent to the next /// free worker. /// /// ## Parameters /// - `pool` - The name of the pool to execute work on /// - `timeout` - Maximum time to wait for completion in milliseconds /// - `work` - The message to send to the worker function /// /// ## 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 pub fn call( in pool: Name(PoolMsg(work, result)), timeout timeout: Int, msg work: work, ) -> result { let assert [result] = do_call(pool, timeout, [work]) result } /// Send multiple pieces of work concurrently to the pool workers, 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 messages describing the work to be done /// /// ## 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 pub fn call_parallel( in pool: Name(PoolMsg(work, result)), timeout timeout: Int, msg work: List(work), ) -> List(result) { list.reverse(do_call(pool, timeout, work)) } @internal pub fn do_call( pool: Name(PoolMsg(work, result)), timeout: Int, work: List(work), ) -> List(result) { let timeout_end = system_time() + timeout let assert Ok(pool_pid) = process.named(pool) as "Pool is not running" let monitor = process.monitor(pool_pid) let receive = process.new_subject() // we can always cast here because we will be blocking below anyways actor.send(process.named_subject(pool), Enqueue(receive, work, None)) let selector = process.new_selector() |> process.select_specific_monitor(monitor, fn(down) { let msg = "Pool exited while waiting for work to complete: " <> string.inspect(down) panic as msg }) |> process.select_map(receive, fn(result) { case result { Ok(value) -> value Error(down) -> { // demonitor here in case the panic is rescued process.demonitor_process(monitor) let msg = "Worker exited: " <> string.inspect(down) panic as msg } } }) let result = receive_loop(work, selector, timeout_end, []) process.demonitor_process(monitor) cancel(pool, receive) case result { Ok(value) -> value Error(_) -> panic as "Pool did not complete work in time" } } fn receive_loop( work: List(work), selector: Selector(result), timeout_end: Int, state: List(result), ) -> Result(List(result), 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) } } } } // -- ASYNC API --------------------------------------------------------------- /// The first time `enqueue*` is called from a process, the pool starts to monitor /// that process and cancels all ongoing work in case it goes down. /// /// Sometimes it is useful to manually unsubscribe and cancel all ongoing work /// for a subject. Doing so will also remove the monitor added in the pool once /// the last `receive` subject got cancelled. /// /// Note that finished work might still arrive on this selector after /// `cancel` got called. pub fn cancel( pool: Name(PoolMsg(work, result)), receive: Subject(Result(result, process.ExitReason)), ) -> Nil { actor.send(process.named_subject(pool), Cancel(receive)) } /// Submit a single piece of work to the pool using a subscription channel. /// The work will be done asynchronously and the result will be sent back /// to the provided subject. The timeout controls how long to wait for the /// work to be added to the queue successfully. /// /// This is a lower-level function for submitting work. It is the callers /// responsibility to handle timeouts, submission order and failures. /// Most users should prefer `call`. /// /// The first time you `enqueue` is called from a process, the pool sets up /// a monitor making sure work is cancelled when the process no longer exists /// to receive a result. You can clean up this monitor early by using `cancel`. pub fn submit( pool: Name(PoolMsg(work, result)), timeout: Int, receive: Subject(Result(result, process.ExitReason)), work: work, ) -> Nil { submit_all(pool, timeout, receive, [work]) } /// Submit multiple pieces of work to the pool using a subscription channel. /// The work will be done asynchronously and the result will be sent back /// to the provided subject. The timeout controls how long to wait for the /// work to be added to the queue successfully. /// /// This is a lower-level function for submitting work. It is the callers /// responsibility to handle timeouts, submission order and failures. /// Most users should prefer the `call_parallel` function. /// /// The first time you `enqueue` is called from a process, the pool sets up /// a monitor making sure work is cancelled when the process no longer exists /// to receive a result. You can clean up this monitor early by using `cancel`. pub fn submit_all( pool: Name(PoolMsg(work, result)), timeout: Int, receive: Subject(Result(result, process.ExitReason)), work: List(work), ) -> Nil { let counter = get_counter(pool) let subject = process.named_subject(pool) case work { [] -> Nil [_, ..] if counter > 0 -> { actor.send(subject, Enqueue(receive:, work:, enqueued: None)) } [_, ..] -> { actor.call(subject, timeout, fn(enqueued) { Enqueue(receive:, work:, enqueued: Some(enqueued)) }) } } } // -- POOL -------------------------------------------------------------------- pub opaque type PoolMsg(work, result) { WorkerStarted(pid: Pid, send: Subject(Work(work, result))) WorkerIdle(pid: Pid) MonitoredProcessExited(reason: process.Down) // GetWorkerCount(reply_to: Subject(Int)) Enqueue( receive: Receiver(result), work: List(work), enqueued: Option(Subject(Nil)), ) Cancel(receive: Receiver(result)) } type State(work, result) { State( name: Name(PoolMsg(work, result)), worker_count: Int, idle_workers: List(Worker(work, result)), active_workers: Dict(Pid, ActiveWorker(work, result)), // callers: Dict(Pid, Caller(result)), channels: Dict(Receiver(result), Channel(work, result)), // queue: Deque(Work(work, result)), overflow_queue: Deque(Request(work, result)), ) } type Worker(work, result) { Worker(pid: Pid, send: Subject(Work(work, result)), monitor: Monitor) } type ActiveWorker(work, result) { ActiveWorker(worker: Worker(work, result), work: Work(work, result)) } type Caller(result) { Caller(pid: Pid, monitor: Monitor, channels: Set(Receiver(result))) } type Channel(work, result) { Channel(from: Pid, workers: Set(Pid), receive: Receiver(result)) } type Request(work, result) { Request( caller: Pid, receive: Receiver(result), work: List(work), enqueued: Option(Subject(Nil)), ) } fn init_pool( builder: Builder(state, work, result), self: Subject(PoolMsg(work, result)), ) { let selector = process.new_selector() |> process.select(self) |> process.select_monitors(MonitoredProcessExited) // it's big, but it will provide backpressure eventually. let max_queue_length = builder.max_queue_length |> option.unwrap(builder.size * 100) set_counter(builder.name, max_queue_length) let state = State( name: builder.name, worker_count: 0, idle_workers: [], active_workers: dict.new(), callers: dict.new(), channels: dict.new(), queue: deque.new(), overflow_queue: deque.new(), ) actor.initialised(state) |> actor.selecting(selector) |> actor.returning(self) |> Ok } fn pool( state: State(work, result), msg: PoolMsg(work, result), ) -> Next(State(work, result), PoolMsg(work, result)) { let next_state = case msg { WorkerStarted(pid:, send:) -> { let monitor = process.monitor(pid) let worker = Worker(pid:, send:, monitor:) State(..state, worker_count: state.worker_count + 1) |> try_dequeue_work(worker) } MonitoredProcessExited(process.ProcessDown(pid:, ..)) -> { 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(_), _ -> { // we do not have an active caller, maybe we still got the cancel just before? use caller <- try_(dict.get(state.callers, pid), state) let channels = set.fold(caller.channels, state.channels, fn(channels, receive) { use channel <- try_(dict.get(state.channels, receive), channels) // we kill all workers currently working on things for this caller. // the WorkerStopped messages will then clean up the workers. set.each(channel.workers, process.kill) dict.delete(channels, receive) }) // the queue is cleaned up while consuming it if the caller is missing. let callers = dict.delete(state.callers, pid) State(..state, callers:, channels:) } // Active worker down Error(_), Ok(_) -> { let active_workers = dict.delete(state.active_workers, pid) let worker_count = state.worker_count - 1 State(..state, active_workers:, worker_count:) } // idle worker down Error(_), Error(_) -> { let idle_workers = list.filter(state.idle_workers, fn(worker) { worker.pid != pid }) let worker_count = state.worker_count - 1 State(..state, idle_workers:, worker_count:) } } } // 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:, work: Work(receive:, ..)) <- try_( dict.get(state.active_workers, pid), state, ) let active_workers = dict.delete(state.active_workers, pid) case dict.get(state.channels, receive) { Ok(channel) -> { let channel = Channel(..channel, workers: set.delete(channel.workers, pid)) let channels = dict.insert(state.channels, receive, channel) State(..state, channels:, active_workers:) |> try_dequeue_work(worker) } // the channels got cancelled while we were still working. // we sent a kill channel while unsubscribing, so we drop the worker here. Error(_) -> State(..state, active_workers:, worker_count: state.worker_count - 1) } } GetWorkerCount(reply_to:) -> { process.send(reply_to, state.worker_count) state } Enqueue(receive: _, work: [], enqueued: None) -> state Enqueue(receive: _, work: [], enqueued: Some(enqueued)) -> { process.send(enqueued, Nil) state } Enqueue(receive:, work:, enqueued:) -> { // caller already exited, do not queue their work. use pid <- try_(process.subject_owner(receive), state) // get the channel data or construct a new entry if needed. let channel = case dict.get(state.channels, receive) { Ok(channel) -> channel Error(_) -> Channel(from: pid, workers: set.new(), receive:) } // get the caller data and insert the new channel, starting to monitor if needed. let caller = case dict.get(state.callers, pid) { Ok(caller) -> Caller(..caller, channels: set.insert(caller.channels, receive)) Error(_) -> { let monitor = process.monitor(pid) Caller(pid: pid, monitor:, channels: set.new() |> set.insert(receive)) } } // enqueue_loop will insert the channel into the state. State(..state, callers: dict.insert(state.callers, pid, caller)) |> enqueue_loop(channel, work, enqueued) } Cancel(receive:) -> { // if the receiver does not exist, we don't need to do anything use channel <- try_(dict.get(state.channels, receive), state) let channels = dict.delete(state.channels, receive) // cancel all running workers set.each(channel.workers, process.kill) use caller <- try_( dict.get(state.callers, channel.from), State(..state, channels:), ) let caller = Caller(..caller, channels: set.delete(caller.channels, receive)) // if this channel was the last one for the caller, demonitor and remove it. let callers = case set.is_empty(caller.channels) { True -> { process.demonitor_process(caller.monitor) dict.delete(state.callers, caller.pid) } False -> { dict.insert(state.callers, caller.pid, caller) } } State(..state, channels:, callers:) } } actor.continue(next_state) } fn enqueue_loop( state: State(work, result), channel: Channel(work, result), work: List(work), enqueued: Option(Subject(Nil)), ) -> State(work, result) { let Channel(from: caller, receive:, workers:) = channel case work, state.idle_workers { [work, ..rest], [worker, ..idle_workers] -> { // got a work item and a worker - it's a match! let work = Work(work:, caller:, receive:) process.send(worker.send, work) let active_worker = ActiveWorker(worker:, work:) let active_workers = dict.insert(state.active_workers, worker.pid, active_worker) let channel = Channel(..channel, workers: set.insert(workers, worker.pid)) State(..state, active_workers:, idle_workers:) |> enqueue_loop(channel, rest, enqueued) } [_, ..], [] -> { let capacity = get_counter(state.name) let #(queue, overflow, new_capacity) = enqueue_loop2(caller, receive, work, state.queue, capacity) decrement_counter(state.name, capacity - new_capacity) // if we have overflow, push it onto the channels overflow queue // if we don't, send back that we enqueued successfully. let overflow_queue = case overflow { [] -> { case enqueued { Some(enqueued) -> process.send(enqueued, Nil) None -> Nil } state.overflow_queue } work -> { let request = Request(caller:, receive:, work:, enqueued:) deque.push_back(state.overflow_queue, request) } } let channels = dict.insert(state.channels, receive, channel) State(..state, queue:, channels:, overflow_queue:) } [], _ -> { case enqueued { Some(enqueued) -> process.send(enqueued, Nil) None -> Nil } State(..state, channels: dict.insert(state.channels, receive, channel)) } } } fn enqueue_loop2( caller: Pid, receive: Receiver(result), work: List(work), queue: Deque(Work(work, result)), capacity: Int, ) { case work { [work, ..rest] if capacity > 0 -> { let queue = deque.push_back(queue, Work(work:, caller:, receive:)) let capacity = capacity - 1 enqueue_loop2(caller, receive, rest, queue, capacity) } _ -> #(queue, work, capacity) } } fn try_dequeue_work( state: State(work, result), worker: Worker(work, result), ) -> State(work, result) { use #(work, state) <- try(pop_work(state), fn(_) { // the queue is empty, this worker becomes idle. State(..state, idle_workers: [worker, ..state.idle_workers]) }) use channel <- try(dict.get(state.channels, work.receive), fn(_) { // this channel got cancelled, try again try_dequeue_work(state, worker) }) process.send(worker.send, work) let active_worker = ActiveWorker(worker:, work:) let active_workers = dict.insert(state.active_workers, worker.pid, active_worker) let channel = Channel(..channel, workers: set.insert(channel.workers, worker.pid)) let channels = dict.insert(state.channels, work.receive, channel) State(..state, active_workers:, channels:) } fn pop_work( state: State(work, result), ) -> Result(#(Work(work, result), State(work, result)), Nil) { use #(work, queue) <- result.try(deque.pop_front(state.queue)) use #(Request(caller:, receive:, work: request, enqueued:), overflow_queue) <- try( deque.pop_front(state.overflow_queue), fn(_) { // overflow queue is empty, the queue got smaller and we can increase the capacity increment_counter(state.name, 1) Ok(#(work, State(..state, queue:))) }, ) case request { [first, ..rest] -> { // we can push one overflow item onto the queue as a replacement for the // item that just got popped off. let queue = deque.push_back(queue, Work(work: first, caller:, receive:)) // push the request back if there is still work to do. // if there is none, send the enqueued signal. let overflow_queue = case rest { [_, ..] -> { let request = Request(caller:, receive:, work: rest, enqueued:) deque.push_front(overflow_queue, request) } [] -> { case enqueued { Some(enqueued) -> process.send(enqueued, Nil) None -> Nil } overflow_queue } } Ok(#(work, State(..state, queue:, overflow_queue:))) } [] -> { // got an empty work request - this should not happen but I guess we can handle it case enqueued { Some(enqueued) -> process.send(enqueued, Nil) None -> Nil } pop_work(State(..state, overflow_queue:)) } } } // -- WORKER ------------------------------------------------------------------ type Receiver(result) = Subject(Result(result, process.ExitReason)) type Work(work, result) { Work(work: work, caller: Pid, receive: Receiver(result)) } fn start_worker( builder: Builder(state, work, result), ) -> Result(actor.Started(_), actor.StartError) { let initialised = process.new_subject() let pid = process.spawn(fn() { worker(builder, initialised) }) let monitor = process.monitor(pid) let selector = process.new_selector() |> process.select_map(initialised, Ok) |> process.select_specific_monitor(monitor, Error) let result = case process.selector_receive(selector, builder.init_timeout) { Ok(Ok(Nil)) -> Ok(actor.Started(pid, Nil)) Ok(Error(down)) -> Error(actor.InitExited(down.reason)) Error(Nil) -> { process.unlink(pid) process.kill(pid) Error(actor.InitTimeout) } } process.demonitor_process(monitor) result } fn worker(builder: Builder(state, work, result), initialised: Subject(Nil)) { let self = process.self() let subject = process.new_subject() let pool_subject = process.named_subject(builder.name) let state = builder.init() process.send(initialised, Nil) process.send(pool_subject, WorkerStarted(self, subject)) worker_loop(pool_subject, self, subject, state, builder.work) } fn worker_loop(pool, self, subject, state, do_work) -> Nil { let Work(work:, receive:, caller: _) = process.receive_forever(subject) let result = do_work(state, work) process.send(pool, WorkerIdle(self)) process.send(receive, Ok(result)) worker_loop(pool, self, subject, state, do_work) } // -- 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", "get_counter") fn get_counter(name: Name(a)) -> Int @external(erlang, "crew_ffi", "increment_counter") fn increment_counter(name: Name(a), amount: Int) -> Nil @external(erlang, "crew_ffi", "decrement_counter") fn decrement_counter(name: Name(a), amount: Int) -> Nil @external(erlang, "crew_ffi", "set_counter") fn set_counter(name: Name(a), value: Int) -> Nil