ensure fn passed to generator completes (#318)

This avoids forcing the generator to abort, which is not thread-safe.

Author: carllerche

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
     });
 }

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),

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())
     }