Fix all sample schedulers to complete the root computation

lib/picos_lwt/picos_lwt.ml

@@ -109,4 +109,8 @@ let run ?(forbid = false) system main =
   if not (Picos_thread.is_main_thread ()) then not_main_thread ();
   let computation = Computation.create ~mode:`LIFO () in
   let fiber = Fiber.create ~forbid computation in
+  let main () =
+    Computation.capture computation main ();
+    Computation.await computation
+  in
   go fiber system (Effect.Shallow.fiber main) (Ok ())

lib/picos_threaded/picos_threaded.ml

@@ -136,5 +136,10 @@ and handler = Handler.{ current; spawn; yield; cancel_after; await }
 
 let run ?(forbid = false) main =
   Select.check_configured ();
-  let packed = Computation.Packed (Computation.create ~mode:`LIFO ()) in
-  Handler.using handler (create_packed ~forbid packed) main
+  let computation = Computation.create ~mode:`LIFO () in
+  let context = create_packed ~forbid (Packed computation) in
+  let main () =
+    Computation.capture computation main ();
+    Computation.await computation
+  in
+  Handler.using handler context main

test/test_schedulers.ml

@@ -7,6 +7,31 @@ let test_returns () =
   let actual = Test_scheduler.run @@ fun () -> 42 in
   assert (actual = 42)
 
+let test_completes () =
+  let packed = ref (Computation.Packed (Computation.create ())) in
+  let result =
+    Test_scheduler.run (fun () ->
+        packed := Fiber.get_computation (Fiber.current ());
+        101)
+  in
+  assert (result = 101);
+  let (Packed computation) = !packed in
+  assert (not (Computation.is_running computation));
+  assert (not (Computation.is_canceled computation));
+  begin
+    match
+      Test_scheduler.run (fun () ->
+          packed := Fiber.get_computation (Fiber.current ());
+          (failwith "42" : unit))
+    with
+    | () -> assert false
+    | exception Failure msg -> assert (msg = "42")
+    | exception _ -> assert false
+  end;
+  let (Packed computation) = !packed in
+  assert (not (Computation.is_running computation));
+  assert (Computation.is_canceled computation)
+
 let test_current () =
   Test_scheduler.run ~max_domains:2 @@ fun () ->
   let fiber_parent = Fiber.current () in
@@ -33,6 +58,7 @@ let test_cancel_after_long_timeout () =
 let () =
   [
     ("Returns", [ Alcotest.test_case "" `Quick test_returns ]);
+    ("Completes", [ Alcotest.test_case "" `Quick test_completes ]);
     ("Current", [ Alcotest.test_case "" `Quick test_current ]);
     ( "Cancel_after",
       [