Make krun-server wait for child processes

Continue to accept and handle incoming connections, until the
server is idle and all child processes have exited.

crates/krun-guest/src/net.rs

@@ -1,5 +1,5 @@
 use std::fs;
-use std::os::unix::process::ExitStatusExt;
+use std::os::unix::process::ExitStatusExt as _;
 use std::process::Command;
 
 use anyhow::{anyhow, Context, Result};

crates/krun-guest/src/pulse.rs

@@ -1,8 +1,7 @@
 use std::collections::HashMap;
-use std::env;
-use std::fs;
 use std::path::Path;
 use std::process::{Command, Stdio};
+use std::{env, fs};
 
 use anyhow::{Context, Result};
 use utils::env::find_in_path;

crates/krun-server/Cargo.toml

@@ -13,7 +13,6 @@ anyhow = { workspace = true, features = ["std"] }
 bpaf = { workspace = true, features = [] }
 env_logger = { workspace = true, features = ["auto-color", "humantime", "unstable-kv"] }
 log = { workspace = true, features = ["kv"] }
-nix = { workspace = true, features = ["socket"] }
 serde = { workspace = true, features = [] }
 serde_json = { workspace = true, features = ["std"] }
 utils = { workspace = true, features = [] }

crates/krun-server/src/bin/krun-server.rs

@@ -1,32 +1,96 @@
 use std::net::TcpListener;
-use std::os::fd::AsRawFd;
-use std::panic;
+use std::os::unix::process::ExitStatusExt as _;
 use std::process::Command;
+use std::sync::mpsc::{channel, TryRecvError};
+use std::sync::{Arc, Condvar, Mutex};
 
-use anyhow::{Context, Result};
+use anyhow::Result;
 use krun_server::cli_options::options;
 use krun_server::server::start_server;
