Fix Windows process waiting and avoid slab allocator shutdown deadlock

dev_tests/src/ratchet.rs

@@ -43,7 +43,7 @@ fn ratchet_globals() -> Result<()> {
             ("litebox_platform_linux_userland/", 5),
             ("litebox_platform_lvbs/", 22),
             ("litebox_platform_multiplex/", 1),
-            ("litebox_platform_windows_userland/", 8),
+            ("litebox_platform_windows_userland/", 7),
             ("litebox_runner_lvbs/", 6),
             ("litebox_runner_snp/", 2),
             ("litebox_shim_linux/", 2),

litebox_platform_windows_userland/src/lib.rs

@@ -1866,40 +1866,6 @@ impl<const ALIGN: usize> litebox::platform::PageManagementProvider<ALIGN> for Wi
     }
 }
 
-#[global_allocator]
-static SLAB_ALLOC: litebox::mm::allocator::SafeZoneAllocator<'static, 28, WindowsUserland> =
-    litebox::mm::allocator::SafeZoneAllocator::new();
-
-impl litebox::mm::allocator::MemoryProvider for WindowsUserland {
-    fn alloc(layout: &std::alloc::Layout) -> Option<(usize, usize)> {
-        let size = core::cmp::max(
-            layout.size().next_power_of_two(),
-            // Note `mmap` provides no guarantee of alignment, so we double the size to ensure we
-            // can always find a required chunk within the returned memory region.
-            core::cmp::max(layout.align(), 0x1000) << 1,
-        );
-
-        match unsafe {
-            VirtualAlloc2(
-                GetCurrentProcess(),
-                core::ptr::null_mut(),
-                size,
-                Win32_Memory::MEM_COMMIT | Win32_Memory::MEM_RESERVE,
-                Win32_Memory::PAGE_READWRITE,
-                core::ptr::null_mut(),
-                0,
-            )
-        } {
-            addr if addr.is_null() => None,
-            addr => Some((addr as usize, size)),
-        }
-    }
-
-    unsafe fn free(_addr: usize) {
-        unimplemented!("Memory deallocation is not implemented for Windows yet.");
-    }
-}
-
 unsafe extern "C" {
     // Defined in asm blocks above
     fn syscall_callback() -> isize;

litebox_shim_linux/src/syscalls/mm.rs

@@ -2011,36 +2011,40 @@ mod tests {
         let task = init_platform();
         let platform = task.global.platform;
         let mut data = alloc::vec::Vec::new();
-        // Find an address that is allocated to the global allocator but not in reserved regions.
-        // LiteBox's page manager is not aware of the global allocator's allocations.
+        let mut count = 0;
+        // Model an external allocator allocation that LiteBox's page manager does not track.
         let addr = loop {
-            #[allow(
-                unused_variables,
-                reason = "the following features are mutually exclusive"
-            )]
-            #[cfg(target_os = "windows")]
-            let addr = {
-                let buf = alloc::vec::Vec::<u8>::with_capacity(0x10_0000);
-                let addr = buf.as_ptr() as usize;
-                data.push(buf);
-                addr
-            };
-            #[cfg(target_os = "linux")]
+            assert!(
+                count < 100,
+                "Failed to find a suitable address after 100 attempts"
+            );
+            count += 1;
             let addr = {
-                let addr = unsafe {
-                    libc::mmap(
-                        core::ptr::null_mut(),
-                        0x10_000,
-                        libc::PROT_READ | libc::PROT_WRITE,
-                        libc::MAP_PRIVATE | libc::MAP_ANONYMOUS,
-                        -1,
-                        0,
+                use litebox::platform::{
+                    RawConstPointer as _,
+                    page_mgmt::{FixedAddressBehavior, MemoryRegionPermissions},
+                };
+
+                let allocation = <Platform as PageManagementProvider<4096>>::allocate_pages(
+                    platform,
+                    0..0x2000,
+                    MemoryRegionPermissions::READ | MemoryRegionPermissions::WRITE,
+                    false,
+                    false,
+                    FixedAddressBehavior::Hint,
+                )
+                .unwrap()
+                .as_usize();
+                // SAFETY: The first page belongs to this test and has no outstanding references.
+                unsafe {
+                    <Platform as PageManagementProvider<4096>>::deallocate_pages(
+                        platform,
+                        allocation..allocation + 0x1000,
                     )
-                } as usize;
-                data.push(alloc::vec::Vec::<u8>::from(unsafe {
-                    core::slice::from_raw_parts(addr as *const u8, 0x10_000)
-                }));
-                addr
+                    .unwrap();
+                }
+                data.push(allocation);
+                allocation + 0x1000
             };
 
             let mut included = false;
@@ -2098,6 +2102,20 @@ mod tests {
             )
             .unwrap_err();
         assert_eq!(err, Errno::ENOMEM);
+
+        task.sys_munmap(res, 0x1000).unwrap();
+        task.sys_munmap(UserPtrMut::from_usize(addr - 0x1000), 0x1000)
+            .unwrap();
+        for allocation in data {
+            // SAFETY: The remaining page belongs to this test and was never mapped by the shim.
+            unsafe {
+                <Platform as PageManagementProvider<4096>>::deallocate_pages(
+                    platform,
+                    allocation + 0x1000..allocation + 0x2000,
+                )
+                .unwrap();
+            }
+        }
     }
 
     #[test]

litebox_shim_windows/src/lib.rs

@@ -16,7 +16,7 @@ use alloc::collections::BTreeMap;
 use alloc::sync::Arc;
 use alloc::vec::Vec;
 use core::marker::PhantomData;
-use core::sync::atomic::{AtomicI32, AtomicU32, AtomicUsize, Ordering};
+use core::sync::atomic::{AtomicU32, AtomicUsize, Ordering};
 use litebox_common_windows::nt_status::NtStatus;
 
 use litebox::LiteBox;
@@ -30,6 +30,7 @@ use litebox::sync::{Mutex, RawSyncPrimitivesProvider};
 use litebox::utils::TruncateExt as _;
 use litebox_common_windows::loader::PAGE_SIZE;
 use litebox_common_windows::{NtSysno, Win32Sysno};
+use litebox_platform::sync::{RawMutex as _, RawMutexProvider};
 use litebox_platform::time::TimeProvider;
 
 use crate::syscalls::event::{EventHandleObject, EventSubsystem};
@@ -593,16 +594,18 @@ pub struct Process<Platform: ShimPlatform> {
     trace_notifications: Mutex<Platform, syscalls::trace::TraceNotifications<Platform>>,
     gdi_state: Mutex<Platform, Option<syscalls::gdi::GdiProcessState>>,
     cookie: u32,
-    exit_code: AtomicI32,
     next_thread_id: AtomicUsize,
+    nr_threads: <Platform as RawMutexProvider>::RawMutex,
     threads: litebox::sync::RwLock<Platform, ProcessThreads<Platform>>,
 }
 
-/// The live threads of a process, and whether the process is tearing down.
+/// The live threads and exit state of a process.
 struct ProcessThreads<Platform: ShimPlatform> {
     /// Set once the process has started exiting. No further threads may be
     /// created, and the exit code is frozen.
     group_exit: bool,
+    /// The process exit code, updated by thread exits until group exit.
+    exit_code: i32,
     /// The thread objects of every thread that has not yet completed, keyed by
     /// thread ID.
     threads: BTreeMap<usize, Arc<syscalls::thread::ThreadObject<Platform>>>,
@@ -630,12 +633,33 @@ impl<Platform: ShimPlatform> Process<Platform> {
         }
         let previous = threads.threads.insert(thread_id, thread.clone());
         debug_assert!(previous.is_none(), "thread ID {thread_id} already exists");
+        let nr_threads = self.nr_threads.underlying_atomic();
+        nr_threads.store(nr_threads.load(Ordering::Relaxed) + 1, Ordering::Release);
         true
     }
 
     /// Unregisters a thread that failed to start or has completed.
     fn detach_thread(&self, thread_id: usize) {
-        self.threads.write().threads.remove(&thread_id);
+        let notify = {
+            let mut threads = self.threads.write();
+            threads.threads.remove(&thread_id);
+
+            let nr_threads = self.nr_threads.underlying_atomic();
+            let count = nr_threads.load(Ordering::Relaxed);
+            let new_count = count
+                .checked_sub(1)
+                .expect("decrementing from zero threads");
+            nr_threads.store(new_count, Ordering::Release);
+            if new_count == 0 {
+                debug_assert!(threads.threads.is_empty());
+                // The last thread exited. Prevent new threads.
+                threads.group_exit = true;
+            }
+            new_count == 0
+        };
+        if notify {
+            self.nr_threads.wake_all();
+        }
     }
 
     /// Returns the number of threads that have not yet completed.
@@ -667,14 +691,17 @@ impl<Platform: ShimPlatform> Process<Platform> {
         self.threads.read().threads.get(&thread_id).cloned()
     }
 
-    /// Wait for the process to exit, returning its exit code.
-    ///
-    /// Currently a placeholder that returns a fixed exit code immediately.
-    /// Once NT process lifecycle exists, this will actually block.
+    /// Waits for all guest threads in the process to complete, returning its exit code.
     #[must_use]
     pub fn wait(&self) -> i32 {
-        // TODO: Wait for the NT process object once process lifecycle exists.
-        self.exit_code.load(Ordering::Relaxed)
+        loop {
+            let remaining = self.nr_threads.underlying_atomic().load(Ordering::Acquire);
+            if remaining == 0 {
+                break;
+            }
+            let _ = self.nr_threads.block(remaining);
+        }
+        self.threads.read().exit_code
     }
 
     fn default(
@@ -714,10 +741,11 @@ impl<Platform: ShimPlatform> Process<Platform> {
             trace_notifications: Mutex::new(syscalls::trace::TraceNotifications::default()),
             gdi_state: Mutex::new(None),
             cookie: syscalls::process::default_process_cookie(),
-            exit_code: AtomicI32::new(DEFAULT_PROCESS_EXIT_CODE),
             next_thread_id: AtomicUsize::new(syscalls::process::INITIAL_THREAD_ID + 1),
+            nr_threads: <Platform as RawMutexProvider>::RawMutex::INIT,
             threads: litebox::sync::RwLock::new(ProcessThreads {
                 group_exit: false,
+                exit_code: DEFAULT_PROCESS_EXIT_CODE,
                 threads: BTreeMap::new(),
             }),
         }
@@ -748,8 +776,13 @@ impl<Platform: ShimPlatform> Task<Platform> {
         });
     }
 
-    /// Marks the current thread as exiting with `exit_status`.
+    /// Updates the process exit status and marks the current thread as exiting.
     fn exit_thread(&self, exit_status: i32) {
+        let mut threads = self.process.threads.write();
+        if self.thread_object.is_exiting() {
+            return;
+        }
+        threads.exit_code = exit_status;
         self.thread_object.exit_thread(exit_status);
     }
 
@@ -761,7 +794,7 @@ impl<Platform: ShimPlatform> Task<Platform> {
             return;
         }
         threads.group_exit = true;
-        self.process.exit_code.store(exit_status, Ordering::Relaxed);
+        threads.exit_code = exit_status;
         // Interrupting the caller is a no-op, because it is running in the
         // host, so this does not need to single it out.
         for thread in threads.threads.values() {