stryke/
ppool.rs

1//! Persistent thread pool (`ppool`) — workers pull jobs from a shared queue and run
2//! each task on a **fresh** [`VMHelper`] on an **existing** OS thread (no rayon task
3//! spawn per item; threads stay alive between jobs).
4
5use std::collections::{HashMap, VecDeque};
6use std::sync::atomic::{AtomicU64, Ordering};
7use std::sync::{Arc, Mutex};
8use std::thread::{self, JoinHandle};
9
10use crossbeam::channel::{unbounded, Receiver, Sender};
11
12use crate::error::{StrykeError, StrykeResult};
13use crate::scope::{AtomicArray, AtomicHash};
14use crate::value::{PerlPpool, StrykeSub, StrykeValue};
15use crate::vm_helper::{Flow, FlowOrError, VMHelper};
16
17/// Shared pool state (jobs in, results out-of-order; `PerlPpool::collect` reorders).
18pub struct PpoolInner {
19    /// `None` after the pool is shut down.
20    pub(crate) job_tx: Mutex<Option<Sender<PoolJob>>>,
21    /// `result_rx` field.
22    result_rx: Mutex<Receiver<(u64, StrykeValue)>>,
23    /// `pending` field.
24    pending: Mutex<VecDeque<(u64, StrykeValue)>>,
25    pub(crate) next_order: AtomicU64,
26    /// `collect_from` field.
27    collect_from: AtomicU64,
28    /// `workers` field.
29    workers: Mutex<Option<Vec<JoinHandle<()>>>>,
30}
31
32pub(crate) struct PoolJob {
33    order: u64,
34    sub: Arc<StrykeSub>,
35    arg: StrykeValue,
36    subs: HashMap<String, Arc<StrykeSub>>,
37    capture: Vec<(String, StrykeValue)>,
38    atomic_arrays: Vec<(String, AtomicArray)>,
39    atomic_hashes: Vec<(String, AtomicHash)>,
40}
41
42impl PerlPpool {
43    pub(crate) fn submit(
44        &self,
45        interp: &mut VMHelper,
46        args: &[StrykeValue],
47        line: usize,
48    ) -> StrykeResult<StrykeValue> {
49        if args.is_empty() {
50            return Err(StrykeError::runtime(
51                "submit() expects a code reference and optional argument for $_",
52                line,
53            ));
54        }
55        let Some(sub) = args[0].as_code_ref() else {
56            return Err(StrykeError::runtime(
57                "submit() first argument must be a CODE ref",
58                line,
59            ));
60        };
61        // One-arg form: bind worker `$_` from the caller's `$_` at submit time (postfix `for @tasks`
62        // sets `$_` each iteration). Two-arg form: explicit binding (may be `undef`).
63        let arg = if args.len() >= 2 {
64            args[1].clone()
65        } else {
66            interp.scope.get_scalar("_").clone()
67        };
68        let order = self.0.next_order.fetch_add(1, Ordering::SeqCst);
69        let subs = interp.subs.clone();
70        let (capture, atomic_arrays, atomic_hashes) = interp.scope.capture_with_atomics();
71        let job = PoolJob {
72            order,
73            sub: Arc::clone(&sub),
74            arg,
75            subs,
76            capture,
77            atomic_arrays,
78            atomic_hashes,
79        };
80        let tx = self
81            .0
82            .job_tx
83            .lock()
84            .map_err(|_| StrykeError::runtime("ppool: job queue poisoned", line))?;
85        let Some(sender) = tx.as_ref() else {
86            return Err(StrykeError::runtime("ppool: pool shut down", line));
87        };
88        sender
89            .send(job)
90            .map_err(|_| StrykeError::runtime("ppool: submit failed (pool shut down)", line))?;
91        Ok(StrykeValue::UNDEF)
92    }
93
94    pub(crate) fn collect(&self, line: usize) -> StrykeResult<StrykeValue> {
95        let start = self.0.collect_from.load(Ordering::SeqCst);
96        let end = self.0.next_order.load(Ordering::SeqCst);
97        let n = (end - start) as usize;
98        if n == 0 {
99            return Ok(StrykeValue::array(vec![]));
100        }
101
102        let mut slots: Vec<Option<StrykeValue>> = vec![None; n];
103        let mut count = 0usize;
104
105        {
106            let mut pending = self
107                .0
108                .pending
109                .lock()
110                .map_err(|_| StrykeError::runtime("ppool: pending buffer poisoned", line))?;
111            let mut keep = VecDeque::new();
112            for (o, v) in pending.drain(..) {
113                if o >= start && o < end {
114                    let idx = (o - start) as usize;
115                    if slots[idx].is_none() {
116                        slots[idx] = Some(v);
117                        count += 1;
118                    }
119                } else {
120                    keep.push_back((o, v));
121                }
122            }
123            *pending = keep;
124        }
125
126        let rx = self
127            .0
128            .result_rx
129            .lock()
130            .map_err(|_| StrykeError::runtime("ppool: collect lock poisoned", line))?;
131
132        while count < n {
133            let (o, v) = rx.recv().map_err(|_| {
134                StrykeError::runtime("ppool: result channel closed (workers stopped)", line)
135            })?;
136            if o < start {
137                continue;
138            }
139            if o >= end {
140                self.0
141                    .pending
142                    .lock()
143                    .map_err(|_| StrykeError::runtime("ppool: pending buffer poisoned", line))?
144                    .push_back((o, v));
145                continue;
146            }
147            let idx = (o - start) as usize;
148            if slots[idx].is_none() {
149                slots[idx] = Some(v);
150                count += 1;
151            }
152        }
153
154        self.0.collect_from.store(end, Ordering::SeqCst);
155        let out: Vec<StrykeValue> = slots
156            .into_iter()
157            .map(|s| s.unwrap_or(StrykeValue::UNDEF))
158            .collect();
159        Ok(StrykeValue::array(out))
160    }
161}
162
163impl Drop for PpoolInner {
164    fn drop(&mut self) {
165        if let Ok(mut g) = self.job_tx.lock() {
166            let _ = g.take();
167        }
168        if let Ok(mut g) = self.workers.lock() {
169            if let Some(handles) = g.take() {
170                for h in handles {
171                    let _ = h.join();
172                }
173            }
174        }
175    }
176}
177
178fn worker_loop(job_rx: Receiver<PoolJob>, result_tx: Sender<(u64, StrykeValue)>) {
179    while let Ok(job) = job_rx.recv() {
180        let mut interp = VMHelper::new();
181        interp.subs = job.subs;
182        interp.scope.restore_capture(&job.capture);
183        interp
184            .scope
185            .restore_atomics(&job.atomic_arrays, &job.atomic_hashes);
186        if let Some(env) = job.sub.closure_env.as_ref() {
187            interp.scope.restore_capture(env);
188        }
189        interp.enable_parallel_guard();
190        interp.scope.set_topic(job.arg);
191        interp.scope_push_hook();
192        let val = match interp.exec_block_no_scope(&job.sub.body) {
193            Ok(v) => v,
194            Err(FlowOrError::Flow(Flow::Return(v))) => v,
195            Err(_) => StrykeValue::UNDEF,
196        };
197        interp.scope_pop_hook();
198        let _ = result_tx.send((job.order, val));
199    }
200}
201
202/// Create a pool with `workers` OS threads (clamped to 1..=256). Each thread runs jobs
203/// sequentially; new [`VMHelper`] values are constructed per job (cheap vs thread spawn).
204pub fn create_pool(workers: usize) -> StrykeResult<StrykeValue> {
205    let workers = workers.clamp(1, 256);
206    let (job_tx, job_rx): (Sender<PoolJob>, Receiver<PoolJob>) = unbounded();
207    type ResultMsg = (u64, StrykeValue);
208    let (result_tx, result_rx): (Sender<ResultMsg>, Receiver<ResultMsg>) = unbounded();
209
210    let mut handles = Vec::with_capacity(workers);
211    for _ in 0..workers {
212        let jrx = job_rx.clone();
213        let rtx = result_tx.clone();
214        handles.push(thread::spawn(move || worker_loop(jrx, rtx)));
215    }
216    drop(job_rx);
217    drop(result_tx);
218
219    let inner = Arc::new(PpoolInner {
220        job_tx: Mutex::new(Some(job_tx)),
221        result_rx: Mutex::new(result_rx),
222        pending: Mutex::new(VecDeque::new()),
223        next_order: AtomicU64::new(0),
224        collect_from: AtomicU64::new(0),
225        workers: Mutex::new(Some(handles)),
226    });
227
228    Ok(StrykeValue::ppool(PerlPpool(inner)))
229}
230
231#[cfg(test)]
232mod tests {
233    use super::*;
234    use crate::parse;
235
236    #[test]
237    fn test_ppool_basic() {
238        let mut interp = VMHelper::new();
239        let pool_val = create_pool(2).expect("create_pool");
240        let pool = pool_val.as_ppool().expect("as_ppool");
241
242        let prog = parse("{ $_ * 2 }").expect("parse");
243        let body = match &prog.statements[0].kind {
244            crate::ast::StmtKind::Block(b) => b.clone(),
245            _ => panic!("expected block"),
246        };
247
248        let sub_val = StrykeValue::code_ref(Arc::new(StrykeSub {
249            name: "anon".to_string(),
250            params: vec![],
251            body,
252            prototype: None,
253            closure_env: None,
254            fib_like: None,
255        }));
256
257        for i in 1..=5 {
258            pool.submit(&mut interp, &[sub_val.clone(), StrykeValue::integer(i)], 1)
259                .expect("submit");
260        }
261
262        let results = pool.collect(1).expect("collect");
263        let arr = results.as_array_vec().expect("array");
264        assert_eq!(arr.len(), 5);
265        let mut ints: Vec<i64> = arr.iter().map(|v| v.to_int()).collect();
266        ints.sort();
267        assert_eq!(ints, vec![2, 4, 6, 8, 10]);
268    }
269}
stryke/ppool.rs

stryke/
ppool.rs
