Add graceful shutdown to wasmtime serve, fix flaky CI

This commit fixes the spurious CI failure found [here][fail]. The
problem here is that the infrastructure for testing `wasmtime serve` was
not properly waiting for all output in `finish`. There's some
infrastructure in the tests for spawning a subprocess and managing it,
but prior to this commit there was no way to shut down the server
gracefully and thus read all pending output from the child tasks.

The specific problem is that a specific error message was expected in
the logs after a request had been processed, but the `finish` method
wasn't reading the message. The reason for this is that `finish` had to
resort to `kill -9` on the child process as there was no other means of
shutting it down. This meant that the print, which happened after
request completion, might be killed and never happen.

The solution ended up in this commit is to (a) add a `--shutdown-addr`
CLI flag to `wasmtime serve` and (b) beef up handling of graceful
shutdown. Previously ctrl-c was used to exit the server but it didn't do
anything but drop all in-progress work. Instead graceful shutdown is now
handled by breaking out of the `accept` loop and then waiting for all
child tasks to complete, meaning that no http requests once received are
cancelled. In addition to ctrl-c the `--shutdown-addr` is used to listen
for a TCP connection which is a second means of cancellation. The reason
for this is that sending ctrl-c to a process is not nearly as trivial on
Windows as it is on Unix, so I didn't want to deal with all the console
bits necessary to get that aligned.

[fail]: https://github.com/bytecodealliance/wasmtime/actions/runs/13833291117/job/38702374924?pr=10390

Author: alexcrichton

src/commands/serve.rs

@@ -6,9 +6,10 @@ use std::{
     path::PathBuf,
     sync::{
         atomic::{AtomicBool, AtomicU64, Ordering},
-        Arc,
+        Arc, Mutex,
     },
 };
+use tokio::sync::Notify;
 use wasmtime::component::Linker;
 use wasmtime::{Engine, Store, StoreLimits};
 use wasmtime_wasi::{IoView, StreamError, StreamResult, WasiCtx, WasiCtxBuilder, WasiView};
