Repository: ocaml-multicore/picos
Pull request #135 (2024-05-10 13:06:58)

Support multi-threaded use in `Picos_randos`
This PR adds a low level interface to `Picos_randos` to allow it to use multiple threads and/or domains.


Patch:
diff --git a/lib/picos_randos/dune b/lib/picos_randos/dune
index ea55fcd..75ea255 100644
--- a/lib/picos_randos/dune
+++ b/lib/picos_randos/dune
@@ -10,4 +10,5 @@
    from
    (picos.select -> select.some.ml)
    (-> select.none.ml))
+  multicore-magic
   picos.htbl))
diff --git a/lib/picos_randos/picos_randos.ml b/lib/picos_randos/picos_randos.ml
index 9289f62..014d9c1 100644
--- a/lib/picos_randos/picos_randos.ml
+++ b/lib/picos_randos/picos_randos.ml
@@ -19,6 +19,11 @@ module Collection = struct
   let pop_exn t =
     let key = Picos_htbl.find_random_exn t in
     Picos_htbl.remove_exn t key
+
+  let is_empty t =
+    match Picos_htbl.find_random_exn t with
+    | _ -> false
+    | exception Not_found -> true
 end
 
 type ready =
@@ -31,16 +36,18 @@ type ready =
 
 type t = {
   ready : ready Collection.t;
-  needs_wakeup : bool Atomic.t;
+  num_waiters_non_zero : bool ref;
   num_alive_fibers : int Atomic.t;
-  mutex : Mutex.t;
-  condition : Condition.t;
   resume :
     Trigger.t ->
     Fiber.t ->
     (Exn_bt.t option, unit) Effect.Deep.continuation ->
     unit;
   retc : unit -> unit;
+  num_waiters : int ref;
+  condition : Condition.t;
+  mutex : Mutex.t;
+  mutable run : bool;
 }
 
 let rec spawn t n forbid packed = function
@@ -48,6 +55,7 @@ let rec spawn t n forbid packed = function
   | main :: mains ->
       let fiber = Fiber.create_packed ~forbid packed in
       Collection.push t.ready (Spawn (fiber, main));
+      if !(t.num_waiters_non_zero) then Condition.signal t.condition;
       spawn t (n + 1) forbid packed mains
 
 let rec next t =
@@ -111,29 +119,51 @@ let rec next t =
   | Resume (fiber, k) -> Fiber.resume fiber k
   | exception Not_found ->
       if Atomic.get t.num_alive_fibers <> 0 then begin
-        if Atomic.get t.needs_wakeup then begin
-          Mutex.lock t.mutex;
-          match
-            if Atomic.get t.needs_wakeup then Condition.wait t.condition t.mutex
-          with
-          | () -> Mutex.unlock t.mutex
-          | exception exn ->
+        Mutex.lock t.mutex;
+        if Collection.is_empty t.ready && Atomic.get t.num_alive_fibers <> 0
+        then begin
+          let n = !(t.num_waiters) + 1 in
+          t.num_waiters := n;
+          if n = 1 then t.num_waiters_non_zero := true;
+          match Condition.wait t.condition t.mutex with
+          | () ->
+              let n = !(t.num_waiters) - 1 in
+              t.num_waiters := n;
+              if n = 0 then t.num_waiters_non_zero := false;
+              Mutex.unlock t.mutex;
+              next t
+          | exception async_exn ->
+              let n = !(t.num_waiters) - 1 in
+              t.num_waiters := n;
+              if n = 0 then t.num_waiters_non_zero := false;
               Mutex.unlock t.mutex;
-              raise exn
+              raise async_exn
         end
-        else Atomic.set t.needs_wakeup true;
-        next t
+        else begin
+          Mutex.unlock t.mutex;
+          next t
+        end
+      end
+      else begin
+        Mutex.lock t.mutex;
+        Mutex.unlock t.mutex;
+        Condition.broadcast t.condition
       end
 
-let run ?(forbid = false) main =
+let context () =
   Select.check_configured ();
-  let ready = Collection.create ()
-  and needs_wakeup = Atomic.make false
-  and num_alive_fibers = Atomic.make 1
-  and mutex = Mutex.create ()
-  and condition = Condition.create () in
   let rec t =
