guest/server: pass exit code back to pid1

libkrun's built-in init supports reporting the exit code back to the host
from the program it launches. Let's not break the chain here!

Signed-off-by: Val Packett <[email protected]>

Author: valpackett

crates/muvm/src/guest/server.rs

@@ -43,7 +43,7 @@ pub async fn server_main(command: PathBuf, command_args: Vec<String>) -> Result<
     let mut state_rx = WatchStream::new(state_rx);
 
     let mut server_died = false;
-    let mut command_exited = false;
+    let mut command_exited_code = None;
 
     loop {
         tokio::select! {
@@ -59,14 +59,16 @@ pub async fn server_main(command: PathBuf, command_args: Vec<String>) -> Result<
                     server_died = true;
                 }
             },
-            res = &mut command_status, if !command_exited => {
+            res = &mut command_status, if command_exited_code.is_none() => {
+                command_exited_code = Some(1);
                 match res {
                     Ok(status) => {
                         if !status.success() {
                             if let Some(code) = status.code() {
                                 eprintln!(
                                     "{command:?} process exited with status code: {code}"
                                 );
+                                command_exited_code = Some(code);
                             } else {
                                 eprintln!(
                                     "{:?} process terminated by signal: {}",
@@ -76,6 +78,8 @@ pub async fn server_main(command: PathBuf, command_args: Vec<String>) -> Result<
                                         .expect("either one of status code or signal should be set")
                                 );
                             }
+                        } else {
+                            command_exited_code = Some(0);
                         }
                     },
                     Err(err) => {
@@ -84,14 +88,13 @@ pub async fn server_main(command: PathBuf, command_args: Vec<String>) -> Result<
                         );
                     },
                 }
-                command_exited = true;
             },
-            Some(state) = state_rx.next(), if command_exited => {
+            Some(state) = state_rx.next(), if command_exited_code.is_some() => {
                 if state.connection_idle() && state.child_processes() == 0 {
                     // Server is idle (not currently handling an accepted
                     // incoming connection) and no more child processes.
                     // We're done.
-                    return Ok(());
+                    std::process::exit(command_exited_code.unwrap());
                 }
                 println!(
                     "Waiting for {} other commands launched through this muvm server to exit...",