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