khora_control/
worker_pool.rs

1// Copyright 2025 eraflo
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! A small persistent thread pool owned by the [`ExecutionScheduler`].
16//!
17//! The scheduler runs a concurrent wave's [`AgentAccess::Isolated`] agents as
18//! independent jobs. Rather than re-spawn threads with `std::thread::scope`
19//! every wave, this pool spawns its workers **once** (in
20//! [`ExecutionScheduler::new`]) and reuses them for the life of the scheduler,
21//! amortising the spawn cost as waves grow.
22//!
23//! The jobs it runs are `'static`: an `Isolated` agent touches no `World` and
24//! reaches its per-frame inputs through an `Arc<LaneBus>` + `Arc<Runtime>` and
25//! its own `Arc<Mutex<dyn Agent>>`, all of which are owned/`Send`. Nothing
26//! borrowed from the frame stack crosses onto a worker — the scheduler's lone
27//! `SharedWorld` agent (which reads `&World`) runs inline on the calling thread
28//! instead — so this needs **no `unsafe`** and no lifetime transmute.
29//!
30//! # RULES §5
31//! This is the third named exception to the "no `std::thread::spawn`" rule (the
32//! DCC tick thread and the scheduler's scoped executor are the other two): a
33//! scheduler-owned pool, spawned once and joined on [`Drop`]. It never spawns
34//! per frame and its threads never outlive the scheduler.
35//!
36//! [`AgentAccess::Isolated`]: khora_core::agent::AgentAccess::Isolated
37//! [`ExecutionScheduler`]: crate::scheduler::ExecutionScheduler
38//! [`ExecutionScheduler::new`]: crate::scheduler::ExecutionScheduler::new
39
40use std::sync::mpsc::{self, Receiver, Sender};
41use std::sync::{Arc, Mutex};
42use std::thread::JoinHandle;
43
44/// A unit of work run on a worker thread.
45type Job = Box<dyn FnOnce() + Send + 'static>;
46
47/// Message sent from the scheduler to a worker thread.
48enum Message {
49    /// Run this job, then wait for the next message.
50    Run(Job),
51    /// Stop the worker loop and let the thread exit.
52    Shutdown,
53}
54
55/// A fixed-size pool of persistent worker threads.
56///
57/// Submit `'static` jobs with [`submit`](Self::submit); they run on the next
58/// free worker. The pool joins all workers on [`Drop`], so a scheduler simply
59/// holds it as a field and lets RAII shut it down.
60pub struct WorkerPool {
61    /// Sender half of the job channel. `None` only transiently during `Drop`.
62    sender: Option<Sender<Message>>,
63    /// Worker thread handles, joined on `Drop`.
64    workers: Vec<JoinHandle<()>>,
65}
66
67impl WorkerPool {
68    /// Spawns a pool of `size` worker threads (at least one).
69    pub fn new(size: usize) -> Self {
70        let size = size.max(1);
71        let (sender, receiver) = mpsc::channel::<Message>();
72        // Shared receiver: each worker locks it only long enough to pull the
73        // next message, then releases it before running the job — the canonical
74        // safe std work-stealing-lite pattern.
75        let receiver = Arc::new(Mutex::new(receiver));
76
77        let mut workers = Vec::with_capacity(size);
78        for _ in 0..size {
79            let receiver: Arc<Mutex<Receiver<Message>>> = Arc::clone(&receiver);
80            // RULES §5 (3rd exception): scheduler-owned persistent pool, spawned
81            // once here and joined on Drop — never spawned per frame.
82            let handle = std::thread::spawn(move || loop {
83                let message = {
84                    let rx = receiver.lock().unwrap_or_else(|e| e.into_inner());
85                    rx.recv()
86                };
87                match message {
88                    Ok(Message::Run(job)) => job(),
89                    // Explicit shutdown, or the sender was dropped: exit.
90                    Ok(Message::Shutdown) | Err(_) => break,
91                }
92            });
93            workers.push(handle);
94        }
95
96        Self {
97            sender: Some(sender),
98            workers,
99        }
100    }
101
102    /// Number of worker threads in the pool.
103    pub fn thread_count(&self) -> usize {
104        self.workers.len()
105    }
106
107    /// Submits a job to run on the next free worker.
108    ///
109    /// Best-effort: if the pool is mid-shutdown (all workers gone) the job is
110    /// dropped rather than run. During normal operation the send always
111    /// succeeds.
112    pub fn submit<F>(&self, job: F)
113    where
114        F: FnOnce() + Send + 'static,
115    {
116        if let Some(sender) = &self.sender {
117            let _ = sender.send(Message::Run(Box::new(job)));
118        }
119    }
120}
121
122impl Drop for WorkerPool {
123    fn drop(&mut self) {
124        // Tell every worker to stop, then drop the sender so any worker racing
125        // on `recv()` also observes a disconnected channel and breaks.
126        if let Some(sender) = &self.sender {
127            for _ in &self.workers {
128                let _ = sender.send(Message::Shutdown);
129            }
130        }
131        self.sender = None;
132
133        for worker in self.workers.drain(..) {
134            let _ = worker.join();
135        }
136    }
137}
138
139#[cfg(test)]
140mod tests {
141    use super::WorkerPool;
142    use std::sync::atomic::{AtomicUsize, Ordering};
143    use std::sync::mpsc;
144    use std::sync::Arc;
145
146    #[test]
147    fn runs_submitted_jobs_and_collects_results() {
148        let pool = WorkerPool::new(4);
149        let (tx, rx) = mpsc::channel::<usize>();
150        let n = 32;
151        for i in 0..n {
152            let tx = tx.clone();
153            pool.submit(move || {
154                tx.send(i * 2).expect("result channel open");
155            });
156        }
157        drop(tx);
158        let mut got: Vec<usize> = rx.iter().take(n).collect();
159        got.sort_unstable();
160        let expected: Vec<usize> = (0..n).map(|i| i * 2).collect();
161        assert_eq!(got, expected);
162    }
163
164    #[test]
165    fn drop_joins_all_workers() {
166        let counter = Arc::new(AtomicUsize::new(0));
167        {
168            let pool = WorkerPool::new(3);
169            for _ in 0..12 {
170                let counter = Arc::clone(&counter);
171                pool.submit(move || {
172                    counter.fetch_add(1, Ordering::SeqCst);
173                });
174            }
175            // Pool drops here: pending jobs finish, workers join.
176        }
177        assert_eq!(counter.load(Ordering::SeqCst), 12);
178    }
179
180    #[test]
181    fn always_has_at_least_one_worker() {
182        let pool = WorkerPool::new(0);
183        assert_eq!(pool.thread_count(), 1);
184    }
185}
khora_control/worker_pool.rs

khora_control/
worker_pool.rs
