feat: Implement WASIX ControlPlane Thread Limits

This commit adds a total task (aka thread) limit for a WasiControlPlane.

The limit is configured at startup, and is checked whenever a new
process or a new thread is created.

The major breaking change is that the new_process and new_thread methods
are now fallible and return an error.

Also adds two tests that validate the limits are repsected.

Along with this feature the commit also does some cleanup and reafactoring:

* Simplify WasiControlPlane

It was using multiple Arcs with locks and atomics with no real justification.
The state is only locked for very short periods, so the complicated
setup doesn't yield any benefits.

The mutable state is consolidated behind a single RwLock, which holds the
process id seed and the process map, and the whole type is behind a
single Arc.

The reservations are removed, since they are not really useful if the
state is only locked for very short periods.

Also adds todo notes for handling pid exhaustion and process de-registration.

* TaskJoinHandle

Adds a new TaskJoinHandle handle that is used by both processes and threads
to track termination.

Replaces a complex, cryptic type.

* Simplify Thread state and make it private

Make internal thread state private to better encapsulate the code.
Also restructures the type to be behind a single Arc instead of multiple
redundant Arcs.

NOTE: should also be done for WasiProcess

lib/wasi/src/os/console/mod.rs

@@ -18,9 +18,9 @@ use tokio::sync::{mpsc, RwLock};
 use tracing::{debug, error, info, trace, warn};
 #[cfg(feature = "sys")]
 use wasmer::Engine;
-use wasmer_vbus::{BusSpawnedProcess, SpawnOptionsConfig};
+use wasmer_vbus::{BusSpawnedProcess, SpawnOptionsConfig, VirtualBusError};
 use wasmer_vfs::{FileSystem, RootFileSystemBuilder, SpecialFile, WasiPipe};
