Repository: ocaml-multicore/picos
Pull request #200 (2024-07-31 12:54:06)

Add `Picos_randos.run_on` helper for running on multiple domains


Patch:
diff --git a/lib/picos/dune b/lib/picos/dune
index 50c4314..cb2e4cb 100644
--- a/lib/picos/dune
+++ b/lib/picos/dune
@@ -24,5 +24,5 @@
 (mdx
  (enabled_if
   (>= %{ocaml_version} 5.0.0))
- (libraries picos picos.exn_bt picos.fifos picos.structured)
+ (libraries picos picos.exn_bt picos.structured picos.randos)
  (files picos.mli))
diff --git a/lib/picos/picos.mli b/lib/picos/picos.mli
index f7ed275..8927fe8 100644
--- a/lib/picos/picos.mli
+++ b/lib/picos/picos.mli
@@ -113,13 +113,14 @@
       open Picos
     ]}
 
-    as well as a simple helper for cleaning up resources
+    as well as the {!Picos_structured} library,
 
     {[
-      open Picos_structured.Finally
+      open Picos_structured
     ]}
 
-    and define a simple scheduler on OCaml 4
+    which we will be using for managing fibers in some of the examples, and
+    define a simple scheduler on OCaml 4
 
     {@ocaml version<5.0.0[
       let run main = Picos_threaded.run main
@@ -128,23 +129,18 @@
     using {{!Picos_threaded} the basic thread based scheduler} and on OCaml 5
 
     {@ocaml version>=5.0.0[
-      let run main = Picos_fifos.run main
+      let run main = Picos_randos.run_on ~n_domains:2 main
     ]}
 
-    using {{!Picos_fifos} the basic effects based scheduler} that come with
-    Picos as samples.
+    using {{!Picos_randos} the randomized effects based scheduler} that come
+    with Picos as samples.
 
     {2 Auxiliary modules} *)
 
 module Exn_bt = Picos_exn_bt
 (** Exceptions with backtraces. *)
 
-(** {2 Core modules}
-
-    Please note that the example code snippets in this documentation may
-    e.g. use the {!Domain} and {!Unix} modules in order to be able to describe
-    Picos concepts in isolation in the absence of a Picos compatible
-    scheduler. *)
+(** {2 Core modules} *)
 
 module Trigger : sig
   (** Ability to await for a signal.
@@ -157,13 +153,14 @@ module Trigger : sig
 
       {[
         run begin fun () ->
+          Bundle.join_after @@ fun bundle ->
+
           let trigger = Trigger.create () in
 
-          let@ signaler =
-            finally Domain.join @@ fun () ->
-            Domain.spawn @@ fun () ->
-              Trigger.signal trigger
-          in
+          Bundle.fork bundle begin fun () ->
+            Trigger.signal trigger
+          end;
+
           match Trigger.await trigger with
           | None ->
             (* We were resumed normally. *)
@@ -431,30 +428,34 @@ module Computation : sig
 
       {[
         run begin fun () ->
+          Bundle.join_after @@ fun bundle ->
+
           let computation =
             Computation.create ()
           in
-          let@ computer =
-            finally Domain.join @@ fun () ->
-            Domain.spawn @@ fun () ->
-              let rec fib i =
-                Computation.check
-                  computation;
-                if i <= 1 then
-                  i
-                else
-                  fib (i - 1) + fib (i - 2)
-              in
-              Computation.capture
-                computation fib 10
-          in
-          let@ canceler =
-            finally Domain.join @@ fun () ->
-            Domain.spawn @@ fun () ->
-              Unix.sleepf 0.1;
-              Computation.cancel computation
-              @@ Exn_bt.get_callstack 2 Exit
+
+          let canceler = Bundle.fork_as_promise bundle @@ fun () ->
+            Fiber.sleep ~seconds:1.0;
+
+            Computation.cancel computation
+              (Exn_bt.get_callstack 0 Exit)
           in
+
+          Bundle.fork bundle begin fun () ->
+            let rec fib i =
+              Computation.check
+                computation;
+              if i <= 1 then
+                i
+              else
+                fib (i - 1) + fib (i - 2)
+            in
+            Computation.capture computation
+              fib 10;
+
+            Promise.terminate canceler
+          end;
+
           Computation.await computation
         end
       ]}
diff --git a/lib/picos_randos/picos_randos.ml b/lib/picos_randos/picos_randos.ml
index 25fe3f0..480e410 100644
--- a/lib/picos_randos/picos_randos.ml
+++ b/lib/picos_randos/picos_randos.ml
@@ -225,3 +225,26 @@ let run ?context:t_opt ?(forbid = false) main =
     Mutex.lock t.mutex;
     await t computation
   end
+
+let rec run_on n ?forbid main context =
+  if n <= 1 then run ~context ?forbid main
+  else
+    let runner =
+      try Domain.spawn @@ fun () -> runner_on_this_thread context
+      with exn ->
+        let bt = Printexc.get_raw_backtrace () in
+        run ~context Fun.id;
+        Printexc.raise_with_backtrace exn bt
+    in
+    match run_on (n - 1) ?forbid main context with
+    | result ->
+        Domain.join runner;
+        result
+    | exception exn ->
+        let bt = Printexc.get_raw_backtrace () in
+        Domain.join runner;
+        Printexc.raise_with_backtrace exn bt
+
+let run_on ~n_domains ?forbid main =
+  if n_domains < 1 then invalid_arg "n_domains must be positive";
+  run_on n_domains ?forbid main (context ())
diff --git a/lib/picos_randos/picos_randos.mli b/lib/picos_randos/picos_randos.mli
index 0e36c3e..6aa0b7a 100644
--- a/lib/picos_randos/picos_randos.mli
+++ b/lib/picos_randos/picos_randos.mli
@@ -36,3 +36,7 @@ val run : ?context:t -> ?forbid:bool -> (unit -> 'a) -> 'a
 
     The optional [forbid] argument defaults to [false] and determines whether
     propagation of cancelation is initially allowed. *)
+
+val run_on : n_domains:int -> ?forbid:bool -> (unit -> 'a) -> 'a
+(** [run_on ~n_domains main] spawns [n_domains - 1] additional domains and runs
+    the [main] on the current domain and those additional domains. *)


Test patch:
diff --git a/test/dune b/test/dune
index 3be5980..650717a 100644
--- a/test/dune
+++ b/test/dune
@@ -10,7 +10,6 @@
    (picos.fifos
     picos.randos
     picos.lwt_unix
-    picos.structured
     picos.thread
     ->
     test_scheduler.ocaml5.ml)
diff --git a/test/test_scheduler.ocaml5.ml b/test/test_scheduler.ocaml5.ml
index b0bb2fe..2d3d19f 100644
--- a/test/test_scheduler.ocaml5.ml
+++ b/test/test_scheduler.ocaml5.ml
@@ -1,8 +1,6 @@
 let () = Printexc.record_backtrace true
 let () = Random.self_init ()
 
-open Picos_structured.Finally
-
 let () =
   Picos_select.check_configured ();
 
@@ -27,20 +25,8 @@ let rec run ?(max_domains = 1) ?(allow_lwt = true) ?forbid main =
         Lwt_main.run (Picos_lwt_unix.run ?forbid main)
       else run ~max_domains ~allow_lwt ?forbid main
   | `Randos ->
-      let context = Picos_randos.context () in
-      let rec spawn n =
-        if n <= 1 then Picos_randos.run ~context ?forbid main
-        else
-          let@ _ =
-            finally Domain.join @@ fun () ->
-            try
-              Domain.spawn @@ fun () ->
-              Picos_randos.runner_on_this_thread context
-            with exn ->
-              Picos_randos.run ~context Fun.id;
-              raise exn
-          in
-          spawn (n - 1)
+      let n_domains =
+        Int.min max_domains (Domain.recommended_domain_count ())
       in
-      spawn (Int.min max_domains (Domain.recommended_domain_count ()))
+      Picos_randos.run_on ~n_domains ?forbid main
   | `Fifos -> Picos_fifos.run ?forbid main