Wasix Control Plane - Thread Limit + Cleanup

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;