Support setting model thread stack size

Instead of creating a generator for each possible thread at the
beginning, create threads and backing generators as they are spawned by
the test cases. This allows us to provide a stack size when spawning a
scheduler thread, which can be set via the thread Builder.

Fixes #309

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/scheduler.rs

@@ -17,9 +17,14 @@ 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 {
@@ -65,17 +70,16 @@ 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,
     {
-        let mut next_thread = 1;
-        let mut threads = spawn_threads(self.max_threads);
-        threads[0].set_para(Some(Box::new(f)));
+        let mut threads = Vec::new();
+        threads.push(spawn_thread(Box::new(f), None));
         threads[0].resume();
 
         loop {
@@ -88,16 +92,18 @@ impl Scheduler {
             let mut queued_spawn = Self::tick(&mut threads[active.as_usize()], execution);
 
             while let Some(th) = queued_spawn.pop_front() {
-                let thread_id = next_thread;
-                next_thread += 1;
+                assert!(threads.len() < self.max_threads);
+
+                let thread_id = threads.len();
+                let QueuedSpawn { f, stack_size } = th;
 
-                threads[thread_id].set_para(Some(th));
+                threads.push(spawn_thread(f, stack_size));
                 threads[thread_id].resume();
             }
         }
     }
 
-    fn tick(thread: &mut Thread, execution: &mut Execution) -> VecDeque<Box<dyn FnOnce()>> {
+    fn tick(thread: &mut Thread, execution: &mut Execution) -> VecDeque<QueuedSpawn> {
         let mut queued_spawn = VecDeque::new();
         let state = RefCell::new(State {
             execution,
@@ -122,22 +128,23 @@ impl Scheduler {
     }
 }
 
-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()
+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()();
+        }
+
+        // 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(|| {