DProxy

lib/wasix/Cargo.toml

@@ -49,6 +49,7 @@ tokio = { version = "1", features = [
 ], default_features = false }
 tokio-stream = { version = "0.1", features = [ "sync" ] }
 futures = { version = "0.3" }
+async-ctrlc = { version = "1.2.0", features = [ "stream" ], optional = true }
 # used by feature='os'
 async-trait = { version = "^0.1" }
 urlencoding = { version = "^2" }
@@ -141,6 +142,7 @@ wasmer = { path = "../api", version = "=4.2.5", default-features = false, featur
 default = ["sys-default"]
 
 time = ["tokio/time"]
+ctrlc = ["dep:async-ctrlc"]
 
 webc_runner_rt_wcgi = ["hyper", "wcgi", "wcgi-host", "tower", "tower-http"]
 webc_runner_rt_dcgi = ["webc_runner_rt_wcgi", "journal"]
@@ -158,6 +160,7 @@ sys-default = [
     "host-vnet",
     "host-threads",
     "host-reqwest",
+    "ctrlc"
 ]
 sys-poll = []
 sys-thread = ["tokio/rt", "tokio/time", "tokio/rt-multi-thread", "rusty_pool"]

lib/wasix/src/os/task/process.rs

@@ -33,6 +33,7 @@ use super::{
     signal::{SignalDeliveryError, SignalHandlerAbi},
     task_join_handle::OwnedTaskStatus,
     thread::WasiMemoryLayout,
+    TaskStatus,
 };
 
 /// Represents the ID of a sub-process
@@ -147,6 +148,8 @@ impl Into<Range<u64>> for MemorySnapshotRegion {
 pub struct WasiProcessInner {
     /// Unique ID of this process
     pub pid: WasiProcessId,
+    /// Number of threads waiting for children to exit
+    pub(crate) waiting: Arc<AtomicU32>,
     /// The threads that make up this process
     pub threads: HashMap<WasiThreadId, WasiThread>,
     /// Number of threads running for this process
@@ -158,6 +161,9 @@ pub struct WasiProcessInner {
     /// Represents a checkpoint which blocks all the threads
     /// and then executes some maintenance action
     pub checkpoint: WasiProcessCheckpoint,
+    /// Flag that indicates that the process should snapshot
+    /// itself and exit when the ctrl-c key is pressed
+    pub snapshot_on_sigint: bool,
     /// Any wakers waiting on this process (for example for a checkpoint)
     pub wakers: Vec<Waker>,
     /// The snapshot memory significantly reduce the amount of
@@ -331,26 +337,47 @@ impl Drop for WasiProcessWait {
 
 impl WasiProcess {
     pub fn new(pid: WasiProcessId, module_hash: ModuleHash, plane: WasiControlPlaneHandle) -> Self {
+        let waiting = Arc::new(AtomicU32::new(0));
+        let inner = Arc::new((
+            Mutex::new(WasiProcessInner {
+                pid,
+                threads: Default::default(),
+                thread_count: Default::default(),
+                signal_intervals: Default::default(),
+                children: Default::default(),
+                checkpoint: WasiProcessCheckpoint::Execute,
+                wakers: Default::default(),
+                waiting: waiting.clone(),
+                snapshot_on_sigint: false,
+                snapshot_memory_hash: Default::default(),
+            }),
+            Condvar::new(),
+        ));
+
+        #[derive(Debug)]
+        struct SignalHandler(LockableWasiProcessInner);
+        impl SignalHandlerAbi for SignalHandler {
+            fn signal(&self, signal: u8) -> Result<(), SignalDeliveryError> {
+                if let Ok(signal) = signal.try_into() {
+                    signal_process_internal(&self.0, signal);
+                    Ok(())
+                } else {
+                    Err(SignalDeliveryError)
+                }
+            }
+        }
+
         WasiProcess {
             pid,
             module_hash,
             parent: None,
             compute: plane,
-            inner: Arc::new((
-                Mutex::new(WasiProcessInner {
-                    pid,
-                    threads: Default::default(),
-                    thread_count: Default::default(),
-                    signal_intervals: Default::default(),
-                    children: Default::default(),
-                    checkpoint: WasiProcessCheckpoint::Execute,
-                    wakers: Default::default(),
-                    snapshot_memory_hash: Default::default(),
-                }),
-                Condvar::new(),
-            )),
-            finished: Arc::new(OwnedTaskStatus::default()),
-            waiting: Arc::new(AtomicU32::new(0)),
+            inner: inner.clone(),
+            finished: Arc::new(
+                OwnedTaskStatus::new(TaskStatus::Pending)
+                    .with_signal_handler(Arc::new(SignalHandler(inner))),
+            ),
+            waiting,
         }
     }
 
@@ -471,26 +498,7 @@ impl WasiProcess {
 
     /// Signals all the threads in this process
     pub fn signal_process(&self, signal: Signal) {
-        let pid = self.pid();
-        tracing::trace!(%pid, "signal-process({:?})", signal);
-
-        {
-            let inner = self.inner.0.lock().unwrap();
-            if self.waiting.load(Ordering::Acquire) > 0 {
-                let mut triggered = false;
-                for child in inner.children.iter() {
-                    child.signal_process(signal);
-                    triggered = true;
-                }
-                if triggered {
-                    return;
-                }
-            }
-        }
-        let inner = self.inner.0.lock().unwrap();
-        for thread in inner.threads.values() {
-            thread.signal(signal);
-        }
+        signal_process_internal(&self.inner, signal);
     }
 
     /// Signals one of the threads every interval
@@ -606,6 +614,47 @@ impl WasiProcess {
     }
 }
 
+/// Signals all the threads in this process
+fn signal_process_internal(process: &LockableWasiProcessInner, signal: Signal) {
+    let mut inner = process.0.lock().unwrap();
+    let pid = inner.pid;
+    tracing::trace!(%pid, "signal-process({:?})", signal);
+
+    // If the snapshot on ctrl-c is currently registered then we need
+    // to take a snapshot and exit
+    if inner.snapshot_on_sigint {
+        tracing::debug!(%pid, "snapshot-on-interrupt-signal");
+
+        // Initiate the checksum
+        inner.checkpoint = WasiProcessCheckpoint::Snapshot {
+            trigger: SnapshotTrigger::Sigint,
+        };
+        for waker in inner.wakers.drain(..) {
+            waker.wake();
+        }
+        process.1.notify_all();
+        return;
+    }
+
+    // Check if there are subprocesses that will receive this signal
+    // instead of this process
+    if inner.waiting.load(Ordering::Acquire) > 0 {
+        let mut triggered = false;
+        for child in inner.children.iter() {
+            child.signal_process(signal);
+            triggered = true;
+        }
+        if triggered {
+            return;
+        }
+    }
+
+    // Otherwise just send the signal to all the threads
+    for thread in inner.threads.values() {
+        thread.signal(signal);
+    }
+}
+
 impl SignalHandlerAbi for WasiProcess {
     fn signal(&self, sig: u8) -> Result<(), SignalDeliveryError> {
         if let Ok(sig) = sig.try_into() {

lib/wasix/src/os/task/signal.rs

@@ -1,4 +1,4 @@
-use std::time::Duration;
+use std::{sync::Arc, time::Duration};
 
 use wasmer_wasix_types::types::Signal;
 
@@ -15,6 +15,8 @@ where
     fn signal(&self, signal: u8) -> Result<(), SignalDeliveryError>;
 }
 
+pub type DynSignalHandlerAbi = dyn SignalHandlerAbi + Send + Sync + 'static;
+
 #[derive(Debug)]
 pub struct WasiSignalInterval {
     /// Signal that will be raised
@@ -26,3 +28,30 @@ pub struct WasiSignalInterval {
     /// Last time that a signal was triggered
     pub last_signal: u128,
 }
+
+pub fn default_signal_handler() -> Arc<DynSignalHandlerAbi> {
+    #[derive(Debug)]
+    struct DefaultHandler {}
+    impl SignalHandlerAbi for DefaultHandler {
+        fn signal(&self, signal: u8) -> Result<(), SignalDeliveryError> {
+            if let Ok(signal) = TryInto::<Signal>::try_into(signal) {
+                match signal {
+                    Signal::Sigkill
+                    | Signal::Sigterm
+                    | Signal::Sigabrt
+                    | Signal::Sigquit
+                    | Signal::Sigint
+                    | Signal::Sigstop => {
+                        tracing::debug!("handling terminate signal");
+                        std::process::exit(1);
+                    }
+                    signal => tracing::info!("unhandled signal - {:?}", signal),
+                }
+            } else {
+                tracing::info!("unknown signal - {}", signal)
+            }
+            Ok(())
+        }
+    }
+    Arc::new(DefaultHandler {})
+}

lib/wasix/src/os/task/task_join_handle.rs

@@ -8,6 +8,8 @@ use wasmer_wasix_types::wasi::{Errno, ExitCode};
 
 use crate::WasiRuntimeError;
 
+use super::signal::{default_signal_handler, DynSignalHandlerAbi};
+
 #[derive(Clone, Debug)]
 pub enum TaskStatus {
     Pending,
@@ -70,6 +72,9 @@ pub trait VirtualTaskHandle: std::fmt::Debug + Send + Sync + 'static {
 /// A handle that allows awaiting the termination of a task, and retrieving its exit code.
 #[derive(Debug)]
 pub struct OwnedTaskStatus {
+    // The signal handler that can be invoked for this owned task
+    signal_handler: Arc<DynSignalHandlerAbi>,
+
     watch_tx: tokio::sync::watch::Sender<TaskStatus>,
     // Even through unused, without this receive there is a race condition
     // where the previously sent values are lost.
@@ -81,11 +86,23 @@ impl OwnedTaskStatus {
     pub fn new(status: TaskStatus) -> Self {
         let (tx, rx) = tokio::sync::watch::channel(status);
         Self {
+            signal_handler: default_signal_handler(),
             watch_tx: tx,
             watch_rx: rx,
         }
     }
 
+    /// Sets the signal handler used for this owned task
+    pub fn set_signal_handler(&mut self, handler: Arc<DynSignalHandlerAbi>) {
+        self.signal_handler = handler;
+    }
+
+    /// Attaches a signal handler
+    pub fn with_signal_handler(mut self, handler: Arc<DynSignalHandlerAbi>) -> Self {
+        self.set_signal_handler(handler);
+        self
+    }
+
     pub fn new_finished_with_code(code: ExitCode) -> Self {
         Self::new(TaskStatus::Finished(Ok(code)))
     }
@@ -144,6 +161,7 @@ impl OwnedTaskStatus {
 
     pub fn handle(&self) -> TaskJoinHandle {
         TaskJoinHandle {
+            signal_handler: self.signal_handler.clone(),
             watch: self.watch_tx.subscribe(),
         }
     }
@@ -158,6 +176,7 @@ impl Default for OwnedTaskStatus {
 /// A handle that allows awaiting the termination of a task, and retrieving its exit code.
 #[derive(Clone, Debug)]
 pub struct TaskJoinHandle {
+    signal_handler: Arc<DynSignalHandlerAbi>,
     watch: tokio::sync::watch::Receiver<TaskStatus>,
 }
 
@@ -167,6 +186,29 @@ impl TaskJoinHandle {
         self.watch.borrow().clone()
     }
 
+    #[cfg(feature = "ctrlc")]
+    pub fn install_ctrlc_handler(&self) {
+        use wasmer::FromToNativeWasmType;
+        use wasmer_wasix_types::wasi::Signal;
+
+        let signal_handler = self.signal_handler.clone();
+
+        let mut ctrlc = async_ctrlc::CtrlC::new().expect("cannot create Ctrl+C handler");
+        let task_handle = self.clone();
+        tokio::spawn(async move {
+            // Loop sending ctrl-c presses as signals to the signal handler
+            loop {
+                let ctrlc_pin = &mut ctrlc;
+                ctrlc_pin.await;
+
+                if let Err(err) = signal_handler.signal(Signal::Sigint.to_native() as u8) {
+                    tracing::error!("failed to process signal - {}", err);
+                    std::process::exit(1);
+                }
+            }
+        });
+    }
+
     /// Wait until the task finishes.
     pub async fn wait_finished(&mut self) -> Result<ExitCode, Arc<WasiRuntimeError>> {
         loop {

lib/wasix/src/runners/wasi.rs

@@ -326,6 +326,9 @@ impl crate::runners::Runner for WasiRunner {
                         .await
                         .context("Spawn failed")?;
 
+                #[cfg(feature = "ctrlc")]
+                task_handle.install_ctrlc_handler();
+
                 task_handle
                     .wait_finished()
                     .await