-    { ready; needs_wakeup; num_alive_fibers; mutex; condition; resume; retc }
+    {
+      ready = Collection.create ();
+      num_waiters_non_zero = ref false |> Multicore_magic.copy_as_padded;
+      num_alive_fibers = Atomic.make 1 |> Multicore_magic.copy_as_padded;
+      resume;
+      retc;
+      num_waiters = ref 0 |> Multicore_magic.copy_as_padded;
+      condition = Condition.create ();
+      mutex = Mutex.create ();
+      run = false;
+    }
   and retc () =
     Atomic.decr t.num_alive_fibers;
     next t
@@ -141,21 +171,48 @@ let run ?(forbid = false) main =
     let resume = Resume (fiber, k) in
     Fiber.unsuspend fiber trigger |> ignore;
     Collection.push t.ready resume;
-    if
-      Atomic.get t.needs_wakeup
-      && Atomic.compare_and_set t.needs_wakeup true false
-    then begin
-      begin
-        match Mutex.lock t.mutex with
-        | () -> Mutex.unlock t.mutex
-        | exception Sys_error _ -> ()
-      end;
-      Condition.broadcast t.condition
-    end
+    let non_zero =
+      match Mutex.lock t.mutex with
+      | () ->
+          let non_zero = !(t.num_waiters_non_zero) in
+          Mutex.unlock t.mutex;
+          non_zero
+      | exception Sys_error _ -> false
+    in
+    if non_zero then Condition.signal t.condition
   in
