stryke/
cluster.rs

1//! Persistent SSH worker pool dispatcher for `pmap_on`.
2//!
3//! ## Architecture
4//!
5//! ```text
6//!                                   ┌── slot 0 (ssh host1) ────┐
7//!                                   │  worker thread + ssh proc │
8//!                                   │  HELLO + SESSION_INIT     │
9//!                                   │  loop: take JOB from work │
10//!                                   │        send + read        │
11//!                                   │        push to results    │
12//!                                   └────────────────────────────┘
13//!                                   ┌── slot 1 (ssh host1) ────┐
14//!                                   │  worker thread + ssh proc │
15//!  main thread                      │  ...                      │
16//!  ┌─────────────────┐              └────────────────────────────┘
17//!  │ enqueue all jobs├──► work_tx ─►┌── slot 2 (ssh host2) ────┐
18//!  │ collect results │              │  ...                      │
19//!  └─────────────────┘              └────────────────────────────┘
20//!         ▲                                    │
21//!         │                                    ▼
22//!         └────────── result_rx ────────────────┘
23//! ```
24//!
25//! Each slot is one persistent `ssh HOST PE_PATH --remote-worker` process. The HELLO and
26//! SESSION_INIT handshakes happen once per slot lifetime, then the slot pulls JOB messages
27//! from a shared crossbeam channel and pushes responses to a result channel. Work-stealing
28//! emerges naturally: fast slots drain the queue faster, slow slots take fewer jobs.
29//!
30//! ## Fault tolerance
31//!
32//! When a slot's read or write fails (ssh died, network blip, remote crash), the worker
33//! thread re-enqueues the in-flight job to the shared queue with `attempts++` and exits.
34//! Other living slots pick the job up. A job is permanently failed when its attempt count
35//! reaches `cluster.max_attempts`. The whole map fails only when **every** slot is dead or
36//! every queued job has exhausted its retry budget.
37//!
38//! ## Per-job timeout
39//!
40//! Each `recv` from a slot's stdout uses a per-slot helper thread + bounded channel so the
41//! main wait is `crossbeam::channel::recv_timeout(cluster.job_timeout_ms)`. On timeout the
42//! ssh child is killed (SIGKILL), the slot is marked dead, and the in-flight job is
43//! re-enqueued just like any other slot failure.
44
45use std::io::Read;
46use std::process::{Child, Command, Stdio};
47use std::sync::Arc;
48use std::thread;
49use std::time::Duration;
50
51use crossbeam::channel::{bounded, select, unbounded, Receiver, RecvTimeoutError, Sender};
52
53use crate::remote_wire::{
54    frame_kind, perl_to_json_value, read_typed_frame, send_msg, HelloAck, HelloMsg, JobMsg,
55    JobRespMsg, SessionAck, SessionInit, PROTO_VERSION,
56};
57use crate::value::{RemoteCluster, RemoteSlot, StrykeValue};
58
59/// One unit of work tracked by the dispatcher. Carries the original sequence number for
60/// order-preserving result collection plus an attempt counter for retry accounting.
61#[derive(Debug, Clone)]
62pub struct DispatchJob {
63    /// `seq` field.
64    pub seq: u64,
65    /// `item` field.
66    pub item: serde_json::Value,
67    /// `attempts` field.
68    pub attempts: u32,
69}
70
71/// One result reported back to the main thread. `seq` matches the originating
72/// [`DispatchJob::seq`] so the dispatcher can stitch results back into source order.
73#[derive(Debug)]
74pub struct DispatchResult {
75    /// `seq` field.
76    pub seq: u64,
77    /// `outcome` field.
78    pub outcome: Result<StrykeValue, String>,
79}
80
81/// Run a `pmap_on` against a [`RemoteCluster`]. Blocks until every job has either succeeded
82/// or exhausted its retry budget. Returns the per-item results in the original list order
83/// or the first permanent failure.
84///
85/// `subs_prelude` and `block_src` are sent **once** per slot at session init.
86/// `capture` is the captured-lexical snapshot from the calling scope.
87/// `items` is the list of work items (already JSON-marshalled).
88pub fn run_cluster(
89    cluster: &RemoteCluster,
90    subs_prelude: String,
91    block_src: String,
92    capture: Vec<(String, serde_json::Value)>,
93    items: Vec<serde_json::Value>,
94) -> Result<Vec<StrykeValue>, String> {
95    if items.is_empty() {
96        return Ok(Vec::new());
97    }
98    if cluster.slots.is_empty() {
99        return Err("cluster: no slots".to_string());
100    }
101
102    // Shared work queue: every slot pulls from here, and slot threads re-enqueue on failure.
103    // Bounded so a misbehaving producer can't memory-blow; size is `slot_count * 2` to give
104    // each slot something to grab on the next iteration without blocking.
105    let work_capacity = (cluster.slots.len() * 2).max(8);
106    let (work_tx, work_rx) = bounded::<DispatchJob>(work_capacity);
107    let (result_tx, result_rx) = unbounded::<DispatchResult>();
108    // Shutdown signal: slot workers hold their own `work_tx` clones for re-enqueue, so the
109    // work channel never closes on its own once every initial job is sent. When all results
110    // have been collected the main thread drops `shutdown_tx`, which closes `shutdown_rx`
111    // and breaks the slot workers out of their blocking `recv` in `select!`.
112    let (shutdown_tx, shutdown_rx) = bounded::<()>(0);
113
114    // Spawn one worker thread per slot.
115    let mut handles = Vec::with_capacity(cluster.slots.len());
116    let session_init = Arc::new(SessionInit {
117        subs_prelude,
118        block_src,
119        capture,
120    });
121    let cluster_arc = Arc::new(cluster.clone());
122
123    for (slot_idx, slot) in cluster.slots.iter().enumerate() {
124        let slot = slot.clone();
125        let work_rx = work_rx.clone();
126        let work_tx = work_tx.clone();
127        let result_tx = result_tx.clone();
128        let shutdown_rx = shutdown_rx.clone();
129        let init = Arc::clone(&session_init);
130        let cluster = Arc::clone(&cluster_arc);
131        handles.push(thread::spawn(move || {
132            slot_worker_loop(
133                slot_idx,
134                slot,
135                init,
136                cluster,
137                work_rx,
138                work_tx,
139                result_tx,
140                shutdown_rx,
141            );
142        }));
143    }
144
145    // Drop the dispatcher-side handles so closing all slot copies signals queue shutdown.
146    drop(work_rx);
147    drop(result_tx);
148    drop(shutdown_rx);
149
150    // Seed the queue with the initial work.
151    for (i, item) in items.iter().enumerate() {
152        let job = DispatchJob {
153            seq: i as u64,
154            item: item.clone(),
155            attempts: 0,
156        };
157        if work_tx.send(job).is_err() {
158            return Err("cluster: all worker slots died before any work was sent".to_string());
159        }
160    }
161    drop(work_tx); // close once initial enqueue is done; slot threads keep their own clones
162
163    // Collect results in seq order. We allocate the full vector up-front and assign by
164    // index so we don't depend on receive order — slot threads complete jobs in any order.
165    let mut results: Vec<Option<Result<StrykeValue, String>>> =
166        (0..items.len()).map(|_| None).collect();
167    let mut received = 0usize;
168    while received < items.len() {
169        match result_rx.recv() {
170            Ok(r) => {
171                let idx = r.seq as usize;
172                if idx < results.len() && results[idx].is_none() {
173                    results[idx] = Some(r.outcome);
174                    received += 1;
175                }
176            }
177            Err(_) => {
178                // All slot threads dropped their senders before we got every result.
179                break;
180            }
181        }
182    }
183
184    // All results (or terminal slot-death) are in. Signal slots to stop pulling new work
185    // from the queue so they can run their SHUTDOWN handshake and exit cleanly. Without
186    // this drop the slot `select!` below would park forever on `work_rx.recv()` because
187    // every slot still holds its own `work_tx` clone for re-enqueue.
188    drop(shutdown_tx);
189
190    // Wait for slot threads to wind down.
191    for h in handles {
192        let _ = h.join();
193    }
194
195    // Stitch results back together; surface the first permanent failure if any.
196    let mut out = Vec::with_capacity(items.len());
197    for (i, slot_result) in results.into_iter().enumerate() {
198        match slot_result {
199            Some(Ok(v)) => out.push(v),
200            Some(Err(e)) => {
201                return Err(format!("cluster: job {i} failed permanently: {e}"));
202            }
203            None => {
204                return Err(format!(
205                    "cluster: job {i} never completed (all slots died?)"
206                ));
207            }
208        }
209    }
210    Ok(out)
211}
212
213/// Per-slot worker thread: spawn ssh, do HELLO + SESSION_INIT, then loop pulling JOBs from
214/// the shared queue. On any I/O failure the in-flight job is re-enqueued (or permanently
215/// failed if it has exhausted its retry budget) and the slot exits.
216#[allow(clippy::too_many_arguments)]
217fn slot_worker_loop(
218    slot_idx: usize,
219    slot: RemoteSlot,
220    init: Arc<SessionInit>,
221    cluster: Arc<RemoteCluster>,
222    work_rx: Receiver<DispatchJob>,
223    work_tx: Sender<DispatchJob>,
224    result_tx: Sender<DispatchResult>,
225    shutdown_rx: Receiver<()>,
226) {
227    // Spawn the ssh child + initial handshake. Failures here mean this slot never makes
228    // any progress; we exit and let other slots drain the queue.
229    let mut session = match SlotSession::open(&slot, &init, &cluster) {
230        Ok(s) => s,
231        Err(e) => {
232            eprintln!(
233                "cluster: slot {slot_idx} ({}) failed to start: {e}",
234                slot.host
235            );
236            return;
237        }
238    };
239
240    loop {
241        // Take one job, or bail out if the dispatcher has signalled shutdown. We can't rely
242        // on `work_rx` closing by itself because every slot holds its own `work_tx` clone
243        // for re-enqueue on transport failure — so the channel would stay open forever once
244        // all initial jobs are drained. The shutdown channel is the explicit wakeup.
245        let job = select! {
246            recv(work_rx) -> r => match r {
247                Ok(j) => j,
248                Err(_) => {
249                    // Queue fully closed (e.g. every slot dropped its `work_tx`) — done.
250                    let _ = session.shutdown();
251                    return;
252                }
253            },
254            recv(shutdown_rx) -> _ => {
255                // Dispatcher collected every result — clean SHUTDOWN frame + child wait.
256                let _ = session.shutdown();
257                return;
258            },
259        };
260
261        match session.run_job(&job, cluster.job_timeout_ms) {
262            Ok(resp) => {
263                if resp.ok {
264                    let pv = match crate::remote_wire::json_to_perl(&resp.result) {
265                        Ok(v) => v,
266                        Err(e) => {
267                            let _ = result_tx.send(DispatchResult {
268                                seq: job.seq,
269                                outcome: Err(format!("decode result: {e}")),
270                            });
271                            continue;
272                        }
273                    };
274                    let _ = result_tx.send(DispatchResult {
275                        seq: job.seq,
276                        outcome: Ok(pv),
277                    });
278                } else {
279                    // Permanent in-script failure — no point retrying, the body is the
280                    // same on every slot. Surface immediately.
281                    let _ = result_tx.send(DispatchResult {
282                        seq: job.seq,
283                        outcome: Err(resp.err_msg),
284                    });
285                }
286            }
287            Err(SlotError::Transport(e)) => {
288                // Wire-level failure — retry on a different slot if budget allows.
289                eprintln!(
290                    "cluster: slot {slot_idx} ({}) transport error: {e}; retrying job {}",
291                    slot.host, job.seq
292                );
293                requeue_or_fail(&work_tx, &result_tx, &cluster, job);
294                let _ = session.kill();
295                return;
296            }
297            Err(SlotError::Timeout) => {
298                eprintln!(
299                    "cluster: slot {slot_idx} ({}) timed out on job {}; retrying",
300                    slot.host, job.seq
301                );
302                requeue_or_fail(&work_tx, &result_tx, &cluster, job);
303                let _ = session.kill();
304                return;
305            }
306        }
307    }
308}
309
310fn requeue_or_fail(
311    work_tx: &Sender<DispatchJob>,
312    result_tx: &Sender<DispatchResult>,
313    cluster: &RemoteCluster,
314    mut job: DispatchJob,
315) {
316    job.attempts += 1;
317    if job.attempts >= cluster.max_attempts {
318        let _ = result_tx.send(DispatchResult {
319            seq: job.seq,
320            outcome: Err(format!(
321                "job exhausted retry budget after {} attempts",
322                job.attempts
323            )),
324        });
325        return;
326    }
327    if work_tx.send(job).is_err() {
328        // No live slots left to take the work — the dispatcher will detect this when
329        // result_rx closes with missing entries.
330    }
331}
332
333/// One persistent ssh child + the framed I/O handles to talk to it. Holds a stderr
334/// drainer thread so a verbose remote `stryke` doesn't fill its pipe and deadlock.
335struct SlotSession {
336    child: Child,
337    stdin: std::process::ChildStdin,
338    /// Channel that receives one `JobRespMsg` per JOB, with a per-job timeout. Backed by a
339    /// helper thread that loops on `read_typed_frame(stdout)` and forwards results.
340    resp_rx: Receiver<Result<JobRespMsg, String>>,
341}
342
343#[derive(Debug)]
344enum SlotError {
345    Transport(String),
346    Timeout,
347}
348
349impl SlotSession {
350    fn open(
351        slot: &RemoteSlot,
352        init: &SessionInit,
353        cluster: &RemoteCluster,
354    ) -> Result<Self, String> {
355        // Test-only bypass: `STRYKE_CLUSTER_LOCAL_BIN=path/to/stryke`
356        // spawns the worker locally instead of going through ssh. Used by
357        // the `~d>` and `pmap_on` test fixtures so we exercise the full
358        // session handshake + JOB-frame wire without requiring an sshd /
359        // ssh keys in CI. Slot `host` is ignored when this is set.
360        let local_override = std::env::var_os("STRYKE_CLUSTER_LOCAL_BIN");
361        let mut child = if let Some(bin) = local_override {
362            Command::new(bin)
363                .arg("--remote-worker")
364                .stdin(Stdio::piped())
365                .stdout(Stdio::piped())
366                .stderr(Stdio::piped())
367                .spawn()
368                .map_err(|e| format!("spawn local worker: {e}"))?
369        } else {
370            // ssh -o ConnectTimeout=N HOST PE_PATH --remote-worker
371            let connect_timeout = (cluster.connect_timeout_ms / 1000).max(1);
372            Command::new("ssh")
373                .arg("-o")
374                .arg(format!("ConnectTimeout={connect_timeout}"))
375                .arg("-o")
376                .arg("BatchMode=yes")
377                .arg(&slot.host)
378                .arg(&slot.pe_path)
379                .arg("--remote-worker")
380                .stdin(Stdio::piped())
381                .stdout(Stdio::piped())
382                .stderr(Stdio::piped())
383                .spawn()
384                .map_err(|e| format!("spawn ssh: {e}"))?
385        };
386        let mut stdin = child
387            .stdin
388            .take()
389            .ok_or_else(|| "ssh stdin missing".to_string())?;
390        let mut stdout = child
391            .stdout
392            .take()
393            .ok_or_else(|| "ssh stdout missing".to_string())?;
394        let mut stderr = child
395            .stderr
396            .take()
397            .ok_or_else(|| "ssh stderr missing".to_string())?;
398
399        // Drain stderr in the background so a verbose worker can't deadlock its pipe.
400        thread::spawn(move || {
401            let mut buf = String::new();
402            let _ = stderr.read_to_string(&mut buf);
403            // Forward to our own stderr prefixed for visibility — operators want to see
404            // remote crashes when debugging cluster runs.
405            if !buf.trim().is_empty() {
406                eprintln!("[remote-worker] {}", buf.trim());
407            }
408        });
409
410        // 1. HELLO. Direct stdin write (the helper-thread response loop hasn't started yet).
411        let hello = HelloMsg {
412            proto_version: PROTO_VERSION,
413            pe_version: env!("CARGO_PKG_VERSION").to_string(),
414        };
415        send_msg(&mut stdin, frame_kind::HELLO, &hello).map_err(|e| format!("send HELLO: {e}"))?;
416        let (kind, body) =
417            read_typed_frame(&mut stdout).map_err(|e| format!("read HELLO_ACK: {e}"))?;
418        if kind != frame_kind::HELLO_ACK {
419            return Err(format!("expected HELLO_ACK, got frame kind {kind:#04x}"));
420        }
421        let _: HelloAck =
422            bincode::deserialize(&body).map_err(|e| format!("decode HELLO_ACK: {e}"))?;
423
424        // 2. SESSION_INIT (`init` is `&SessionInit` via deref coercion from `&Arc<SessionInit>`).
425        send_msg(&mut stdin, frame_kind::SESSION_INIT, init)
426            .map_err(|e| format!("send SESSION_INIT: {e}"))?;
427        let (kind, body) =
428            read_typed_frame(&mut stdout).map_err(|e| format!("read SESSION_ACK: {e}"))?;
429        if kind != frame_kind::SESSION_ACK {
430            return Err(format!("expected SESSION_ACK, got frame kind {kind:#04x}"));
431        }
432        let ack: SessionAck =
433            bincode::deserialize(&body).map_err(|e| format!("decode SESSION_ACK: {e}"))?;
434        if !ack.ok {
435            return Err(format!("worker rejected session: {}", ack.err_msg));
436        }
437
438        // 3. Spin up the response helper thread. Each iteration reads one frame and
439        // forwards either the parsed JobRespMsg or an error string.
440        let (resp_tx, resp_rx) = bounded::<Result<JobRespMsg, String>>(1);
441        thread::spawn(move || loop {
442            match read_typed_frame(&mut stdout) {
443                Ok((kind, body)) if kind == frame_kind::JOB_RESP => {
444                    match bincode::deserialize::<JobRespMsg>(&body) {
445                        Ok(r) => {
446                            if resp_tx.send(Ok(r)).is_err() {
447                                return;
448                            }
449                        }
450                        Err(e) => {
451                            let _ = resp_tx.send(Err(format!("decode JOB_RESP: {e}")));
452                            return;
453                        }
454                    }
455                }
456                Ok((other, _)) => {
457                    let _ = resp_tx.send(Err(format!(
458                        "unexpected frame kind {other:#04x} in resp loop"
459                    )));
460                    return;
461                }
462                Err(e) => {
463                    let _ = resp_tx.send(Err(format!("read frame: {e}")));
464                    return;
465                }
466            }
467        });
468
469        Ok(Self {
470            child,
471            stdin,
472            resp_rx,
473        })
474    }
475
476    fn run_job(&mut self, job: &DispatchJob, timeout_ms: u64) -> Result<JobRespMsg, SlotError> {
477        let msg = JobMsg {
478            seq: job.seq,
479            item: job.item.clone(),
480        };
481        send_msg(&mut self.stdin, frame_kind::JOB, &msg)
482            .map_err(|e| SlotError::Transport(format!("send JOB: {e}")))?;
483        match self.resp_rx.recv_timeout(Duration::from_millis(timeout_ms)) {
484            Ok(Ok(r)) => Ok(r),
485            Ok(Err(e)) => Err(SlotError::Transport(e)),
486            Err(RecvTimeoutError::Timeout) => Err(SlotError::Timeout),
487            Err(RecvTimeoutError::Disconnected) => {
488                Err(SlotError::Transport("response channel closed".to_string()))
489            }
490        }
491    }
492
493    fn shutdown(&mut self) -> Result<(), String> {
494        // Best-effort SHUTDOWN frame; ignore errors because we're tearing down anyway.
495        let _ = send_msg::<_, ()>(&mut self.stdin, frame_kind::SHUTDOWN, &());
496        let _ = self.child.wait();
497        Ok(())
498    }
499
500    fn kill(&mut self) -> Result<(), String> {
501        let _ = self.child.kill();
502        let _ = self.child.wait();
503        Ok(())
504    }
505}
506
507/// Convenience: marshal a `Vec<StrykeValue>` into the JSON values the dispatcher needs.
508pub fn perl_items_to_json(items: &[StrykeValue]) -> Result<Vec<serde_json::Value>, String> {
509    items.iter().map(perl_to_json_value).collect()
510}
511
512#[cfg(test)]
513mod tests {
514    use super::*;
515
516    // ─── perl_items_to_json ──────────────────────────────────────────────
517
518    #[test]
519    fn perl_items_to_json_empty_slice_yields_empty_vec() {
520        let r = perl_items_to_json(&[]).unwrap();
521        assert!(r.is_empty());
522    }
523
524    #[test]
525    fn perl_items_to_json_preserves_order_and_length() {
526        let items = [
527            StrykeValue::integer(1),
528            StrykeValue::integer(2),
529            StrykeValue::integer(3),
530        ];
531        let r = perl_items_to_json(&items).unwrap();
532        assert_eq!(r.len(), 3);
533        assert_eq!(r[0], serde_json::json!(1));
534        assert_eq!(r[1], serde_json::json!(2));
535        assert_eq!(r[2], serde_json::json!(3));
536    }
537
538    #[test]
539    fn perl_items_to_json_marshals_string_values() {
540        let items = [
541            StrykeValue::string("a".into()),
542            StrykeValue::string("b".into()),
543        ];
544        let r = perl_items_to_json(&items).unwrap();
545        assert_eq!(r, vec![serde_json::json!("a"), serde_json::json!("b")]);
546    }
547
548    #[test]
549    fn perl_items_to_json_undef_becomes_json_null() {
550        let items = [StrykeValue::UNDEF];
551        let r = perl_items_to_json(&items).unwrap();
552        assert_eq!(r, vec![serde_json::Value::Null]);
553    }
554
555    #[test]
556    fn perl_items_to_json_float_preserved() {
557        let items = [StrykeValue::float(2.5)];
558        let r = perl_items_to_json(&items).unwrap();
559        assert!((r[0].as_f64().unwrap() - 2.5).abs() < 1e-12);
560    }
561
562    #[test]
563    fn perl_items_to_json_mixed_types_in_one_call() {
564        let items = [
565            StrykeValue::integer(7),
566            StrykeValue::string("x".into()),
567            StrykeValue::UNDEF,
568        ];
569        let r = perl_items_to_json(&items).unwrap();
570        assert_eq!(
571            r,
572            vec![
573                serde_json::json!(7),
574                serde_json::json!("x"),
575                serde_json::Value::Null,
576            ]
577        );
578    }
579
580    #[test]
581    fn perl_items_to_json_handles_array_ref_element() {
582        use parking_lot::RwLock;
583        use std::sync::Arc;
584        let inner = StrykeValue::array_ref(Arc::new(RwLock::new(vec![
585            StrykeValue::integer(1),
586            StrykeValue::integer(2),
587        ])));
588        let r = perl_items_to_json(&[inner]).unwrap();
589        assert_eq!(r, vec![serde_json::json!([1, 2])]);
590    }
591
592    #[test]
593    fn perl_items_to_json_length_matches_input_for_large_batch() {
594        // Pin contract: output length is always input length.
595        let items: Vec<StrykeValue> = (0..256).map(StrykeValue::integer).collect();
596        let r = perl_items_to_json(&items).unwrap();
597        assert_eq!(r.len(), 256);
598        assert_eq!(r[0], serde_json::json!(0));
599        assert_eq!(r[255], serde_json::json!(255));
600    }
601}
stryke/cluster.rs

stryke/
cluster.rs
