Serve concurrent terminal sessions with exact teardown

crates/lazaret-host-coordinator/src/bootstrap.rs

@@ -20,8 +20,9 @@ use lazaret_storage_catalog::Catalog;
 
 use crate::control::ControlSessionAdmission;
 use crate::control::{
-    ControlError, InteractiveSessionExecutor, NonInteractiveSessionExecutor, PublishedControlFace,
-    SessionRegistry, StagedControlFace, TrustAnchor,
+    AcceptedControlConnection, ControlError, InteractiveSessionExecutor,
+    NonInteractiveSessionExecutor, PublishedControlFace, SessionRegistry, StagedControlFace,
+    TrustAnchor,
 };
 use crate::domain_store::{DomainStoreMutation, HostCoordinatorDomainStore};
 use crate::git_routing::GitRoutingCoordinator;
@@ -682,7 +683,17 @@ pub fn serve_production_control(
         .adapter()
         .published_control()
         .ok_or(ControlError::StaleGeneration)?;
-    face.serve_next_coordinated(
+    let connection = face.accept_connection()?;
+    serve_production_control_connection(lifecycle, &face, connection)
+}
+
+pub(crate) fn serve_production_control_connection(
+    lifecycle: &Mutex<ProductionHostCoordinatorLifecycle<'_, '_>>,
+    face: &PublishedControlFace,
+    connection: AcceptedControlConnection,
+) -> Result<(), ControlError> {
+    face.serve_connection_coordinated(
+        connection,
         |admission| {
             lifecycle
                 .lock()

crates/lazaret-host-coordinator/src/client_bootstrap.rs

@@ -20,13 +20,15 @@ use lazaret_linux::filesystem::try_exclusive_flock;
 use lazaret_runtime_supervisor::sandbox::production_sandbox_layout;
 use rustix::fs::{Mode, OFlags, openat};
 
-use crate::bootstrap::{ProductionHostCoordinatorLifecycle, serve_production_control};
+use crate::bootstrap::{ProductionHostCoordinatorLifecycle, serve_production_control_connection};
+use crate::control::ControlError;
 use crate::domain_store::HostCoordinatorDomainStore;
 use crate::production::{BootstrapOwnerGuard, ProductionColdStart, ProductionColdStartError};
 use crate::production_selector::{InstalledProductPaths, select_production_cold_start};
 use crate::waiter::BootstrapWaiterFace;
 
 pub const PRIVATE_CLIENT_BOOTSTRAP_ROLE: &str = "--lazaret-private-client-bootstrap-v1";
+const MAX_CONCURRENT_CONNECTIONS: usize = 32;
 
 #[derive(Debug)]
 pub struct PrivateClientBootstrapError(String);
@@ -169,11 +171,30 @@ fn serve_published(
         .ok_or_else(|| ProductionColdStartError::service("control face is absent"))?
         .set_nonblocking(true)
         .map_err(ProductionColdStartError::service)?;
+    std::thread::scope(|scope| {
+        let (completed, completions) =
+            std::sync::mpsc::channel::<(u64, Result<(), ControlError>)>();
+        let mut active: BTreeMap<u64, std::os::unix::net::UnixStream> = BTreeMap::new();
+        let mut next_connection = 0_u64;
         loop {
-        {
+            while let Ok((connection_id, result)) = completions.try_recv() {
+                active.remove(&connection_id);
+                connection_completion(result, false)?;
+            }
+            let face = {
                 let mut lifecycle = lifecycle.lock().map_err(|_| {
                     ProductionColdStartError::service("production lifecycle lock is poisoned")
                 })?;
+                let Some(face) = lifecycle.adapter().published_control() else {
+                    for socket in active.values() {
+                        let _ = socket.shutdown(std::net::Shutdown::Both);
+                    }
+                    drop(lifecycle);
+                    while !active.is_empty() {
+                        wait_for_connection_completion(&completions, &mut active, true)?;
+                    }
+                    return Ok(());
+                };
                 lifecycle
                     .reconcile_published_bootstrap_owner(
                         bootstrap_owner_session_id,
@@ -182,16 +203,33 @@ fn serve_published(
                             .map_err(ProductionColdStartError::service)?,
                     )
                     .map_err(ProductionColdStartError::service)?;
-            if lifecycle.adapter().published_control().is_none() {
-                return Ok(());
-            }
                 waiters
                     .borrow_mut()
                     .reconcile_published(&mut lifecycle)
                     .map_err(ProductionColdStartError::service)?;
+                face
+            };
+            if active.len() == MAX_CONCURRENT_CONNECTIONS {
+                wait_for_connection_completion(&completions, &mut active, false)?;
+                continue;
+            }
+            match face.accept_connection() {
+                Ok(connection) => {
+                    let cancellation = connection
+                        .cancellation_handle()
+                        .map_err(ProductionColdStartError::service)?;
+                    let connection_id = next_connection;
+                    next_connection = next_connection.checked_add(1).ok_or_else(|| {
+                        ProductionColdStartError::service("control connection counter overflow")
+                    })?;
+                    active.insert(connection_id, cancellation);
+                    let completed = completed.clone();
+                    scope.spawn(move || {
+                        let result =
+                            serve_production_control_connection(lifecycle, &face, connection);
+                        let _ = completed.send((connection_id, result));
+                    });
                 }
-        match serve_production_control(lifecycle) {
-            Ok(()) => {}
                 Err(crate::control::ControlError::Io(error))
                     if error.kind() == std::io::ErrorKind::WouldBlock =>
                 {
@@ -199,17 +237,42 @@ fn serve_published(
                 }
                 Err(error) => return Err(ProductionColdStartError::service(error)),
             }
-        if lifecycle
-            .lock()
-            .map_err(|_| {
-                ProductionColdStartError::service("production lifecycle lock is poisoned")
-            })?
-            .adapter()
-            .published_control()
-            .is_none()
-        {
-            return Ok(());
         }
+    })
+}
+
+fn wait_for_connection_completion(
+    completions: &std::sync::mpsc::Receiver<(u64, Result<(), ControlError>)>,
+    active: &mut BTreeMap<u64, std::os::unix::net::UnixStream>,
+    cancellation_requested: bool,
+) -> Result<(), ProductionColdStartError> {
+    let (connection_id, result) = completions
+        .recv()
+        .map_err(ProductionColdStartError::service)?;
+    active.remove(&connection_id);
+    connection_completion(result, cancellation_requested)
+}
+
+fn connection_completion(
+    result: Result<(), ControlError>,
+    cancellation_requested: bool,
+) -> Result<(), ProductionColdStartError> {
+    match result {
+        Ok(()) => Ok(()),
+        Err(ControlError::Io(error))
+            if cancellation_requested
+                && matches!(
+                    error.kind(),
+                    std::io::ErrorKind::UnexpectedEof
+                        | std::io::ErrorKind::ConnectionReset
+                        | std::io::ErrorKind::ConnectionAborted
+                        | std::io::ErrorKind::BrokenPipe
+                        | std::io::ErrorKind::NotConnected
+                ) =>
+        {
+            Ok(())
+        }
+        Err(error) => Err(ProductionColdStartError::service(error)),
     }
 }
 
@@ -463,4 +526,32 @@ mod tests {
         .map(OsString::from);
         assert!(parse_arguments(&extra).is_err());
     }
+
+    #[test]
+    fn withdrawal_suppresses_only_cancelled_socket_io() {
+        assert!(
+            connection_completion(
+                Err(ControlError::Io(std::io::Error::new(
+                    std::io::ErrorKind::UnexpectedEof,
+                    "withdrawn socket",
+                ))),
+                true,
+            )
+            .is_ok()
+        );
+        assert!(
+            connection_completion(
+                Err(ControlError::Io(std::io::Error::new(
+                    std::io::ErrorKind::UnexpectedEof,
+                    "ordinary socket failure",
+                ))),
+                false,
+            )
+            .is_err()
+        );
+        assert!(
+            connection_completion(Err(ControlError::Lifecycle("teardown failed".into())), true,)
+                .is_err()
+        );
+    }
 }

crates/lazaret-host-coordinator/src/control.rs

@@ -838,6 +838,16 @@ pub struct PublishedControlFace {
     interactive_executor: Arc<dyn InteractiveSessionExecutor>,
 }
 
+pub(crate) struct AcceptedControlConnection {
+    socket: std::os::unix::net::UnixStream,
+}
+
+impl AcceptedControlConnection {
+    pub(crate) fn cancellation_handle(&self) -> std::io::Result<std::os::unix::net::UnixStream> {
+        self.socket.try_clone()
+    }
+}
+
 pub trait NonInteractiveSessionExecutor: Send + Sync {
     /// # Errors
     /// Returns a runtime, stream, or exact-session execution failure.
@@ -1514,6 +1524,25 @@ impl PublishedControlFace {
     /// failure.
     pub fn serve_next_coordinated<Admit, Disconnect>(
         &self,
+        admit: Admit,
+        disconnect: Disconnect,
+    ) -> Result<(), ControlError>
+    where
+        Admit: FnMut(ControlSessionAdmission) -> Result<(), ControlError>,
+        Disconnect: FnMut(SessionId) -> Result<(), ControlError>,
+    {
+        let connection = self.accept_connection()?;
+        self.serve_connection_coordinated(connection, admit, disconnect)
+    }
+
+    pub(crate) fn accept_connection(&self) -> Result<AcceptedControlConnection, ControlError> {
+        let (socket, _) = self.listener.accept()?;
+        Ok(AcceptedControlConnection { socket })
+    }
+
+    pub(crate) fn serve_connection_coordinated<Admit, Disconnect>(
+        &self,
+        connection: AcceptedControlConnection,
         mut admit: Admit,
         mut disconnect: Disconnect,
     ) -> Result<(), ControlError>
@@ -1521,7 +1550,7 @@ impl PublishedControlFace {
         Admit: FnMut(ControlSessionAdmission) -> Result<(), ControlError>,
         Disconnect: FnMut(SessionId) -> Result<(), ControlError>,
     {
-        let (socket, _) = self.listener.accept()?;
+        let socket = connection.socket;
         let peer = socket_peercred(&socket).map_err(io::Error::from)?;
         require_peer_uid(peer.uid.as_raw(), self.uid)?;
         self.validate_publication()?;

crates/lazaret-host-coordinator/src/terminal.rs

@@ -1713,15 +1713,18 @@ impl SessionTerminalProcess for RunscSessionTerminal {
     }
 
     fn reap_exec(&mut self) -> io::Result<ProcessExit> {
-        if let Some(exit) = self.exit {
-            return Ok(exit);
-        }
+        let exit = if let Some(exit) = self.exit {
+            exit
+        } else {
             let exit = child_broker().reap_status(
                 self.exec
                     .take()
                     .ok_or_else(|| io::Error::other("exec client exit is unavailable"))?,
             )?;
             self.exit = Some(exit);
+            exit
+        };
+        remove_pid_file(&self.request.internal_pid_file)?;
         Ok(exit)
     }
 

crates/lazaret-host-coordinator/tests/runtime-composition.rs

@@ -302,6 +302,23 @@ fn run_session_terminal_matrix() -> Result<(), MatrixError> {
         .store
         .apply(&[DomainStoreMutation::PutSession(session.clone())])
         .map_err(failure)?;
+    let second_session_id = SessionId::new([41; 32]);
+    let mut second_session = Session {
+        generation_id: GenerationId::new(7),
+        session_id: second_session_id,
+        cli_id: ClientId::new([42; 32]),
+        session_anchor_id: SessionAnchorId::new([43; 32]),
+        session_phase: SessionPhase::Registering,
+    };
+    fixture
+        .store
+        .apply(&[DomainStoreMutation::PutSession(second_session.clone())])
+        .map_err(failure)?;
+    second_session.session_phase = SessionPhase::Attached;
+    fixture
+        .store
+        .apply(&[DomainStoreMutation::PutSession(second_session)])
+        .map_err(failure)?;
     session.session_phase = SessionPhase::Attached;
     fixture
         .store
@@ -317,6 +334,9 @@ fn run_session_terminal_matrix() -> Result<(), MatrixError> {
     let count = workspace.join("session-count");
     let finish = workspace.join("session-finish");
     let signals = workspace.join("session-signals");
+    let second_count = workspace.join("session-second-count");
+    let second_finish = workspace.join("session-second-finish");
+    let second_signals = workspace.join("session-second-signals");
     let line = "x".repeat(1024);
     let script = format!(
         "trap 'echo INT >> /workspace/session-signals' INT; \
@@ -326,6 +346,15 @@ fn run_session_terminal_matrix() -> Result<(), MatrixError> {
          n=0; while test ! -e /workspace/session-finish; do \
          echo '{line}'; n=$((n+1)); echo $n > /workspace/session-count; done; exit 23"
     );
+    let second_line = "y".repeat(1024);
+    let second_script = format!(
+        "trap 'echo INT >> /workspace/session-second-signals' INT; \
+         trap 'echo QUIT >> /workspace/session-second-signals' QUIT; \
+         trap 'echo TSTP >> /workspace/session-second-signals' TSTP; \
+         trap 'echo WINCH >> /workspace/session-second-signals' WINCH; \
+         n=0; while test ! -e /workspace/session-second-finish; do \
+         echo '{second_line}'; n=$((n+1)); echo $n > /workspace/session-second-count; done; exit 42"
+    );
     let request = TerminalAcquisition {
         domain_id: fixture.configuration.domain_id(),
         generation_id: GenerationId::new(7),
@@ -344,7 +373,20 @@ fn run_session_terminal_matrix() -> Result<(), MatrixError> {
         Arc::new(NoActivePrompt),
     );
     engine.acquire(&request).map_err(failure)?;
+    engine
+        .acquire(&TerminalAcquisition {
+            session_id: second_session_id,
+            internal_pid_file: runtime_domain.join("session-second-internal.pid"),
+            argv: vec![
+                b"/bin/sh".to_vec(),
+                b"-c".to_vec(),
+                second_script.into_bytes(),
+            ],
+            ..request.clone()
+        })
+        .map_err(failure)?;
     wait_for_path(&count)?;
+    wait_for_path(&second_count)?;
     let running = fixture
         .store
         .snapshot()
@@ -384,6 +426,9 @@ fn run_session_terminal_matrix() -> Result<(), MatrixError> {
             .map_err(failure)?;
     }
     wait_for_signal_lines(&mut engine, session_id, &signals, 4)?;
+    if second_signals.exists() {
+        return Err(failure("terminal signals crossed Session IDs"));
+    }
     let mut forwarded_before_prompt = 0_usize;
     for _ in 0..2_000 {
         let messages = engine
@@ -417,11 +462,21 @@ fn run_session_terminal_matrix() -> Result<(), MatrixError> {
         .map_err(failure)?;
     thread::sleep(Duration::from_millis(300));
     let before = fs::read_to_string(&count)?;
-    thread::sleep(Duration::from_millis(300));
+    let second_before = fs::read_to_string(&second_count)?;
+    for _ in 0..60 {
+        engine
+            .poll(second_session_id, GenerationId::new(7))
+            .map_err(failure)?;
+        thread::sleep(Duration::from_millis(5));
+    }
     let after = fs::read_to_string(&count)?;
+    let second_after = fs::read_to_string(&second_count)?;
     if before != after {
         return Err(failure("prompt route did not apply PTY backpressure"));
     }
+    if second_before == second_after {
+        return Err(failure("one Prompt paused another Session Terminal"));
+    }
     session.session_phase = SessionPhase::Attached;
     fixture
         .store
@@ -466,11 +521,49 @@ fn run_session_terminal_matrix() -> Result<(), MatrixError> {
     }
     if exit.and_then(lazaret_core::records::SessionExit::exit_code) != Some(23)
         || engine.contains(session_id)
+        || !engine.contains(second_session_id)
     {
         return Err(failure(
             "Session Terminal exit or no-reattachment cleanup failed",
         ));
     }
+    let second_before_first_exit = fs::read_to_string(&second_count)?;
+    for _ in 0..200 {
+        engine
+            .poll(second_session_id, GenerationId::new(7))
+            .map_err(failure)?;
+        if fs::read_to_string(&second_count).ok().as_deref()
+            != Some(second_before_first_exit.as_str())
+        {
+            break;
+        }
+        thread::sleep(Duration::from_millis(5));
+    }
+    if fs::read_to_string(&second_count).ok().as_deref() == Some(second_before_first_exit.as_str())
+    {
+        return Err(failure("first Session exit paused the second Session"));
+    }
+    fs::write(&second_finish, b"finish")?;
+    let mut second_exit = None;
+    for _ in 0..2_000 {
+        for message in engine
+            .poll(second_session_id, GenerationId::new(7))
+            .map_err(failure)?
+        {
+            if let SessionTerminalMessage::Exit(outcome) = message {
+                second_exit = Some(outcome);
+            }
+        }
+        if second_exit.is_some() {
+            break;
+        }
+        thread::sleep(Duration::from_millis(5));
+    }
+    if second_exit.and_then(lazaret_core::records::SessionExit::exit_code) != Some(42)
+        || engine.contains(second_session_id)
+    {
+        return Err(failure("second Session exit routing or cleanup failed"));
+    }
     runtime
         .stop(&SandboxCoordinator::new(&fixture.store))
         .map_err(failure)