refactor(rpc): move server lifecycle into rpc library

Move generic RPC server run, ready, and stop handling from dune_rpc_impl into Rpc.Server.Lifecycle. Keep Dune-specific startup-failure handling and the status-line wrapper in dune_rpc_impl.

Signed-off-by: Rudi Grinberg <me@rgrinberg.com>

Author: rgrinberg

src/dune_rpc_impl/dune

@@ -9,7 +9,6 @@
   source
   memo
   dune_util
-  dune_trace
   fiber
   stdune
   unix)

src/dune_rpc_impl/server.ml

@@ -27,53 +27,6 @@ module Session = Rpc.Server.Session
 module Handler = Rpc.Server.Handler
 module Csexp_rpc = Rpc.Csexp_rpc
 
-module Run = struct
-  type t =
-    { handler : Rpc.Server.t
-    ; action_runner : Action_runner.Rpc_server.t
-    ; pool : Fiber.Pool.t
-    ; root : string
-    ; where : Dune_rpc.Where.t
-    ; server : Csexp_rpc.Server.t Lazy.t
-    ; startup_ivar : (Csexp_rpc.Server.t, Exn_with_backtrace.t) result Fiber.Ivar.t
-    ; registry : [ `Add | `Skip ]
-    ; watch_mode : Watch_mode_config.t
-    }
-
-  let run t =
-    let registry = Rpc.Registry.create ~root:t.root ~where:t.where t.registry in
-    let print_uncaught_rpc_error exn =
-      Console.print [ Pp.text "Uncaught RPC Error"; Exn_with_backtrace.pp exn ]
-    in
-    let with_print_errors f () =
-      Fiber.with_error_handler f ~on_error:(fun exn ->
-        print_uncaught_rpc_error exn;
-        Exn_with_backtrace.reraise exn)
-    in
-    let run () =
-      let open Fiber.O in
-      match Exn_with_backtrace.try_with (fun () -> Lazy.force t.server) with
-      | Error exn ->
-        Dune_trace.emit Rpc (fun () -> Dune_trace.Event.Rpc.startup_failure exn);
-        print_uncaught_rpc_error exn;
-        Fiber.Ivar.fill t.startup_ivar (Error exn)
-      | Ok server ->
-        let* () = Fiber.Ivar.fill t.startup_ivar (Ok server) in
-        Fiber.fork_and_join_unit
-          (fun () ->
-             let* sessions = Csexp_rpc.Server.serve server in
-             let () = Rpc.Registry.register registry in
-             let* () = Rpc.Server.serve sessions t.handler in
-             Fiber.Pool.close t.pool)
-          (fun () -> Fiber.Pool.run t.pool)
-    in
-    with_print_errors run
-    |> Fiber.finalize ~finally:(fun () ->
-      Rpc.Registry.cleanup registry;
-      Fiber.return ())
-  ;;
-end
-
 type build_request =
   | Build of Dune_lang.Dep_conf.t list
   | Runtest of string list
@@ -114,7 +67,11 @@ module Clients = struct
 end
 
 type server =