-  let computation = Computation.create () in
-  let fiber = Fiber.create ~forbid computation in
-  let main = Computation.capture computation main in
-  Collection.push t.ready (Spawn (fiber, main));
-  next t;
-  Computation.await computation
+  t
+
+let runner_on_this_thread = next
+
+let rec await t computation =
+  if !(t.num_waiters_non_zero) then begin
+    match Condition.wait t.condition t.mutex with
+    | () -> await t computation
+    | exception async_exn ->
+        Mutex.unlock t.mutex;
+        raise async_exn
+  end
+  else begin
+    Mutex.unlock t.mutex;
+    Computation.await computation
+  end
+
+let run ?context:t_opt ?(forbid = false) main =
+  let t = match t_opt with Some t -> t | None -> context () in
+  Mutex.lock t.mutex;
+  if t.run then begin
+    Mutex.unlock t.mutex;
+    invalid_arg "already run"
+  end
+  else begin
+    t.run <- true;
+    Mutex.unlock t.mutex;
+    let computation = Computation.create () in
+    let fiber = Fiber.create ~forbid computation in
+    let main = Computation.capture computation main in
+    Collection.push t.ready (Spawn (fiber, main));
+    next t;
+    Mutex.lock t.mutex;
+    await t computation
+  end
diff --git a/lib/picos_randos/picos_randos.mli b/lib/picos_randos/picos_randos.mli
index db61f55..0e36c3e 100644
--- a/lib/picos_randos/picos_randos.mli
+++ b/lib/picos_randos/picos_randos.mli
@@ -1,4 +1,4 @@
-(** Randomized single-threaded effects based {!Picos} compatible scheduler for
+(** Randomized multi-threaded effects based {!Picos} compatible scheduler for
     OCaml 5.
 
     ℹ️ This scheduler implementation is specifically intended for testing
@@ -10,9 +10,29 @@
     bugs in programs implemented in Picos that make invalid scheduling
     assumptions. *)
 
-val run : ?forbid:bool -> (unit -> 'a) -> 'a
+type t
+(** Represents a shared context for randomized runners. *)
+
+val context : unit -> t
+(** [context ()] creates a new context for randomized runners.  The context
+    should be consumed by a call of {{!run} [run ~context ...]}. *)
+
+val runner_on_this_thread : t -> unit
+(** [runner_on_this_thread context] starts a runner on the current thread to run
+    fibers on the context.  The runner returns when {{!run} [run ~context ...]}
+    returns. *)
+
+val run : ?context:t -> ?forbid:bool -> (unit -> 'a) -> 'a
 (** [run main] runs the [main] thunk with the scheduler.  Returns after [main]
     and all of the fibers spawned by [main] have returned.
 
+    The optional [context] argument specifies a context in which to run the
+    [main] thunk.  If unspecified, a new context is automatically created and
+    the scheduler will be single-threaded.  By {{!context} creating a context},
+    spawning concurrent or parallel {{!runner_on_this_thread} runners} on to the
+    context, and then explicitly passing the context to [run ~context ...] one
+    can create a multi-threaded scheduler.  Only a single call of {!run} per
+    context is allowed.
+
     The optional [forbid] argument defaults to [false] and determines whether
     propagation of cancelation is initially allowed. *)


Test patch:
diff --git a/test/dune b/test/dune
index 05eb5f3..5e8403e 100644
--- a/test/dune
+++ b/test/dune
@@ -6,7 +6,7 @@
   (select
    test_scheduler.ml
    from
-   (picos.fifos picos.randos -> test_scheduler.ocaml5.ml)
+   (picos.fifos picos.randos picos.structured -> test_scheduler.ocaml5.ml)
    (picos.threaded -> test_scheduler.ocaml4.ml))))
 
 (library
@@ -31,7 +31,7 @@
 (test
  (name test_sync)
  (modules test_sync)
- (libraries picos.sync test_scheduler alcotest domain_shims))
+ (libraries picos.structured picos.sync test_scheduler alcotest domain_shims))
 
 ;;
 
diff --git a/test/test_picos.ml b/test/test_picos.ml
index aa92d7c..efa9046 100644
--- a/test/test_picos.ml
+++ b/test/test_picos.ml
@@ -19,7 +19,7 @@ let test_fls_basics =
   in
 
   fun () ->
-    Test_scheduler.run @@ fun () ->
+    Test_scheduler.run ~max_domains:2 @@ fun () ->
     let first =
       run_in_fiber @@ fun () ->
       let fiber = Fiber.current () in
@@ -103,7 +103,7 @@ let test_thread_cancelation () =
 
 let test_cancel_after () =
   Alcotest.check_raises "should be canceled" Not_found @@ fun () ->
-  Test_scheduler.run @@ fun () ->
+  Test_scheduler.run ~max_domains:2 @@ fun () ->
   let computation = Computation.create () in
   let main =
     Computation.capture computation @@ fun () ->
diff --git a/test/test_scheduler.ocaml4.ml b/test/test_scheduler.ocaml4.ml
index 20d8010..447826b 100644
--- a/test/test_scheduler.ocaml4.ml
+++ b/test/test_scheduler.ocaml4.ml
@@ -1,2 +1,2 @@
 let () = Random.self_init ()
-let run ?forbid main = Picos_threaded.run ?forbid main
+let run ?max_domains:_ ?forbid main = Picos_threaded.run ?forbid main
diff --git a/test/test_scheduler.ocaml5.ml b/test/test_scheduler.ocaml5.ml
index 14d906d..0c347be 100644
--- a/test/test_scheduler.ocaml5.ml
+++ b/test/test_scheduler.ocaml5.ml
@@ -1,5 +1,22 @@
+open Picos_structured.Finally
+
 let use_randos = Random.bool (Random.self_init ())
 
-let run ?forbid main =
-  if use_randos then Picos_randos.run ?forbid main
+let run ?(max_domains = 1) ?forbid main =
+  if use_randos then
+    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)
+    in
+    spawn (Int.min max_domains (Domain.recommended_domain_count ()))
   else Picos_fifos.run ?forbid main
diff --git a/test/test_schedulers.ml b/test/test_schedulers.ml
index 5b6ce97..c5940b6 100644
--- a/test/test_schedulers.ml
+++ b/test/test_schedulers.ml
@@ -7,7 +7,7 @@ let test_returns () =
   assert (actual = 42)
 
 let test_current () =
-  Test_scheduler.run @@ fun () ->
+  Test_scheduler.run ~max_domains:2 @@ fun () ->
   let fiber_parent = Fiber.current () in
   let fiber_child = ref fiber_parent in
   let computation = Computation.create () in
diff --git a/test/test_server_and_client.ml b/test/test_server_and_client.ml
index b498962..aa83277 100644
--- a/test/test_server_and_client.ml
+++ b/test/test_server_and_client.ml
@@ -128,6 +128,6 @@ let () =
   Printf.printf "Using %sblocking sockets and %s:\n%!"
     (if use_nonblock then "non-" else "")
     (if is_ocaml4 then "threads on OCaml 4" else "fibers on OCaml 5");
