From 43e97ce86adafd8e4e13a69fb62ddfc9099d4549 Mon Sep 17 00:00:00 2001
From: Carl Lerche <me@carllerche.com>
Date: Fri, 21 Jul 2023 12:31:58 -0700
Subject: [PATCH] ensure fn passed to generator completes

This avoids forcing the generator to abort, which is not thread-safe.
---
 src/rt/mod.rs       |  3 ---
 src/rt/scheduler.rs | 19 ++++++++++++++-----
 src/rt/thread.rs    | 17 +++++++++++++++++
 3 files changed, 31 insertions(+), 8 deletions(-)

diff --git a/src/rt/mod.rs b/src/rt/mod.rs
index 0b417fad..b4347539 100644
--- a/src/rt/mod.rs
+++ b/src/rt/mod.rs
@@ -192,9 +192,6 @@ pub fn thread_done() {
         execution.threads.active_mut().operation = None;
         execution.threads.active_mut().set_terminated();
         let switch = execution.schedule();
-
         trace!(?thread, ?switch, "thread_done: terminate");
-
-        switch
     });
 }
diff --git a/src/rt/scheduler.rs b/src/rt/scheduler.rs
index 7ea38d95..09ee196f 100644
--- a/src/rt/scheduler.rs
+++ b/src/rt/scheduler.rs
@@ -83,7 +83,11 @@ impl Scheduler {
         threads[0].resume();
 
         loop {
-            if !execution.threads.is_active() {
+            if execution.threads.is_complete() {
+                for thread in &mut threads {
+                    thread.resume();
+                    assert!(thread.is_done());
+                }
                 return;
             }
 
@@ -131,12 +135,17 @@ impl Scheduler {
 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()();
+            let f: Option<Option<Box<dyn FnOnce()>>> = generator::yield_(());
+
+            if let Some(f) = f {
+                generator::yield_with(());
+                f.unwrap()();
+            } else {
+                break;
+            }
         }
 
-        // done!();
+        generator::done!();
     };
     let mut g = match stack_size {
         Some(stack_size) => Gn::new_opt(stack_size, body),
diff --git a/src/rt/thread.rs b/src/rt/thread.rs
index 626cc552..f325c4d1 100644
--- a/src/rt/thread.rs
+++ b/src/rt/thread.rs
@@ -234,6 +234,23 @@ impl Set {
         self.active.is_some()
     }
 
+    pub(crate) fn is_complete(&self) -> bool {
+        if self.active.is_none() {
+            // All threads should be terminated
+            for thread in &self.threads {
+                assert!(
+                    thread.is_terminated(),
+                    "thread not terminated; {:#?}",
+                    thread
+                );
+            }
+
+            true
+        } else {
+            false
+        }
+    }
+
     pub(crate) fn active_id(&self) -> Id {
         Id::new(self.execution_id, self.active.unwrap())
     }