@@ -85,12 +86,19 @@ pub struct ServeCommand {
     run: RunCommon,
 
     /// Socket address for the web server to bind to.
-    #[arg(long = "addr", value_name = "SOCKADDR", default_value_t = DEFAULT_ADDR)]
+    #[arg(long , value_name = "SOCKADDR", default_value_t = DEFAULT_ADDR)]
     addr: SocketAddr,
 
+    /// Socket address where, when connected to, will initiate a graceful
+    /// shutdown.
+    ///
+    /// Note that graceful shutdown is also supported on ctrl-c.
+    #[arg(long, value_name = "SOCKADDR")]
+    shutdown_addr: Option<SocketAddr>,
+
     /// Disable log prefixes of wasi-http handlers.
     /// if unspecified, logs will be prefixed with 'stdout|stderr [{req_id}] :: '
-    #[arg(long = "no-logging-prefix")]
+    #[arg(long)]
     no_logging_prefix: bool,
 
     /// The WebAssembly component to run.
@@ -134,17 +142,7 @@ impl ServeCommand {
             .enable_io()
             .build()?;
 
-        runtime.block_on(async move {
-            tokio::select! {
-                _ = tokio::signal::ctrl_c() => {
-                    Ok::<_, anyhow::Error>(())
-                }
-
-                res = self.serve() => {
-                    res
-                }
-            }
-        })?;
+        runtime.block_on(self.serve())?;
 
         Ok(())
     }
@@ -371,6 +369,30 @@ impl ServeCommand {
         let instance = linker.instantiate_pre(&component)?;
         let instance = ProxyPre::new(instance)?;
 
+        // Spawn background task(s) waiting for graceful shutdown signals. This
+        // always listens for ctrl-c but additionally can listen for a TCP
+        // connection to the specified address.
+        let shutdown = Arc::new(GracefulShutdown::default());
+        tokio::task::spawn({
+            let shutdown = shutdown.clone();
+            async move {
+                tokio::signal::ctrl_c().await.unwrap();
+                shutdown.requested.notify_one();
+            }
+        });
+        if let Some(addr) = self.shutdown_addr {
+            let listener = tokio::net::TcpListener::bind(addr).await?;
+            eprintln!(
+                "Listening for shutdown on tcp://{}/",
+                listener.local_addr()?
+            );
+            let shutdown = shutdown.clone();
+            tokio::task::spawn(async move {
+                let _ = listener.accept().await;
+                shutdown.requested.notify_one();
+            });
+        }
+
         let socket = match &self.addr {
             SocketAddr::V4(_) => tokio::net::TcpSocket::new_v4()?,
             SocketAddr::V6(_) => tokio::net::TcpSocket::new_v6()?,
@@ -403,9 +425,16 @@ impl ServeCommand {
         let handler = ProxyHandler::new(self, engine, instance);
 
         loop {
-            let (stream, _) = listener.accept().await?;
+            // Wait for a socket, but also "race" against shutdown to break out
+            // of this loop. Once the graceful shutdown signal is received then
+            // this loop exits immediately.
+            let (stream, _) = tokio::select! {
+                _ = shutdown.requested.notified() => break,
+                v = listener.accept() => v?,
+            };
             let stream = TokioIo::new(stream);
             let h = handler.clone();
+            let shutdown_guard = shutdown.clone().increment();
             tokio::task::spawn(async {
                 if let Err(e) = http1::Builder::new()
                     .keep_alive(true)
@@ -417,8 +446,75 @@ impl ServeCommand {
                 {
                     eprintln!("error: {e:?}");
                 }
+                drop(shutdown_guard);
             });
         }
+
+        // Upon exiting the loop we'll no longer process any more incoming
+        // connections but there may still be outstanding connections
+        // processing in child tasks. If there are wait for those to complete
+        // before shutting down completely. Also enable short-circuiting this
+        // wait with a second ctrl-c signal.
+        if shutdown.close() {
+            return Ok(());
+        }
+        eprintln!("Waiting for child tasks to exit, ctrl-c again to quit sooner...");
+        tokio::select! {
+            _ = tokio::signal::ctrl_c() => {}
+            _ = shutdown.complete.notified() => {}
+        }
+
+        Ok(())
+    }
+}
+
+/// Helper structure to manage graceful shutdown int he accept loop above.
+#[derive(Default)]
+struct GracefulShutdown {
+    /// Async notification that shutdown has been requested.
+    requested: Notify,
+    /// Async notification that shutdown has completed, signaled when
+    /// `notify_when_done` is `true` and `active_tasks` reaches 0.
+    complete: Notify,
+    /// Internal state related to what's in progress when shutdown is requested.
+    state: Mutex<GracefulShutdownState>,
+}
+
+#[derive(Default)]
+struct GracefulShutdownState {
+    active_tasks: u32,
+    notify_when_done: bool,
+}
+
+impl GracefulShutdown {
+    /// Increments the number of active tasks and returns a guard indicating
+    fn increment(self: Arc<Self>) -> impl Drop {
+        struct Guard(Arc<GracefulShutdown>);
+
+        let mut state = self.state.lock().unwrap();
+        assert!(!state.notify_when_done);
+        state.active_tasks += 1;
+        drop(state);
+
+        return Guard(self);
+
+        impl Drop for Guard {
+            fn drop(&mut self) {
+                let mut state = self.0.state.lock().unwrap();
+                state.active_tasks -= 1;
+                if state.notify_when_done && state.active_tasks == 0 {
+                    self.0.complete.notify_one();
+                }
+            }
+        }
+    }
+
+    /// Flags this state as done spawning tasks and returns whether there are no
+    /// more child tasks remaining.
+    fn close(&self) -> bool {
+        let mut state = self.state.lock().unwrap();
+        state.notify_when_done = true;
+        state.active_tasks == 0
     }
 }
 
@@ -533,7 +629,7 @@ async fn handle_request(
             // which should more clearly tell the user what went wrong. Note
             // that we assume the task has already exited at this point so the
             // `await` should resolve immediately.
-            let e = match task.await {
+            let e = match dbg!(task.await) {
                 Ok(Ok(())) => {
                     bail!("guest never invoked `response-outparam::set` method")
                 }

tests/all/cli_tests.rs

@@ -1530,6 +1530,7 @@ mod test_programs {
     struct WasmtimeServe {
         child: Option<Child>,
         addr: SocketAddr,
+        shutdown_addr: SocketAddr,
     }
 
     impl WasmtimeServe {
@@ -1548,61 +1549,73 @@ mod test_programs {
         }
 
         fn spawn(cmd: &mut Command) -> Result<WasmtimeServe> {
+            cmd.arg("--shutdown-addr=127.0.0.1:0");
             cmd.stdin(Stdio::null());
             cmd.stdout(Stdio::piped());
             cmd.stderr(Stdio::piped());
             let mut child = cmd.spawn()?;
 
-            // Read the first line of stderr which will say which address it's
-            // listening on.
-            //
-            // NB: this intentionally discards any extra buffered data in the
-            // `BufReader` once the newline is found. The server shouldn't print
-            // anything interesting other than the address so once we get a line
-            // all remaining output is left to be captured by future requests
-            // send to the server.
+            // Read the first few lines of stderr which will say which address
+            // it's listening on. The first line is the shutdown line (with
+            // `--shutdown-addr`) and the second is what `--addr` was bound to.
+            // This is done to figure out what `:0` was bound to in the child
+            // process.
             let mut line = String::new();
             let mut reader = BufReader::new(child.stderr.take().unwrap());
-            reader.read_line(&mut line)?;
-
-            match line.find("127.0.0.1").and_then(|addr_start| {
-                let addr = &line[addr_start..];
-                let addr_end = addr.find("/")?;
-                addr[..addr_end].parse().ok()
-            }) {
-                Some(addr) => {
-                    assert!(reader.buffer().is_empty());
-                    child.stderr = Some(reader.into_inner());
-                    Ok(WasmtimeServe {
-                        child: Some(child),
-                        addr,
-                    })
+            let mut read_addr_from_line = |prefix: &str| -> Result<SocketAddr> {
+                reader.read_line(&mut line)?;
+
+                if !line.starts_with(prefix) {
+                    bail!("input line `{line}` didn't start with `{prefix}`");
+                }
+                match line.find("127.0.0.1").and_then(|addr_start| {
+                    let addr = &line[addr_start..];
+                    let addr_end = addr.find("/")?;
+                    addr[..addr_end].parse().ok()
+                }) {
+                    Some(addr) => {
+                        line.truncate(0);
+                        Ok(addr)
+                    }
+                    None => bail!("failed to address from: {line}"),
                 }
-                None => {
+            };
+            let shutdown_addr = read_addr_from_line("Listening for shutdown");
+            let addr = read_addr_from_line("Serving HTTP on");
+            let (shutdown_addr, addr) = match (shutdown_addr, addr) {
+                (Ok(a), Ok(b)) => (a, b),
+                // If either failed kill the child and otherwise try to shepherd
+                // along any contextual information we have.
+                (Err(a), _) | (_, Err(a)) => {
                     child.kill()?;
                     child.wait()?;
                     reader.read_to_string(&mut line)?;
-                    bail!("failed to start child: {line}")
+                    return Err(a.context(line));
                 }
-            }
+            };
+            assert!(reader.buffer().is_empty());
+            child.stderr = Some(reader.into_inner());
+            Ok(WasmtimeServe {
+                child: Some(child),
+                addr,
+                shutdown_addr,
+            })
         }
 
         /// Completes this server gracefully by printing the output on failure.
         fn finish(mut self) -> Result<(String, String)> {
             let mut child = self.child.take().unwrap();
 
-            // If the child process has already exited then collect the output
-            // and test if it succeeded. Otherwise it's still running so kill it
-            // and then reap it. Assume that if it's still running then the test
-            // has otherwise passed so no need to print the output.
-            let known_failure = if child.try_wait()?.is_some() {
-                false
-            } else {
-                child.kill()?;
-                true
-            };
+            // If the child process has already exited, then great! Otherwise
+            // the server is still running so send it a graceful ctrl-c shutdown
+            // signal. If this platform doesn't support graceful ctrl-c then we
+            // can't expect the process to exit successfully.
+            if child.try_wait()?.is_none() {
+                std::net::TcpStream::connect(&self.shutdown_addr)
+                    .context("failed to initiate graceful shutdown")?;
+            }
             let output = child.wait_with_output()?;
-            if !known_failure && !output.status.success() {
+            if !output.status.success() {
                 bail!("child failed {output:?}");
             }