-  try Test_scheduler.run main with
+  try Test_scheduler.run ~max_domains:4 main with
   | Exit -> Printf.printf "Server and Client test: SKIPPED\n%!"
   | exn -> Printf.printf "Error: %s\n%!" (Printexc.to_string exn)
diff --git a/test/test_stdio.ml b/test/test_stdio.ml
index 6151bd3..135a6d7 100644
--- a/test/test_stdio.ml
+++ b/test/test_stdio.ml
@@ -1,4 +1,3 @@
-open Picos
 open Picos_structured.Finally
 open Picos_structured
 open Picos_stdio
@@ -24,15 +23,9 @@ let test_openfile_and_read () =
   assert (Bytes.to_string bytes = "open Picos")
 
 let test_sleepf () =
-  Test_scheduler.run @@ fun () ->
-  let children = Computation.create () in
-  let n = Atomic.make 100 in
+  Test_scheduler.run ~max_domains:2 @@ fun () ->
   let start = Unix.gettimeofday () in
-  Fiber.spawn ~forbid:false children
-    ( List.init (Atomic.get n) @@ fun _ () ->
-      Unix.sleepf 0.01;
-      if 1 = Atomic.fetch_and_add n (-1) then Computation.finish children );
-  Computation.await children;
+  Run.all (List.init 100 @@ fun _ () -> Unix.sleepf 0.01);
   let d = Unix.gettimeofday () -. start in
   (* This is non-deterministic and might need to be changed if flaky *)
   assert (0.01 <= d && d < 5.0)
@@ -46,7 +39,7 @@ let test_select_empty_timeout () =
   assert (0.1 <= d && d <= 5.0)
 
 let test_select_empty_forever () =
-  Test_scheduler.run @@ fun () ->
+  Test_scheduler.run ~max_domains:2 @@ fun () ->
   Bundle.join_after @@ fun bundle ->
   begin
     Bundle.fork bundle @@ fun () ->
@@ -57,7 +50,7 @@ let test_select_empty_forever () =
   Bundle.terminate bundle
 
 let test_select () =
-  Test_scheduler.run @@ fun () ->
+  Test_scheduler.run ~max_domains:2 @@ fun () ->
   let@ msg_inn1, msg_out1 =
     finally Unix.close_pair @@ fun () ->
     Unix.socketpair PF_UNIX SOCK_STREAM 0 ~cloexec:true
diff --git a/test/test_structured.ml b/test/test_structured.ml
index 2b3d915..2626f60 100644
--- a/test/test_structured.ml
+++ b/test/test_structured.ml
@@ -1,4 +1,4 @@
-open Picos
+module Exn_bt = Picos.Exn_bt
 open Picos_structured
 open Picos_sync
 
@@ -37,7 +37,7 @@ let test_fork_after_escape () =
 
 let test_exception_in_child_terminates () =
   match
-    Test_scheduler.run @@ fun () ->
+    Test_scheduler.run ~max_domains:2 @@ fun () ->
     let mutex = Mutex.create () in
     let condition = Condition.create () in
     let blocked = ref false in
@@ -65,7 +65,7 @@ let test_exception_in_child_terminates () =
   | exception Exit -> ()
 
 let test_cancelation_awaits_children () =
-  Test_scheduler.run @@ fun () ->
+  Test_scheduler.run ~max_domains:3 @@ fun () ->
   let blocked = ref false in
   let slept = ref false in
   begin
@@ -97,7 +97,7 @@ let test_cancelation_awaits_children () =
   assert !slept
 
 let test_termination_nests () =
-  Test_scheduler.run @@ fun () ->
+  Test_scheduler.run ~max_domains:3 @@ fun () ->
   let mutex = Mutex.create () in
   let condition = Condition.create () in
   let blocked = ref false in
@@ -131,7 +131,7 @@ let test_termination_nests () =
   end
 
 let test_promise_cancelation_does_not_terminate () =
-  Test_scheduler.run @@ fun () ->
+  Test_scheduler.run ~max_domains:2 @@ fun () ->
   Bundle.join_after @@ fun bundle ->
   let promise = Bundle.fork_as_promise bundle @@ fun () -> Control.block () in
   Control.yield ();
@@ -140,7 +140,7 @@ let test_promise_cancelation_does_not_terminate () =
 
 let test_error_in_promise_terminates () =
   match
