From 0c728e378e83f4245850da9be0b1d76d3c7421a4 Mon Sep 17 00:00:00 2001
From: Leandro Ostera <leandro@abstractmachines.dev>
Date: Sat, 19 Oct 2024 20:12:26 +0200
Subject: [PATCH] feat: make runtime shutdowns fast

---
 bench/http_server.ml                  |  2 +-
 packages/riot-runtime/Config.ml       | 11 ++++++++++
 packages/riot-runtime/riot_runtime.ml |  5 +++--
 packages/riot/riot.ml                 | 31 ++++++++++++---------------
 packages/riot/riot.mli                | 18 +++++++++-------
 test/add_monitor_test.ml              |  2 +-
 test/process_priority_test.ml         |  2 +-
 test/selective_receive_test.ml        |  2 +-
 test/send_after_test.ml               |  2 +-
 9 files changed, 43 insertions(+), 32 deletions(-)
 create mode 100644 packages/riot-runtime/Config.ml

diff --git a/bench/http_server.ml b/bench/http_server.ml
index f55c0b84..b3e912cb 100644
--- a/bench/http_server.ml
+++ b/bench/http_server.ml
@@ -51,4 +51,4 @@ let main () =
   let _ = List.init 99 (fun _ -> spawn_link acceptor) in
   acceptor ()
 
-let () = Riot.run ~workers:0 @@ main
+let () = Riot.run ~config:(Config.make ~workers:0 ()) @@ main
diff --git a/packages/riot-runtime/Config.ml b/packages/riot-runtime/Config.ml
new file mode 100644
index 00000000..b4f3d0a7
--- /dev/null
+++ b/packages/riot-runtime/Config.ml
@@ -0,0 +1,11 @@
+type t = { rnd : Random.State.t; max_workers : int; workers : int }
+
+let make ?workers () =
+  let max_workers = Int.max 0 (Stdlib.Domain.recommended_domain_count () - 2) in
+  let workers =
+    match workers with Some w -> Int.min w max_workers | None -> max_workers
+  in
+  let rnd = Random.State.make_self_init () in
+  { rnd; max_workers; workers }
+
+let default () = make ()
diff --git a/packages/riot-runtime/riot_runtime.ml b/packages/riot-runtime/riot_runtime.ml
index 11870954..ae3d39f2 100644
--- a/packages/riot-runtime/riot_runtime.ml
+++ b/packages/riot-runtime/riot_runtime.ml
@@ -1,8 +1,9 @@
-module Log = Log
+module Config = Config
 module Core = Core
 module Import = Import
-module Util = Util
+module Log = Log
 module Scheduler = Scheduler
 module Time = Time
+module Util = Util
 
 let set_log_level = Log.set_log_level
diff --git a/packages/riot/riot.ml b/packages/riot/riot.ml
index 89a58f3b..a6b31a7f 100644
--- a/packages/riot/riot.ml
+++ b/packages/riot/riot.ml
@@ -2,6 +2,7 @@ include Riot_stdlib
 
 open struct
   open Riot_runtime
+  module Config = Config
   module Log = Log
   module Core = Core
   module Import = Import
@@ -10,6 +11,8 @@ open struct
   module Time = Time
 end
 
+module Config = Config
+
 open Logger.Make (struct
   let namespace = [ "riot" ]
 end)
@@ -23,13 +26,10 @@ let shutdown ?(status = 0) () =
 
 let started = ref false
 
-let run ?(rnd = Random.State.make_self_init ()) ?workers main =
+let run ?(config = Config.default ()) main =
   if !started then raise Riot_already_started else started := true;
 
-  let max_workers = Int.max 0 (Stdlib.Domain.recommended_domain_count () - 2) in
-  let workers =
-    match workers with Some w -> Int.min w max_workers | None -> max_workers
-  in
+  let Config.{ workers; rnd; _ } = config in
 
   Log.debug (fun f -> f "Initializing Riot runtime...");
   Printexc.record_backtrace true;
@@ -37,7 +37,7 @@ let run ?(rnd = Random.State.make_self_init ()) ?workers main =
   Scheduler.Uid.reset ();
 
   let sch0 = Scheduler.make ~rnd () in
-  let pool, domains = Scheduler.Pool.make ~main:sch0 ~domains:workers () in
+  let pool, _domains = Scheduler.Pool.make ~main:sch0 ~domains:workers () in
 
   Scheduler.set_current_scheduler sch0;
   Scheduler.Pool.set_pool pool;
@@ -45,8 +45,6 @@ let run ?(rnd = Random.State.make_self_init ()) ?workers main =
   let _pid = _spawn ~pool ~scheduler:sch0 main in
   Scheduler.run pool sch0 ();
 