-  { config : Run.t
+  { lifecycle : Rpc.Server.Lifecycle.t
+  ; action_runner : Action_runner.Rpc_server.t
+  ; where : Dune_rpc.Where.t
+  ; registry : [ `Add | `Skip ]
+  ; watch_mode : Watch_mode_config.t
   ; mutable clients : Clients.t
   }
 
@@ -139,7 +96,7 @@ let pp_client_count t =
 ;;
 
 let refresh_client_count_status_line t =
-  match t.server.config.registry with
+  match t.server.registry with
   | `Skip -> ()
   | `Add -> Console.Status_line.refresh ()
 ;;
@@ -152,9 +109,9 @@ let () =
 ;;
 
 let ready (t : t) =
-  Fiber.Ivar.read t.server.config.startup_ivar
+  Rpc.Server.Lifecycle.ready t.server.lifecycle
   >>= function
-  | Ok server -> Csexp_rpc.Server.ready server
+  | Ok () -> Fiber.return ()
   | Error exn ->
     Dune_util.Report_error.report exn;
     Scheduler.shutdown `Failure;
@@ -163,13 +120,8 @@ let ready (t : t) =
 
 let stop (t : t) =
   Fiber.fork_and_join_unit
-    (fun () -> Action_runner.Rpc_server.stop t.server.config.action_runner)
-    (fun () ->
-       Fiber.of_thunk (fun () ->
-         match Fiber.Ivar.peek t.server.config.startup_ivar with
-         | None -> Fiber.return ()
-         | Some (Error _) -> Fiber.return ()
-         | Some (Ok server) -> Csexp_rpc.Server.stop server))
+    (fun () -> Action_runner.Rpc_server.stop t.server.action_runner)
+    (fun () -> Rpc.Server.Lifecycle.stop t.server.lifecycle)
 ;;
 
 let current_errors () =
@@ -373,7 +325,7 @@ let handler (t : t Fdecl.t) action_runner_server : unit Handler.t =
   let () =
     let f _ () =
       let t = Fdecl.get t in
-      match t.server.config.watch_mode with
+      match t.server.watch_mode with
       | No -> Fiber.return `Not_in_watch_mode
       | Yes _ ->
         let+ () = Scheduler.flush_file_watcher () in
@@ -393,7 +345,7 @@ let handler (t : t Fdecl.t) action_runner_server : unit Handler.t =
                Session.Stage1.close entry.session))
       in
       let shutdown () =
