feat(supervisor): expose concurrent UDS auto-start

dstack/supervisor/client/Cargo.toml

@@ -27,6 +27,7 @@ serde.workspace = true
 http-body-util.workspace = true
 tracing-subscriber.workspace = true
 log.workspace = true
+libc.workspace = true
 fs-err.workspace = true
 futures.workspace = true
 

dstack/supervisor/client/src/lib.rs

@@ -11,6 +11,37 @@ use supervisor::{ProcessConfig, ProcessInfo, Response};
 
 pub use supervisor;
 
+#[cfg(unix)]
+fn acquire_uds_start_lock(uds: &Path) -> Result<std::fs::File> {
+    use std::fs::OpenOptions;
+    use std::os::unix::fs::{MetadataExt as _, OpenOptionsExt as _};
+
+    let lock_path = uds.with_extension("lock");
+    let lock = OpenOptions::new()
+        .read(true)
+        .write(true)
+        .create(true)
+        .mode(0o600)
+        .custom_flags(libc::O_CLOEXEC | libc::O_NOFOLLOW)
+        .open(&lock_path)
+        .with_context(|| {
+            format!(
+                "Failed to open Supervisor start lock {}",
+                lock_path.display()
+            )
+        })?;
+    let metadata = lock.metadata()?;
+    let effective_uid = unsafe { libc::geteuid() };
+    if metadata.uid() != effective_uid || metadata.mode() & 0o077 != 0 {
+        anyhow::bail!("Supervisor start lock is not owner-only");
+    }
+    let result = unsafe { libc::flock(std::os::fd::AsRawFd::as_raw_fd(&lock), libc::LOCK_EX) };
+    if result != 0 {
+        return Err(std::io::Error::last_os_error()).context("Failed to lock Supervisor startup");
+    }
+    Ok(lock)
+}
+
 #[derive(Debug, Clone)]
 pub struct SupervisorClient {
     base_url: Arc<String>,
@@ -41,6 +72,12 @@ impl SupervisorClient {
             anyhow::bail!("Failed to connect to supervisor at {uri}");
         }
         info!("Failed to connect to supervisor at {uri}, trying to start supervisor");
+        let _start_lock = acquire_uds_start_lock(uds.as_ref())?;
+        // Another caller may have completed startup while this caller waited.
+        if client.probe(Duration::from_millis(500)).await.is_ok() {
+            info!("Connected to supervisor at {uri} after waiting for startup lock");
+            return Ok(client);
+        }
         // if the uds exists, remove it
         if std::path::Path::new(uds.as_ref()).exists() {
             fs_err::remove_file(uds.as_ref())?;
@@ -51,16 +88,28 @@ impl SupervisorClient {
         let log_file = log_file.as_ref().to_path_buf();
         std::thread::spawn(move || {
             // start supervisor
-            let result = std::process::Command::new(supervisor_path)
+            let mut command = std::process::Command::new(supervisor_path);
+            command
                 .arg("--uds")
                 .arg(uds)
                 .arg("--pid-file")
                 .arg(pid_file)
                 .arg("--log-file")
                 .arg(log_file)
                 .args(if detached { &["--detach"][..] } else { &[] })
-                .env("RUST_LOG", "info,rocket=warn")
-                .output();
+                .env("RUST_LOG", "info,rocket=warn");
+            #[cfg(unix)]
+            {
+                use std::os::unix::process::CommandExt as _;
+
+                unsafe {
+                    command.pre_exec(|| {
+                        libc::umask(0o077);
+                        Ok(())
+                    });
+                }
+            }
+            let result = command.output();
             let output = match result {
                 Ok(output) => output,
                 Err(err) => {

dstack/supervisor/client/src/main.rs

@@ -13,6 +13,26 @@ struct Cli {
     #[arg(long, default_value = "unix:/var/run/supervisor.sock")]
     base_url: String,
 
+    /// Start a missing Supervisor before connecting to a trusted Unix socket.
+    #[arg(long, requires_all = ["supervisor_path", "pid_file", "log_file"])]
+    auto_start: bool,
+
+    /// Supervisor executable used only with --auto-start.
+    #[arg(long)]
+    supervisor_path: Option<std::path::PathBuf>,
+
+    /// PID file passed to an auto-started Supervisor.
+    #[arg(long)]
+    pid_file: Option<std::path::PathBuf>,
+
+    /// Log file passed to an auto-started Supervisor.
+    #[arg(long)]
+    log_file: Option<std::path::PathBuf>,
+
+    /// Detach an auto-started Supervisor process.
+    #[arg(long, requires = "auto_start")]
+    detached: bool,
+
     #[command(subcommand)]
     command: Commands,
 }
@@ -50,11 +70,37 @@ async fn main() -> Result<()> {
     {
         use tracing_subscriber::{fmt, EnvFilter};
         let filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info"));
-        fmt().with_env_filter(filter).with_ansi(false).init();
+        fmt()
+            .with_env_filter(filter)
+            .with_ansi(false)
+            .with_writer(std::io::stderr)
+            .init();
     }
 
     let cli = Cli::parse();
-    let client = SupervisorClient::new(&cli.base_url);
+    let client = if cli.auto_start {
+        let uds = cli
+            .base_url
+            .strip_prefix("unix:")
+            .ok_or_else(|| anyhow::anyhow!("--auto-start requires a unix: base URL"))?;
+        SupervisorClient::start_and_connect_uds(
+            cli.supervisor_path
+                .as_deref()
+                .ok_or_else(|| anyhow::anyhow!("missing --supervisor-path"))?,
+            uds,
+            cli.pid_file
+                .as_deref()
+                .ok_or_else(|| anyhow::anyhow!("missing --pid-file"))?,
+            cli.log_file
+                .as_deref()
+                .ok_or_else(|| anyhow::anyhow!("missing --log-file"))?,
+            cli.detached,
+            true,
+        )
+        .await?
+    } else {
+        SupervisorClient::new(&cli.base_url)
+    };
 
     match cli.command {
         Commands::Deploy { id, command, args } => {

dstack/supervisor/src/web_api.rs

@@ -6,7 +6,7 @@ use anyhow::{anyhow, Result};
 use or_panic::ResultOrPanic;
 use rocket::figment::Figment;
 use rocket::serde::json::Json;
-use rocket::{delete, get, post, routes, Build, Rocket, State};
+use rocket::{delete, get, post, routes, Build, Rocket, Shutdown, State};
 use serde::{Deserialize, Serialize};
 use tokio::signal;
 use tracing::info;
@@ -81,8 +81,12 @@ fn clear(supervisor: &State<Supervisor>) -> Json<Response<()>> {
 }
 
 #[post("/shutdown")]
-async fn shutdown(supervisor: &State<Supervisor>) -> Json<Response<()>> {
-    to_json(perform_shutdown(supervisor, false).await)
+async fn shutdown(supervisor: &State<Supervisor>, shutdown: Shutdown) -> Json<Response<()>> {
+    let result = supervisor.shutdown().await;
+    if result.is_ok() {
+        shutdown.notify();
+    }
+    to_json(result)
 }
 
 async fn perform_shutdown(supervisor: &Supervisor, force: bool) -> Result<()> {