gioser

tawasuyu/brahman

Rust 582TOML 222Markdown 73JSON 198635 Sutichaykuna

⎇ b08d5cbe0a6222f5b1479f40e68346b5d38246a8 crates/modules/shuma/shuma-exec/src/lib.rs Rust · 551 siq'ikuna · 18.1 KB Ñawpa willay Pi imata llamkarqan ch'uya

feat(shuma): ejecución directa + captura configurable por sesión b08d5cbe0 · sergio ·

1//! `shuma-exec` — ejecución de comandos del shell con salida en streaming.
2//!
3//! Dos modos de ejecución, un mismo contrato de eventos:
4//!
5//! - [`Exec::Direct`] — brahman lanza y **conecta los procesos él mismo**:
6//!   un `Command` por etapa del pipe, los pipes cableados con descriptores
7//!   reales. Control total del árbol de procesos (matar todo el pipe de
8//!   un golpe). Es el modo preferido.
9//! - [`Exec::Shell`] — delega a un shell externo (`bash -c "<line>"`).
10//!   Reservado para sintaxis que el modo directo aún no absorbe (globs,
11//!   `$VAR`, redirecciones, `&&`). bash es **sólo un parser de sintaxis**,
12//!   no el ejecutor por defecto.
13//!
14//! **Captura acotada.** [`CommandSpec::capture_limit`] topa los bytes en
15//! RAM; pasado el tope, o se **descarta** ([`RunEvent::Truncated`]) o se
16//! **vuelca a un archivo** si hay [`CommandSpec::spill_path`]
17//! ([`RunEvent::Spilled`]). En ambos casos el pipe se sigue drenando, así
18//! el proceso no se bloquea.
19//!
20//! **Reproceso.** [`CommandSpec::stdin_data`] alimenta un texto por la
21//! entrada estándar: reprocesa la salida capturada de un comando previo
22//! sin volver a correr el original.
23
24#![forbid(unsafe_code)]
25
26use std::fs::File;
27use std::io::{BufRead, BufReader, Read, Write};
28use std::path::PathBuf;
29use std::process::{Child, Command, Stdio};
30use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
31use std::sync::mpsc::{self, Receiver, Sender, TryRecvError};
32use std::sync::{Arc, Mutex};
33use std::thread::JoinHandle;
34
35/// Una etapa del pipe en ejecución directa: un binario y sus argumentos
36/// ya resueltos (sin comillas, sin metacaracteres).
37#[derive(Debug, Clone)]
38pub struct StageSpec {
39    pub program: String,
40    pub args: Vec<String>,
41}
42
43/// Cómo ejecutar.
44#[derive(Debug, Clone)]
45pub enum Exec {
46    /// Vía un shell externo — `program -c "<line>"`.
47    Shell { line: String, program: String },
48    /// Directo — brahman lanza y conecta cada etapa.
49    Direct { stages: Vec<StageSpec> },
50}
51
52/// Qué ejecutar y con qué política de captura.
53#[derive(Debug, Clone)]
54pub struct CommandSpec {
55    pub exec: Exec,
56    pub cwd: String,
57    /// Tope de captura en bytes; `0` = sin límite.
58    pub capture_limit: usize,
59    /// Si está, la salida que excede el tope se vuelca a este archivo.
60    pub spill_path: Option<PathBuf>,
61    /// Texto a alimentar por stdin — para reprocesar una salida previa.
62    pub stdin_data: Option<String>,
63}
64
65impl CommandSpec {
66    /// Ejecución vía `bash -c "<line>"`.
67    pub fn shell(line: impl Into<String>, cwd: impl Into<String>) -> Self {
68        Self {
69            exec: Exec::Shell { line: line.into(), program: "bash".into() },
70            cwd: cwd.into(),
71            capture_limit: 0,
72            spill_path: None,
73            stdin_data: None,
74        }
75    }
76
77    /// Ejecución directa de un pipe de etapas.
78    pub fn direct(stages: Vec<StageSpec>, cwd: impl Into<String>) -> Self {
79        Self {
80            exec: Exec::Direct { stages },
81            cwd: cwd.into(),
82            capture_limit: 0,
83            spill_path: None,
84            stdin_data: None,
85        }
86    }
87
88    /// Fija el tope de captura en bytes (encadenable).
89    pub fn with_limit(mut self, bytes: usize) -> Self {
90        self.capture_limit = bytes;
91        self
92    }
93
94    /// Vuelca la salida excedente a `path` en vez de descartarla.
95    pub fn with_spill(mut self, path: PathBuf) -> Self {
96        self.spill_path = Some(path);
97        self
98    }
99
100    /// Alimenta `data` por la entrada estándar del proceso (encadenable).
101    pub fn with_stdin(mut self, data: impl Into<String>) -> Self {
102        self.stdin_data = Some(data.into());
103        self
104    }
105}
106
107/// Un evento de la ejecución de un comando.
108#[derive(Debug, Clone, PartialEq, Eq)]
109pub enum RunEvent {
110    /// Una línea de salida estándar.
111    Stdout(String),
112    /// Una línea de salida de error.
113    Stderr(String),
114    /// La captura alcanzó su tope; lo que sigue se descarta.
115    Truncated,
116    /// La captura alcanzó su tope; el resto se vuelca al archivo dado.
117    Spilled(String),
118    /// El proceso terminó con este código de salida.
119    Exited(i32),
120    /// El proceso no pudo siquiera lanzarse.
121    Failed(String),
122}
123
124impl RunEvent {
125    /// `true` si el evento cierra la ejecución (`Exited` o `Failed`).
126    pub fn is_terminal(&self) -> bool {
127        matches!(self, RunEvent::Exited(_) | RunEvent::Failed(_))
128    }
129}
130
131/// Destino de volcado de la salida excedente — compartido entre lectores.
132struct SpillSink {
133    path: PathBuf,
134    file: Mutex<Option<File>>,
135}
136
137impl SpillSink {
138    /// Escribe una línea excedente al archivo (lo abre perezosamente).
139    fn write_line(&self, line: &str) {
140        let mut g = self.file.lock().expect("spill lock");
141        if g.is_none() {
142            *g = File::create(&self.path).ok();
143        }
144        if let Some(f) = g.as_mut() {
145            let _ = writeln!(f, "{line}");
146        }
147    }
148}
149
150/// Asa de un comando en ejecución. El consumidor la conserva y drena sus
151/// eventos cuando le conviene.
152pub struct RunHandle {
153    rx: Receiver<RunEvent>,
154    finished: bool,
155    /// Los procesos, compartidos con el hilo coordinador para poder
156    /// matarlos — todas las etapas de un pipe directo.
157    children: Arc<Mutex<Vec<Child>>>,
158}
159
160impl RunHandle {
161    /// Mata todos los procesos del comando. No hace nada si ya terminaron.
162    pub fn kill(&self) {
163        if let Ok(mut guard) = self.children.lock() {
164            for c in guard.iter_mut() {
165                let _ = c.kill();
166            }
167        }
168    }
169
170    /// Drena todos los eventos disponibles ahora mismo, sin bloquear.
171    pub fn try_events(&mut self) -> Vec<RunEvent> {
172        let mut out = Vec::new();
173        loop {
174            match self.rx.try_recv() {
175                Ok(ev) => {
176                    if ev.is_terminal() {
177                        self.finished = true;
178                    }
179                    out.push(ev);
180                }
181                Err(TryRecvError::Empty) => break,
182                Err(TryRecvError::Disconnected) => {
183                    self.finished = true;
184                    break;
185                }
186            }
187        }
188        out
189    }
190
191    /// Bloquea hasta que el proceso termine y devuelve todos sus eventos.
192    pub fn wait_all(&mut self) -> Vec<RunEvent> {
193        let mut out = Vec::new();
194        while let Ok(ev) = self.rx.recv() {
195            let terminal = ev.is_terminal();
196            out.push(ev);
197            if terminal {
198                self.finished = true;
199            }
200        }
201        self.finished = true;
202        out
203    }
204
205    /// `true` si ya se observó el evento terminal.
206    pub fn is_finished(&self) -> bool {
207        self.finished
208    }
209}
210
211/// Lanza un hilo lector de un flujo, con captura acotada. Pasado el tope
212/// emite (una vez) `Truncated` o `Spilled` y deriva el resto al sumidero
213/// de volcado o a la basura — pero **sigue drenando** el pipe.
214#[allow(clippy::too_many_arguments)]
215fn spawn_reader<R: Read + Send + 'static>(
216    stream: R,
217    tx: Sender<RunEvent>,
218    make: fn(String) -> RunEvent,
219    limit: usize,
220    counter: Arc<AtomicUsize>,
221    announced: Arc<AtomicBool>,
222    spill: Option<Arc<SpillSink>>,
223) -> JoinHandle<()> {
224    std::thread::spawn(move || {
225        for line in BufReader::new(stream).lines().map_while(Result::ok) {
226            let total =
227                counter.fetch_add(line.len() + 1, Ordering::Relaxed) + line.len() + 1;
228            if limit != 0 && total > limit {
229                let first = !announced.swap(true, Ordering::Relaxed);
230                match &spill {
231                    Some(sink) => {
232                        if first {
233                            let _ = tx
234                                .send(RunEvent::Spilled(sink.path.display().to_string()));
235                        }
236                        sink.write_line(&line);
237                    }
238                    None => {
239                        if first {
240                            let _ = tx.send(RunEvent::Truncated);
241                        }
242                    }
243                }
244                continue; // descarta/vuelca, pero sigue leyendo el pipe
245            }
246            if tx.send(make(line)).is_err() {
247                break;
248            }
249        }
250    })
251}
252
253/// Resultado de lanzar los procesos: lo que el coordinador necesita.
254struct Spawned {
255    children: Vec<Child>,
256    stdin: Option<std::process::ChildStdin>,
257    stdout: Option<std::process::ChildStdout>,
258    stderrs: Vec<std::process::ChildStderr>,
259}
260
261/// Lanza un único proceso shell (`program -c "<line>"`).
262fn spawn_shell(line: &str, program: &str, cwd: &str, want_stdin: bool) -> std::io::Result<Spawned> {
263    let mut child = Command::new(program)
264        .arg("-c")
265        .arg(line)
266        .current_dir(cwd)
267        .stdin(if want_stdin { Stdio::piped() } else { Stdio::null() })
268        .stdout(Stdio::piped())
269        .stderr(Stdio::piped())
270        .spawn()?;
271    let stdin = child.stdin.take();
272    let stdout = child.stdout.take();
273    let stderrs = child.stderr.take().into_iter().collect();
274    Ok(Spawned { children: vec![child], stdin, stdout, stderrs })
275}
276
277/// Lanza un pipe de etapas conectándolas con descriptores reales.
278fn spawn_direct(stages: &[StageSpec], cwd: &str, want_stdin: bool) -> std::io::Result<Spawned> {
279    if stages.is_empty() {
280        return Err(std::io::Error::new(
281            std::io::ErrorKind::InvalidInput,
282            "pipe vacío",
283        ));
284    }
285    let n = stages.len();
286    let mut children: Vec<Child> = Vec::with_capacity(n);
287    let mut prev_stdout: Option<std::process::ChildStdout> = None;
288
289    for (i, st) in stages.iter().enumerate() {
290        let mut cmd = Command::new(&st.program);
291        cmd.args(&st.args)
292            .current_dir(cwd)
293            .stdout(Stdio::piped())
294            .stderr(Stdio::piped());
295        if i == 0 {
296            cmd.stdin(if want_stdin { Stdio::piped() } else { Stdio::null() });
297        } else {
298            // La etapa anterior alimenta a ésta por su stdout.
299            cmd.stdin(Stdio::from(prev_stdout.take().expect("stdout previo")));
300        }
301        match cmd.spawn() {
302            Ok(mut child) => {
303                if i + 1 < n {
304                    prev_stdout = child.stdout.take();
305                }
306                children.push(child);
307            }
308            Err(e) => {
309                // Si una etapa no arranca, se matan las ya lanzadas.
310                for mut c in children {
311                    let _ = c.kill();
312                }
313                return Err(std::io::Error::new(
314                    e.kind(),
315                    format!("{}: {e}", st.program),
316                ));
317            }
318        }
319    }
320
321    let stdin = children.first_mut().and_then(|c| c.stdin.take());
322    let stdout = children.last_mut().and_then(|c| c.stdout.take());
323    let stderrs = children.iter_mut().filter_map(|c| c.stderr.take()).collect();
324    Ok(Spawned { children, stdin, stdout, stderrs })
325}
326
327/// Lanza `spec` y devuelve un [`RunHandle`] desde el que drenar la
328/// salida. La función vuelve de inmediato: el proceso corre en hilos.
329pub fn run(spec: &CommandSpec) -> RunHandle {
330    let (tx, rx) = mpsc::channel();
331    let spec = spec.clone();
332    let cell: Arc<Mutex<Vec<Child>>> = Arc::new(Mutex::new(Vec::new()));
333    let cell_thread = Arc::clone(&cell);
334
335    std::thread::spawn(move || {
336        let want_stdin = spec.stdin_data.is_some();
337        let spawned = match &spec.exec {
338            Exec::Shell { line, program } => {
339                spawn_shell(line, program, &spec.cwd, want_stdin)
340            }
341            Exec::Direct { stages } => spawn_direct(stages, &spec.cwd, want_stdin),
342        };
343        let Spawned { children, stdin, stdout, stderrs } = match spawned {
344            Ok(s) => s,
345            Err(e) => {
346                let _ = tx.send(RunEvent::Failed(e.to_string()));
347                return;
348            }
349        };
350
351        // Alimenta stdin (reproceso) en su propio hilo.
352        if let (Some(data), Some(mut sink)) = (spec.stdin_data.clone(), stdin) {
353            std::thread::spawn(move || {
354                let _ = sink.write_all(data.as_bytes());
355            });
356        }
357
358        // Comparte los procesos para que `kill` los alcance.
359        if let Ok(mut g) = cell_thread.lock() {
360            *g = children;
361        }
362
363        // Captura acotada: contador y aviso compartidos por todos los
364        // lectores; un sumidero de volcado opcional.
365        let counter = Arc::new(AtomicUsize::new(0));
366        let announced = Arc::new(AtomicBool::new(false));
367        let spill = spec
368            .spill_path
369            .clone()
370            .map(|path| Arc::new(SpillSink { path, file: Mutex::new(None) }));
371        let limit = spec.capture_limit;
372
373        let mut readers: Vec<JoinHandle<()>> = Vec::new();
374        if let Some(s) = stdout {
375            readers.push(spawn_reader(
376                s,
377                tx.clone(),
378                RunEvent::Stdout,
379                limit,
380                Arc::clone(&counter),
381                Arc::clone(&announced),
382                spill.clone(),
383            ));
384        }
385        for s in stderrs {
386            readers.push(spawn_reader(
387                s,
388                tx.clone(),
389                RunEvent::Stderr,
390                limit,
391                Arc::clone(&counter),
392                Arc::clone(&announced),
393                spill.clone(),
394            ));
395        }
396        for h in readers {
397            let _ = h.join();
398        }
399
400        // Cosecha todas las etapas; el código de salida es el de la última.
401        let code = {
402            let mut g = cell_thread.lock().expect("children lock");
403            let mut last = -1;
404            for c in g.iter_mut() {
405                last = c.wait().ok().and_then(|s| s.code()).unwrap_or(-1);
406            }
407            last
408        };
409        let _ = tx.send(RunEvent::Exited(code));
410    });
411
412    RunHandle { rx, finished: false, children: cell }
413}
414
415#[cfg(test)]
416mod tests {
417    use super::*;
418
419    /// Ejecución directa de un único programa.
420    fn direct(program: &str, args: &[&str]) -> CommandSpec {
421        CommandSpec::direct(
422            vec![StageSpec {
423                program: program.into(),
424                args: args.iter().map(|s| s.to_string()).collect(),
425            }],
426            ".",
427        )
428    }
429
430    /// Pipe directo de varias etapas.
431    fn pipe(stages: &[(&str, &[&str])]) -> CommandSpec {
432        CommandSpec::direct(
433            stages
434                .iter()
435                .map(|(p, a)| StageSpec {
436                    program: p.to_string(),
437                    args: a.iter().map(|s| s.to_string()).collect(),
438                })
439                .collect(),
440            ".",
441        )
442    }
443
444    fn stdout_of(events: Vec<RunEvent>) -> Vec<String> {
445        events
446            .into_iter()
447            .filter_map(|e| match e {
448                RunEvent::Stdout(l) => Some(l),
449                _ => None,
450            })
451            .collect()
452    }
453
454    #[test]
455    fn direct_runs_a_single_program() {
456        let mut h = run(&direct("echo", &["hola", "mundo"]));
457        let events = h.wait_all();
458        assert!(events.contains(&RunEvent::Stdout("hola mundo".into())));
459        assert!(events.contains(&RunEvent::Exited(0)));
460    }
461
462    #[test]
463    fn direct_wires_a_pipeline_itself() {
464        // printf … | sort — brahman conecta los procesos, sin shell.
465        let mut h = run(&pipe(&[
466            ("printf", &["b\\na\\nc\\n"]),
467            ("sort", &[]),
468        ]));
469        assert_eq!(stdout_of(h.wait_all()), vec!["a", "b", "c"]);
470    }
471
472    #[test]
473    fn direct_three_stage_pipeline() {
474        let mut h = run(&pipe(&[
475            ("printf", &["3\\n1\\n2\\n1\\n"]),
476            ("sort", &[]),
477            ("uniq", &[]),
478        ]));
479        assert_eq!(stdout_of(h.wait_all()), vec!["1", "2", "3"]);
480    }
481
482    #[test]
483    fn direct_nonzero_exit_is_the_last_stage() {
484        let mut h = run(&direct("false", &[]));
485        assert!(h.wait_all().contains(&RunEvent::Exited(1)));
486    }
487
488    #[test]
489    fn direct_missing_program_fails_gracefully() {
490        let mut h = run(&direct("no-existe-binario-xyz", &[]));
491        let events = h.wait_all();
492        assert!(matches!(events.first(), Some(RunEvent::Failed(_))));
493    }
494
495    #[test]
496    fn shell_mode_still_works_for_complex_syntax() {
497        let mut h = run(&CommandSpec {
498            exec: Exec::Shell { line: "echo $((2 + 3))".into(), program: "sh".into() },
499            ..CommandSpec::shell("", ".")
500        });
501        assert!(h.wait_all().contains(&RunEvent::Stdout("5".into())));
502    }
503
504    #[test]
505    fn capture_limit_truncates_but_process_finishes() {
506        let mut h = run(&direct("seq", &["1", "20000"]).with_limit(400));
507        let events = h.wait_all();
508        assert!(events.contains(&RunEvent::Truncated));
509        assert!(events.contains(&RunEvent::Exited(0)));
510        assert!(stdout_of(events).len() < 20000);
511    }
512
513    #[test]
514    fn spill_writes_overflow_to_a_file() {
515        let path = std::env::temp_dir()
516            .join(format!("shuma-exec-spill-{}.log", std::process::id()));
517        let _ = std::fs::remove_file(&path);
518        let mut h = run(&direct("seq", &["1", "5000"])
519            .with_limit(200)
520            .with_spill(path.clone()));
521        let events = h.wait_all();
522        assert!(events.iter().any(|e| matches!(e, RunEvent::Spilled(_))));
523        assert!(events.contains(&RunEvent::Exited(0)));
524        // El archivo de volcado existe y tiene contenido.
525        let spilled = std::fs::read_to_string(&path).unwrap_or_default();
526        assert!(spilled.contains("5000"), "la cola se volcó al archivo");
527        let _ = std::fs::remove_file(&path);
528    }
529
530    #[test]
531    fn stdin_data_reprocessed_by_a_filter() {
532        let mut h = run(&direct("grep", &["beta"]).with_stdin("alfa\nbeta\nbetabel\ngamma"));
533        assert_eq!(stdout_of(h.wait_all()), vec!["beta", "betabel"]);
534    }
535
536    #[test]
537    fn kill_stops_a_long_running_pipeline() {
538        let mut h = run(&pipe(&[("sleep", &["30"]), ("cat", &[])]));
539        std::thread::sleep(std::time::Duration::from_millis(250));
540        h.kill();
541        let events = h.wait_all();
542        assert!(events.last().map(|e| e.is_terminal()).unwrap_or(false));
543    }
544
545    #[test]
546    fn terminal_event_detection() {
547        assert!(RunEvent::Exited(0).is_terminal());
548        assert!(!RunEvent::Truncated.is_terminal());
549        assert!(!RunEvent::Spilled("x".into()).is_terminal());
550    }
551}