-use nix::sys::socket::{shutdown, Shutdown};
+use log::debug;
 
 fn main() -> Result<()> {
     env_logger::init();
 
     let options = options().run();
 
     let listener = TcpListener::bind(format!("0.0.0.0:{}", options.server_port))?;
-    let listener_fd = listener.as_raw_fd();
+    let (child_tx, child_rx) = channel();
+    let idle_pair = Arc::new((Mutex::new(true), Condvar::new()));
 
-    let server_thread = start_server(listener);
+    start_server(listener, child_tx, Arc::clone(&idle_pair));
 
-    Command::new(&options.command)
+    match Command::new(&options.command)
         .args(options.command_args)
         .status()
-        .with_context(|| format!("Failed to execute command {:?}", options.command))?;
-
-    shutdown(listener_fd, Shutdown::Both)?;
-    if let Err(err) = server_thread.join() {
-        panic::resume_unwind(err);
+    {
+        Ok(status) => {
+            if !status.success() {
+                if let Some(code) = status.code() {
+                    eprintln!(
+                        "{:?} process exited with status code: {code}",
+                        options.command
+                    );
+                } else {
+                    eprintln!(
+                        "{:?} process terminated by signal: {}",
+                        options.command,
+                        status
+                            .signal()
+                            .expect("either one of status code or signal should be set")
+                    );
+                }
+            }
+        },
+        Err(err) => {
+            eprintln!(
+                "Failed to execute {:?} as child process: {err}",
+                options.command
+            );
+        },
     }
 
-    Ok(())
+    println!("Waiting for other commands launched through this krun server to exit...");
+    println!("Press Ctrl+C to force quit");
+
+    let (idle_lock, idle_cvar) = &*idle_pair;
+
+    loop {
+        let mut child = {
+            let _guard = idle_cvar
+                .wait_while(idle_lock.lock().unwrap(), |idle| !*idle)
+                .unwrap();
+            match child_rx.try_recv() {
+                Ok(child) => child,
+                Err(TryRecvError::Empty) => {
+                    // Server is idle and no more child processes. We're done.
+                    return Ok(());
+                },
+                Err(TryRecvError::Disconnected) => {
+                    panic!("child_tx in server thread should not be gone");
+                },
+            }
+        };
+        let child_pid = child.id();
+        match child.wait() {
+            Ok(status) => {
+                debug!("child process {child_pid} exited");
+                if !status.success() {
+                    if let Some(code) = status.code() {
+                        eprintln!("child process {child_pid} exited with status code: {code}");
+                    } else {
+                        eprintln!(
+                            "child process {child_pid} terminated by signal: {}",
+                            status
+                                .signal()
+                                .expect("either one of status code or signal should be set")
+                        );
+                    }
+                }
+            },
+            Err(err) => {
+                eprintln!("Failed to wait for child process {child_pid} to exit: {err}");
+            },
+        }
+    }
 }

crates/krun-server/src/server.rs

@@ -1,32 +1,57 @@
 use std::collections::HashMap;
-use std::env;
 use std::io::{BufRead, BufReader, Write};
 use std::net::{TcpListener, TcpStream};
-use std::process::{Command, Stdio};
-use std::thread::{self, JoinHandle};
+use std::process::{Child, Command, Stdio};
+use std::sync::mpsc::Sender;
+use std::sync::{Arc, Condvar, Mutex};
+use std::{env, thread};
 
-use anyhow::{anyhow, Result};
+use anyhow::{anyhow, Context, Result};
 use log::debug;
-use utils::{launch::Launch, stdio::make_stdout_stderr};
+use utils::launch::Launch;
+use utils::stdio::make_stdout_stderr;
 
-pub fn start_server(listener: TcpListener) -> JoinHandle<()> {
+pub fn start_server(
+    listener: TcpListener,
+    child_tx: Sender<Child>,
+    idle_pair: Arc<(Mutex<bool>, Condvar)>,
+) {
     thread::spawn(move || {
-        if let Err(err) = work(listener) {
-            debug!(err:?; "server thread is terminating")
-        }
-    })
-}
+        let (idle_lock, idle_cvar) = &*idle_pair;
+        for stream in listener.incoming() {
+            {
+                let mut idle = idle_lock.lock().unwrap();
+                *idle = false;
+                idle_cvar.notify_one();
+            }
+            let stream = match stream {
+                Ok(stream) => stream,
+                Err(err) => {
+                    eprintln!("Failed to accept incoming connection: {err}");
+                    let mut idle = idle_lock.lock().unwrap();
+                    *idle = true;
+                    idle_cvar.notify_one();
+                    continue;
+                },
+            };
 
-fn work(listener: TcpListener) -> Result<()> {
-    for stream in listener.incoming() {
-        let stream = stream?;
-
-        if let Err(err) = handle_connection(stream) {
-            println!("Error processing client request: {err:?}");
+            match handle_connection(stream) {
+                Ok(child) => {
+                    child_tx
+                        .send(child)
+                        .expect("child_rx in main thread should not be gone");
+                },
+                Err(err) => {
+                    eprintln!("Failed to process client request: {err:?}");
+                },
+            }
+            let mut idle = idle_lock.lock().unwrap();
+            *idle = true;
+            idle_cvar.notify_one();
         }
-    }
 
-    Ok(())
+        unreachable!("server accept loop should not exit");
+    });
 }
 
 fn read_request(mut stream: &TcpStream) -> Result<Launch> {
@@ -43,32 +68,34 @@ fn read_request(mut stream: &TcpStream) -> Result<Launch> {
     }
 }
 
-fn handle_connection(mut stream: TcpStream) -> Result<()> {
+fn handle_connection(mut stream: TcpStream) -> Result<Child> {
     let mut envs: HashMap<String, String> = env::vars().collect();
 
     let Launch {
         command,
         command_args,
         env,
     } = read_request(&stream)?;
+    debug!(command:?, command_args:?, env:?; "received launch request");
     envs.extend(env);
 
     let (stdout, stderr) = make_stdout_stderr(&command, &envs)?;
 
-    let err = Command::new(&command)
+    let spawn_result = Command::new(&command)
         .args(command_args)
         .envs(envs)
         .stdin(Stdio::null())
         .stdout(stdout)
         .stderr(stderr)
-        .spawn();
-    if let Err(err) = err {
-        let msg = format!("Failed to execute command {command:?}: {err}");
+        .spawn()
+        .with_context(|| format!("Failed to execute {command:?} as child process"));
+    if let Err(err) = &spawn_result {
+        let msg = format!("{err:?}");
         stream.write_all(msg.as_bytes()).ok();
     } else {
         stream.write_all(b"OK").ok();
     }
     stream.flush().ok();
 
-    Ok(())
+    spawn_result
 }

rustfmt.toml

@@ -1,6 +1,6 @@
 edition = "2021"
 
-# empty_item_single_line = false
+# empty_item_single_line = true
 # error_on_line_overflow = true
 # format_code_in_doc_comments = true
 # format_strings = true