-use wasmer_wasi_types::types::__WASI_STDIN_FILENO;
+use wasmer_wasi_types::{types::__WASI_STDIN_FILENO, wasi::BusErrno};
 
 use super::{cconst::ConsoleConst, common::*};
 use crate::{
@@ -161,8 +161,12 @@ impl Console {
 
         // Create the control plane, process and thread
         let control_plane = WasiControlPlane::default();
-        let wasi_process = control_plane.new_process();
-        let wasi_thread = wasi_process.new_thread();
+        let wasi_process = control_plane
+            .new_process()
+            .expect("creating processes on new control planes should always work");
+        let wasi_thread = wasi_process
+            .new_thread()
+            .expect("creating the main thread should always work");
 
         // Create the state
         let mut state = WasiState::builder(prog);

lib/wasi/src/os/task/control_plane.rs

@@ -1,99 +1,206 @@
 use std::{
-    collections::{HashMap, HashSet},
+    collections::HashMap,
     sync::{
-        atomic::{AtomicU32, Ordering},
-        Arc, Mutex, RwLock,
+        atomic::{AtomicUsize, Ordering},
+        Arc, RwLock,
     },
 };
 
-use crate::{os::task::process::WasiProcessInner, WasiProcess, WasiProcessId};
+use crate::{WasiProcess, WasiProcessId};
 
 #[derive(Debug, Clone)]
 pub struct WasiControlPlane {
-    /// The processes running on this machine
-    pub(crate) processes: Arc<RwLock<HashMap<WasiProcessId, WasiProcess>>>,
-    /// Seed used to generate process ID's
-    pub(crate) process_seed: Arc<AtomicU32>,
-    /// Allows for a PID to be reserved
-    pub(crate) reserved: Arc<Mutex<HashSet<WasiProcessId>>>,
+    state: Arc<State>,
 }
 
-impl Default for WasiControlPlane {
-    fn default() -> Self {
+#[derive(Debug, Clone)]
+pub struct ControlPlaneConfig {
+    /// Total number of tasks (processes + threads) that can be spawned.
+    pub max_task_count: Option<usize>,
+}
+
+impl ControlPlaneConfig {
+    pub fn new() -> Self {
         Self {
-            processes: Default::default(),
-            process_seed: Arc::new(AtomicU32::new(0)),
-            reserved: Default::default(),
+            max_task_count: None,
         }
     }
 }
 
+impl Default for ControlPlaneConfig {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+#[derive(Debug)]
+struct State {
+    config: ControlPlaneConfig,
+
+    /// Total number of active tasks (threads) across all processes.
+    task_count: Arc<AtomicUsize>,
+
+    /// Mutable state.
+    mutable: RwLock<MutableState>,
+}
+
+#[derive(Debug)]
+struct MutableState {
+    /// Seed used to generate process ID's
+    process_seed: u32,
+    /// The processes running on this machine
+    processes: HashMap<WasiProcessId, WasiProcess>,
+    // TODO: keep a queue of terminated process ids for id reuse.
+}
+
 impl WasiControlPlane {
-    /// Reserves a PID and returns it
-    pub fn reserve_pid(&self) -> WasiProcessId {
-        let mut pid: WasiProcessId;
-        loop {
-            pid = self.process_seed.fetch_add(1, Ordering::AcqRel).into();
-
-            {
-                let mut guard = self.reserved.lock().unwrap();
-                if guard.contains(&pid) {
-                    continue;
-                }
-                guard.insert(pid);
-            }
+    pub fn new(config: ControlPlaneConfig) -> Self {
+        Self {
+            state: Arc::new(State {
+                config,
+                task_count: Arc::new(AtomicUsize::new(0)),
+                mutable: RwLock::new(MutableState {
+                    process_seed: 0,
+                    processes: Default::default(),
+                }),
+            }),
+        }
+    }
 
-            {
-                let guard = self.processes.read().unwrap();
-                if guard.contains_key(&pid) == false {
-                    break;
-                }
-            }
+    /// Get the current count of active tasks (threads).
+    fn active_task_count(&self) -> usize {
+        self.state.task_count.load(Ordering::SeqCst)
+    }
 
-            {
-                let mut guard = self.reserved.lock().unwrap();
-                guard.remove(&pid);
+    /// Register a new task.
+    ///
+    // Currently just increments the task counter.
+    pub(super) fn register_task(&self) -> Result<TaskCountGuard, ControlPlaneError> {
+        let count = self.state.task_count.fetch_add(1, Ordering::SeqCst);
+        if let Some(max) = self.state.config.max_task_count {
+            if count > max {
+                self.state.task_count.fetch_sub(1, Ordering::SeqCst);
+                return Err(ControlPlaneError::TaskLimitReached { max: count });
             }
         }
-        pid
+        Ok(TaskCountGuard(self.state.task_count.clone()))
     }
 
     /// Creates a new process
-    pub fn new_process(&self) -> WasiProcess {
-        let pid = self.reserve_pid();
-        let ret = WasiProcess {
-            pid,
-            ppid: 0u32.into(),
-            compute: self.clone(),
-            inner: Arc::new(RwLock::new(WasiProcessInner {
-                threads: Default::default(),
-                thread_count: Default::default(),
-                thread_seed: Default::default(),
-                thread_local: Default::default(),
-                thread_local_user_data: Default::default(),
-                thread_local_seed: Default::default(),
-                signal_intervals: Default::default(),
-                bus_processes: Default::default(),
-                bus_process_reuse: Default::default(),
-            })),
-            children: Arc::new(RwLock::new(Default::default())),
-            finished: Arc::new(Mutex::new((None, tokio::sync::broadcast::channel(1).0))),
-            waiting: Arc::new(AtomicU32::new(0)),
-        };
-        {
-            let mut guard = self.processes.write().unwrap();
-            guard.insert(pid, ret.clone());
-        }
-        {
-            let mut guard = self.reserved.lock().unwrap();
-            guard.remove(&pid);
+    // FIXME: De-register terminated processes!
+    // Currently they just accumulate.
+    pub fn new_process(&self) -> Result<WasiProcess, ControlPlaneError> {
+        if let Some(max) = self.state.config.max_task_count {
+            if self.active_task_count() >= max {
+                // NOTE: task count is not incremented here, only when new threads are spawned.
+                // A process will always have a main thread.
+                return Err(ControlPlaneError::TaskLimitReached { max });
+            }
         }
-        ret
+
+        // Create the process first to do all the allocations before locking.
+        let mut proc = WasiProcess::new(WasiProcessId::from(0), self.clone());
+
+        let mut mutable = self.state.mutable.write().unwrap();
+
+        let pid = mutable.next_process_id()?;
+        proc.set_pid(pid);
+        mutable.processes.insert(pid, proc.clone());
+        Ok(proc)
     }
 
     /// Gets a reference to a running process
     pub fn get_process(&self, pid: WasiProcessId) -> Option<WasiProcess> {
-        let guard = self.processes.read().unwrap();
-        guard.get(&pid).map(|a| a.clone())
+        self.state
+            .mutable
+            .read()
+            .unwrap()
+            .processes
+            .get(&pid)
+            .cloned()
+    }
+}
+
+impl MutableState {
+    fn next_process_id(&mut self) -> Result<WasiProcessId, ControlPlaneError> {
+        // TODO: reuse terminated ids, handle wrap-around, ...
+        let id = self.process_seed.checked_add(1).ok_or_else(|| {
+            ControlPlaneError::TaskLimitReached {
+                max: u32::MAX as usize,
+            }
+        })?;
+        self.process_seed = id;
+        Ok(WasiProcessId::from(id))
+    }
+}
+
+impl Default for WasiControlPlane {
+    fn default() -> Self {
+        let config = ControlPlaneConfig::default();
+        Self::new(config)
+    }
+}
+
+/// Guard that ensures the [`WasiControlPlane`] task counter is decremented when dropped.
+#[derive(Debug)]
+pub struct TaskCountGuard(Arc<AtomicUsize>);
+
+impl Drop for TaskCountGuard {
+    fn drop(&mut self) {
+        self.0.fetch_sub(1, Ordering::SeqCst);
+    }
+}
+
+#[derive(thiserror::Error, PartialEq, Eq, Clone, Debug)]
+pub enum ControlPlaneError {
+    /// The maximum number of execution tasks has been reached.
+    #[error("The maximum number of execution tasks has been reached ({max})")]
+    TaskLimitReached {
+        /// The maximum number of tasks.
+        max: usize,
+    },
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// Simple test to ensure task limits are respected.
+    #[test]
+    fn test_control_plane_task_limits() {
+        let p = WasiControlPlane::new(ControlPlaneConfig {
+            max_task_count: Some(2),
+        });
+
+        let p1 = p.new_process().unwrap();
+        let _t1 = p1.new_thread().unwrap();
+        let _t2 = p1.new_thread().unwrap();
+
+        assert_eq!(
+            p.new_process().unwrap_err(),
+            ControlPlaneError::TaskLimitReached { max: 2 }
+        );
+    }
+
+    /// Simple test to ensure task limits are respected and that thread drop guards work.
+    #[test]
+    fn test_control_plane_task_limits_with_dropped_threads() {
+        let p = WasiControlPlane::new(ControlPlaneConfig {
+            max_task_count: Some(2),
+        });
+
+        let p1 = p.new_process().unwrap();
+
+        for _ in 0..10 {
+            let _thread = p1.new_thread().unwrap();
+        }
+
+        let _t1 = p1.new_thread().unwrap();
+        let _t2 = p1.new_thread().unwrap();
+
+        assert_eq!(
+            p.new_process().unwrap_err(),
+            ControlPlaneError::TaskLimitReached { max: 2 }
+        );
     }
 }

lib/wasi/src/os/task/mod.rs

@@ -3,4 +3,5 @@
 pub mod control_plane;
 pub mod process;
 pub mod signal;
+pub mod task_join_handle;
 pub mod thread;