Support setting model thread stack size

src/rt/mod.rs

@@ -64,18 +64,21 @@ pub const MAX_THREADS: usize = 4;
 /// Maximum number of atomic store history to track per-cell.
 pub(crate) const MAX_ATOMIC_HISTORY: usize = 7;
 
-pub(crate) fn spawn<F>(f: F) -> crate::rt::thread::Id
+pub(crate) fn spawn<F>(stack_size: Option<usize>, f: F) -> crate::rt::thread::Id
 where
     F: FnOnce() + 'static,
 {
     let id = execution(|execution| execution.new_thread());
 
     trace!(thread = ?id, "spawn");
 
-    Scheduler::spawn(Box::new(move || {
-        f();
-        thread_done();
-    }));
+    Scheduler::spawn(
+        stack_size,
+        Box::new(move || {
+            f();
+            thread_done();
+        }),
+    );
 
     id
 }

src/rt/rwlock.rs

@@ -171,7 +171,7 @@ impl RwLock {
 
     fn post_acquire_read_lock(&self) -> bool {
         super::execution(|execution| {
-            let mut state = self.state.get_mut(&mut execution.objects);
+            let state = self.state.get_mut(&mut execution.objects);
             let thread_id = execution.threads.active_id();
 
             // Set the lock to the current thread

src/rt/scheduler.rs

@@ -1,20 +1,14 @@
 #![allow(deprecated)]
 
-use crate::rt::{thread, Execution};
+use crate::rt::Execution;
 
 use generator::{self, Generator, Gn};
 use scoped_tls::scoped_thread_local;
 use std::cell::RefCell;
 use std::collections::VecDeque;
-use std::fmt;
 
 pub(crate) struct Scheduler {
-    /// Threads
-    threads: Vec<Thread>,
-
-    next_thread: usize,
-
-    queued_spawn: VecDeque<Box<dyn FnOnce()>>,
+    max_threads: usize,
 }
 
 type Thread = Generator<'static, Option<Box<dyn FnOnce()>>, ()>;
@@ -23,20 +17,21 @@ scoped_thread_local! {
     static STATE: RefCell<State<'_>>
 }
 
+struct QueuedSpawn {
+    f: Box<dyn FnOnce()>,
+    stack_size: Option<usize>,
+}
+
 struct State<'a> {
     execution: &'a mut Execution,
-    queued_spawn: &'a mut VecDeque<Box<dyn FnOnce()>>,
+    queued_spawn: &'a mut VecDeque<QueuedSpawn>,
 }
 
 impl Scheduler {
     /// Create an execution
     pub(crate) fn new(capacity: usize) -> Scheduler {
-        let threads = spawn_threads(capacity);
-
         Scheduler {
-            threads,
-            next_thread: 0,
-            queued_spawn: VecDeque::new(),
+            max_threads: capacity,
         }
     }
 
@@ -75,17 +70,17 @@ impl Scheduler {
         assert!(switch.poll(&mut cx).is_ready());
     }
 
-    pub(crate) fn spawn(f: Box<dyn FnOnce()>) {
-        Self::with_state(|state| state.queued_spawn.push_back(f));
+    pub(crate) fn spawn(stack_size: Option<usize>, f: Box<dyn FnOnce()>) {
+        Self::with_state(|state| state.queued_spawn.push_back(QueuedSpawn { stack_size, f }));
     }
 
     pub(crate) fn run<F>(&mut self, execution: &mut Execution, f: F)
     where
         F: FnOnce() + Send + 'static,
     {
-        self.next_thread = 1;
-        self.threads[0].set_para(Some(Box::new(f)));
-        self.threads[0].resume();
+        let mut threads = Vec::new();
+        threads.push(spawn_thread(Box::new(f), None));
+        threads[0].resume();
 
         loop {
             if !execution.threads.is_active() {
@@ -94,29 +89,31 @@ impl Scheduler {
 
             let active = execution.threads.active_id();
 
-            self.tick(active, execution);
+            let mut queued_spawn = Self::tick(&mut threads[active.as_usize()], execution);
+
+            while let Some(th) = queued_spawn.pop_front() {
+                assert!(threads.len() < self.max_threads);
 
-            while let Some(th) = self.queued_spawn.pop_front() {
-                let thread_id = self.next_thread;
-                self.next_thread += 1;
+                let thread_id = threads.len();
+                let QueuedSpawn { f, stack_size } = th;
 
-                self.threads[thread_id].set_para(Some(th));
-                self.threads[thread_id].resume();
+                threads.push(spawn_thread(f, stack_size));
+                threads[thread_id].resume();
             }
         }
     }
 
-    fn tick(&mut self, thread: thread::Id, execution: &mut Execution) {
+    fn tick(thread: &mut Thread, execution: &mut Execution) -> VecDeque<QueuedSpawn> {
+        let mut queued_spawn = VecDeque::new();
         let state = RefCell::new(State {
             execution,
-            queued_spawn: &mut self.queued_spawn,
+            queued_spawn: &mut queued_spawn,
         });
 
-        let threads = &mut self.threads;
-
         STATE.set(unsafe { transmute_lt(&state) }, || {
-            threads[thread.as_usize()].resume();
+            thread.resume();
         });
+        queued_spawn
     }
 
     fn with_state<F, R>(f: F) -> R
@@ -131,30 +128,23 @@ impl Scheduler {
     }
 }
 
-impl fmt::Debug for Scheduler {
-    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
-        fmt.debug_struct("Schedule")
-            .field("threads", &self.threads)
-            .finish()
-    }
-}
+fn spawn_thread(f: Box<dyn FnOnce()>, stack_size: Option<usize>) -> Thread {
+    let body = move || {
+        loop {
+            let f: Option<Box<dyn FnOnce()>> = generator::yield_(()).unwrap();
+            generator::yield_with(());
+            f.unwrap()();
+        }
 
-fn spawn_threads(n: usize) -> Vec<Thread> {
-    (0..n)
-        .map(|_| {
-            let mut g = Gn::new(move || {
-                loop {
-                    let f: Option<Box<dyn FnOnce()>> = generator::yield_(()).unwrap();
-                    generator::yield_with(());
-                    f.unwrap()();
-                }
-
-                // done!();
-            });
-            g.resume();
-            g
-        })
-        .collect()
+        // done!();
+    };
+    let mut g = match stack_size {
+        Some(stack_size) => Gn::new_opt(stack_size, body),
+        None => Gn::new(body),
+    };
+    g.resume();
+    g.set_para(Some(f));
+    g
 }
 
 unsafe fn transmute_lt<'a, 'b>(state: &'a RefCell<State<'b>>) -> &'a RefCell<State<'static>> {

src/thread.rs

@@ -82,6 +82,7 @@ pub struct LocalKey<T> {
 #[derive(Debug)]
 pub struct Builder {
     name: Option<String>,
+    stack_size: Option<usize>,
 }
 
 static CURRENT_THREAD_KEY: LocalKey<Thread> = LocalKey {
@@ -128,7 +129,7 @@ where
     F: 'static,
     T: 'static,
 {
-    spawn_internal(f, None, location!())
+    spawn_internal(f, None, None, location!())
 }
 
 /// Mock implementation of `std::thread::park`.
@@ -142,7 +143,12 @@ pub fn park() {
     rt::park(location!());
 }
 
-fn spawn_internal<F, T>(f: F, name: Option<String>, location: Location) -> JoinHandle<T>
+fn spawn_internal<F, T>(
+    f: F,
+    name: Option<String>,
+    stack_size: Option<usize>,
+    location: Location,
+) -> JoinHandle<T>
 where
     F: FnOnce() -> T,
     F: 'static,
@@ -154,7 +160,7 @@ where
     let id = {
         let name = name.clone();
         let result = result.clone();
-        rt::spawn(move || {
+        rt::spawn(stack_size, move || {
             rt::execution(|execution| {
                 init_current(execution, name);
             });
@@ -181,7 +187,10 @@ impl Builder {
     // not either, as it's a mock version of the `std` type.
     #[allow(clippy::new_without_default)]
     pub fn new() -> Builder {
-        Builder { name: None }
+        Builder {
+            name: None,
+            stack_size: None,
+        }
     }
 
     /// Names the thread-to-be. Currently the name is used for identification
@@ -193,7 +202,9 @@ impl Builder {
     }
 
     /// Sets the size of the stack (in bytes) for the new thread.
-    pub fn stack_size(self, _size: usize) -> Builder {
+    pub fn stack_size(mut self, size: usize) -> Builder {
+        self.stack_size = Some(size);
+
         self
     }
 
@@ -206,7 +217,7 @@ impl Builder {
         F: Send + 'static,
         T: Send + 'static,
     {
-        Ok(spawn_internal(f, self.name, location!()))
+        Ok(spawn_internal(f, self.name, self.stack_size, location!()))
     }
 }
 

tests/thread_api.rs

@@ -104,6 +104,26 @@ fn thread_names() {
     })
 }
 
+#[test]
+fn thread_stack_size() {
+    const STACK_SIZE: usize = 1 << 16;
+    loom::model(|| {
+        let body = || {
+            // Allocate a large array on the stack.
+            std::hint::black_box(&mut [0usize; STACK_SIZE]);
+        };
+        thread::Builder::new()
+            .stack_size(
+                // Include space for function calls in addition to the array.
+                2 * STACK_SIZE,
+            )
+            .spawn(body)
+            .unwrap()
+            .join()
+            .unwrap()
+    })
+}
+
 #[test]
 fn park_unpark_loom() {
     loom::model(|| {