-        let* () = Csexp_rpc.Server.stop (Lazy.force t.server.config.server) in
+        let* () = stop t in
         Scheduler.shutdown `Ok;
         Fiber.return ()
       in
@@ -490,18 +442,11 @@ let create ~registry ~root ~build watch_mode =
     in
     let action_runner = Action_runner.Rpc_server.create () in
     let handler = Rpc.Server.make (handler t action_runner) in
-    { Run.handler
-    ; action_runner
-    ; pool = Fiber.Pool.create ()
-    ; root
-    ; where
-    ; server
-    ; registry
-    ; startup_ivar = Fiber.Ivar.create ()
-    ; watch_mode
-    }
+    let lifecycle = Rpc.Server.Lifecycle.create ~handler ~root ~where ~registry ~server in
+    action_runner, lifecycle
   in
-  let server = { config; clients = Clients.empty } in
+  let action_runner, lifecycle = config in
+  let server = { lifecycle; action_runner; where; registry; watch_mode; clients = Clients.empty } in
   let res = { server; build } in
   current := Some server;
   Fdecl.set t res;
@@ -511,10 +456,10 @@ let create ~registry ~root ~build watch_mode =
 let run t =
   let run () =
     Fiber.fork_and_join_unit
-      (fun () -> Run.run t.server.config)
-      (fun () -> Action_runner.Rpc_server.run t.server.config.action_runner)
+      (fun () -> Rpc.Server.Lifecycle.run t.server.lifecycle)
+      (fun () -> Action_runner.Rpc_server.run t.server.action_runner)
   in
-  match t.server.config.registry with
+  match t.server.registry with
   | `Skip -> run ()
   | `Add ->
     let section = Console.Status_line.add_section (Live (fun () -> pp_client_count t)) in
@@ -578,5 +523,5 @@ end
 
 let with_background_rpc = Background.with_background_rpc
 let ensure_ready = Background.ensure_ready
-let listening_address t = t.server.config.where
-let action_runner t = t.server.config.action_runner
+let listening_address t = t.server.where
+let action_runner t = t.server.action_runner

src/dune_rpc_impl/server.mli

@@ -21,11 +21,6 @@ val create
   -> Watch_mode_config.t
   -> t
 
-(** Stop accepting new rpc connections. Fiber returns when all existing
-    connections terminate *)
-val stop : t -> unit Fiber.t
-
-val ready : t -> unit Fiber.t
 val run : t -> unit Fiber.t
 val with_background_rpc : t -> (unit -> 'a Fiber.t) -> 'a Fiber.t
 val ensure_ready : unit -> unit Fiber.t

src/rpc/server.ml

@@ -1,4 +1,5 @@
 open Stdune
+module Rpc_registry = Registry
 open Dune_rpc.Private
 open Fiber.O
 module Session_id = Stdune.Id.Make ()
@@ -883,3 +884,63 @@ let serve =
   in
   M.serve
 ;;
+
+module Lifecycle = struct
+  type nonrec t =
+    { handler : t
+    ; registry : Rpc_registry.t
+    ; server : Csexp_rpc.Server.t Lazy.t
+    ; startup_ivar : (Csexp_rpc.Server.t, Exn_with_backtrace.t) result Fiber.Ivar.t
+    }
+
+  let create ~handler ~root ~where ~registry ~server =
+    let registry = Rpc_registry.create ~root ~where registry in
+    { handler; registry; server; startup_ivar = Fiber.Ivar.create () }
+  ;;
+
+  let print_uncaught_rpc_error exn =
+    Console.print [ Pp.text "Uncaught RPC Error"; Exn_with_backtrace.pp exn ]
+  ;;
+
+  let run t =
+    let with_print_errors f () =
+      Fiber.with_error_handler f ~on_error:(fun exn ->
+        print_uncaught_rpc_error exn;
+        Exn_with_backtrace.reraise exn)
+    in
+    let run () =
+      match Exn_with_backtrace.try_with (fun () -> Lazy.force t.server) with
+      | Error exn ->
+        Dune_trace.emit Rpc (fun () -> Dune_trace.Event.Rpc.startup_failure exn);
+        print_uncaught_rpc_error exn;
+        Fiber.Ivar.fill t.startup_ivar (Error exn)
+      | Ok server ->
+        let* () = Fiber.Ivar.fill t.startup_ivar (Ok server) in
+        let* sessions = Csexp_rpc.Server.serve server in
+        let () = Rpc_registry.register t.registry in
+        serve sessions t.handler
+    in
+    with_print_errors run
+    |> Fiber.finalize ~finally:(fun () ->
+      Rpc_registry.cleanup t.registry;
+      Fiber.return ())
+  ;;
+
+  let ready t =
+    Fiber.Ivar.read t.startup_ivar
+    >>= function
+    | Ok server ->
+      let+ () = Csexp_rpc.Server.ready server in
+      Ok ()
+    | Error _ as error -> Fiber.return error
+  ;;
+
+  let stop t =
+    Fiber.of_thunk (fun () ->
+      (* CR-soon rgrinberg: this is weird. we shouldn't be ignoring [stop] like this *)
+      match Fiber.Ivar.peek t.startup_ivar with
+      | None -> Fiber.return ()
+      | Some (Error _) -> Fiber.return ()
+      | Some (Ok server) -> Csexp_rpc.Server.stop server)
+  ;;
+end

src/rpc/server.mli

@@ -154,6 +154,24 @@ type t
 val make : 'a Handler.t -> t
 val serve : Csexp_rpc.Session.t Fiber.Stream.In.t -> t -> unit Fiber.t
 
+module Lifecycle : sig
+    type handler
+    type t
+
+    val create
+      :  handler:handler
+      -> root:string
+      -> where:Where.t
+      -> registry:[ `Add | `Skip ]
+      -> server:Csexp_rpc.Server.t Lazy.t
+      -> t
+
+    val run : t -> unit Fiber.t
+    val ready : t -> (unit, Exn_with_backtrace.t) result Fiber.t
+    val stop : t -> unit Fiber.t
+  end
+  with type handler := t
+
 (** Test only things below *)
 
 module type Session = Server_intf.Session