-    Test_scheduler.run @@ fun () ->
+    Test_scheduler.run ~max_domains:2 @@ fun () ->
     Bundle.join_after @@ fun bundle ->
     let promise =
       Bundle.fork_as_promise bundle @@ fun () -> failwith "I failed"
@@ -152,7 +152,7 @@ let test_error_in_promise_terminates () =
   | exception Failure _ -> ()
 
 let test_can_wait_promises () =
-  Test_scheduler.run @@ fun () ->
+  Test_scheduler.run ~max_domains:2 @@ fun () ->
   Bundle.join_after @@ fun bundle ->
   let promise =
     Bundle.fork_as_promise bundle @@ fun () ->
@@ -164,7 +164,7 @@ let test_can_wait_promises () =
 let test_any_and_all_errors () =
   [ Run.all; Run.any ]
   |> List.iter @@ fun run_op ->
-     Test_scheduler.run @@ fun () ->
+     Test_scheduler.run ~max_domains:6 @@ fun () ->
      let raised = Picos_mpsc_queue.create () in
      let raiser exn () =
        Picos_mpsc_queue.push raised exn;
@@ -203,7 +203,8 @@ let test_any_and_all_returns () =
      |> List.iter @@ fun n_incr ->
         [ (Run.all, n_incr, n_incr); (Run.any, Int.min 1 n_incr, n_incr) ]
         |> List.iter @@ fun (run_op, min, max) ->
-           Test_scheduler.run @@ fun () ->
+           Test_scheduler.run ~max_domains:(n_terminates + n_incr + 1)
+           @@ fun () ->
            let count = Atomic.make 0 in
            let ops =
              List.init n_terminates (fun _ () ->
@@ -217,7 +218,7 @@ let test_any_and_all_returns () =
            assert (n <= max)
 
 let test_race_any () =
-  Test_scheduler.run @@ fun () ->
+  Test_scheduler.run ~max_domains:4 @@ fun () ->
   let winner = ref 0 in
   let ops =
     Test_util.shuffle
diff --git a/test/test_sync.ml b/test/test_sync.ml
index 39e4f36..598a9e5 100644
--- a/test/test_sync.ml
+++ b/test/test_sync.ml
@@ -1,5 +1,6 @@
 open Picos
 open Picos_sync
+open Picos_structured
 
 module Fiber = struct
   include Fiber
@@ -9,14 +10,12 @@ module Fiber = struct
     Fiber.spawn ~forbid:false computation
       [ Computation.capture computation thunk ];
     computation
-
-  let run thunk = Computation.await (start thunk)
 end
 
 let is_ocaml4 = String.starts_with ~prefix:"4." Sys.ocaml_version
 
 let test_mutex_and_condition_basics () =
-  Test_scheduler.run @@ fun () ->
+  Test_scheduler.run ~max_domains:2 @@ fun () ->
   let mutex = Mutex.create () in
   let condition = Condition.create () in
 
@@ -43,17 +42,20 @@ let test_mutex_and_condition_basics () =
            done
 
 let test_mutex_and_condition_errors () =
-  Test_scheduler.run @@ fun () ->
+  Test_scheduler.run ~max_domains:2 @@ fun () ->
   let mutex = Mutex.create () in
   let condition = Condition.create () in
   Mutex.protect mutex @@ fun () ->
+  Bundle.join_after @@ fun bundle ->
   begin
-    match Fiber.run (fun () -> Mutex.unlock mutex) with
+    Bundle.fork bundle @@ fun () ->
+    match Mutex.unlock mutex with
     | () -> assert false
     | exception Sys_error _ -> ()
   end;
   begin
-    match Fiber.run (fun () -> Condition.wait condition mutex) with
+    Bundle.fork bundle @@ fun () ->
+    match Condition.wait condition mutex with
     | () -> assert false
     | exception Sys_error _ -> ()
   end
@@ -165,7 +167,7 @@ let test_lazy_basics () =
   | exception Stdlib.Lazy.Undefined -> ()
 
 let test_lazy_cancelation () =
-  Test_scheduler.run @@ fun () ->
+  Test_scheduler.run ~max_domains:2 @@ fun () ->
   let susp ?to_signal ~to_await result =
     Lazy.from_fun @@ fun () ->
     Option.iter Trigger.signal to_signal;