-  Log.debug (fun f -> f "Riot runtime shutting down...");
-  List.iter Stdlib.Domain.join domains;
   Log.debug (fun f -> f "Riot runtime shutdown");
   Stdlib.exit pool.status
 
@@ -58,16 +56,15 @@ let on_error (error : [ `Msg of string ]) =
   Log.error (fun f -> f "Riot raised an error: %s\n" error_string);
   1
 
-let run_with_status ?(rnd = Random.State.make_self_init ()) ?workers ~on_error
-    main =
-  run ~rnd ?workers (fun _ ->
-      let status =
-        match main () with Ok code -> code | Error reason -> on_error reason
-      in
-      shutdown ~status ())
+let run_with_status ?config ~on_error main =
+  run ?config @@ fun _ ->
+  let status =
+    match main () with Ok code -> code | Error reason -> on_error reason
+  in
+  shutdown ~status ()
 
-let start ?rnd ?workers ~apps () =
-  run ?rnd ?workers @@ fun () ->
+let start ?config ~apps () =
+  run ?config @@ fun () ->
   let child_specs =
     List.map
       (fun (module App : Application.Intf) ->
diff --git a/packages/riot/riot.mli b/packages/riot/riot.mli
index b2e16597..b878647d 100644
--- a/packages/riot/riot.mli
+++ b/packages/riot/riot.mli
@@ -301,14 +301,20 @@ val receive_any : ?after:int64 -> ?ref:unit Ref.t -> unit -> Message.t
 val shutdown : ?status:int -> unit -> unit
 (** Gracefully shuts down the runtime. Any non-yielding process will block this. *)
 
-val run : ?rnd:Random.State.t -> ?workers:int -> (unit -> unit) -> unit
+module Config : sig
+  type t = { rnd : Random.State.t; max_workers : int; workers : int }
+
+  val default : unit -> t
+  val make : ?workers:int -> unit -> t
+end
+
+val run : ?config:Config.t -> (unit -> unit) -> unit
 (** Start the Riot runtime using function [main] to boot the system *)
 
 val on_error : [ `Msg of string ] -> int
 
 val run_with_status :
-  ?rnd:Random.State.t ->
-  ?workers:int ->
+  ?config:Config.t ->
   on_error:('error -> int) ->
   (unit -> (int, 'error) result) ->
   unit
@@ -320,11 +326,7 @@ val run_with_status :
 *)
 
 val start :
-  ?rnd:Random.State.t ->
-  ?workers:int ->
-  apps:(module Application.Intf) list ->
-  unit ->
-  unit
+  ?config:Config.t -> apps:(module Application.Intf) list -> unit -> unit
 (** Start the Riot runtime with a series of applications.
 
     Each application will be started in the same order as specified, and
diff --git a/test/add_monitor_test.ml b/test/add_monitor_test.ml
index 9d18672e..bdeb9586 100644
--- a/test/add_monitor_test.ml
+++ b/test/add_monitor_test.ml
@@ -18,4 +18,4 @@ let main () =
       sleep 0.2;
       raise Fail
 
-let () = Riot.run ~workers:1 @@ main
+let () = Riot.run ~config:(Config.make ~workers:1 ()) @@ main
diff --git a/test/process_priority_test.ml b/test/process_priority_test.ml
index a5c054ef..934a394e 100644
--- a/test/process_priority_test.ml
+++ b/test/process_priority_test.ml
@@ -53,4 +53,4 @@ let main () =
 
    That behavior _is expected_.
 *)
-let () = Riot.run ~workers:0 @@ main
+let () = Riot.run ~config:(Config.make ~workers:0 ()) @@ main
diff --git a/test/selective_receive_test.ml b/test/selective_receive_test.ml
index dc3adbd0..c9211443 100644
--- a/test/selective_receive_test.ml
+++ b/test/selective_receive_test.ml
@@ -40,4 +40,4 @@ let main () =
       sleep 1.;
       Stdlib.exit 1
 
-let () = Riot.run ~workers:0 @@ main
+let () = Riot.run ~config:(Config.make ~workers:0 ()) @@ main
diff --git a/test/send_after_test.ml b/test/send_after_test.ml
index 83719b6f..4ac44f79 100644
--- a/test/send_after_test.ml
+++ b/test/send_after_test.ml
@@ -47,4 +47,4 @@ let main () =
       sleep 0.1;
       shutdown ~status:1 ()
 
-let () = Riot.run ~workers:0 @@ main
+let () = Riot.run ~config:(Config.make ~workers:0 ()) @@ main