fix: disallow multiple rpc servers

Signed-off-by: Rudi Grinberg <[email protected]>

ps-id: 491f09b0-a316-43b8-8df5-a05f53fa331a

Author: rgrinberg

Diff for: bin/build_cmd.ml

@@ -85,7 +85,12 @@ let run_build_command_poll_passive ~(common : Common.t) ~config ~request:_ :
   (* CR-someday aalekseyev: It would've been better to complain if [request] is
      non-empty, but we can't check that here because [request] is a function.*)
   let open Fiber.O in
-  let rpc = Common.rpc common in
+  let rpc =
+    match Common.rpc common with
+    | `Allow server -> server
+    | `Forbid_builds ->
+      Code_error.raise "rpc server must be allowed in passive mode" []
+  in
   Scheduler.go_with_rpc_server_and_console_status_reporting ~common ~config
     (fun () ->
       Scheduler.Run.poll_passive

Diff for: bin/common.ml

@@ -37,7 +37,7 @@ type t =
   ; build_dir : string
   ; no_print_directory : bool
   ; store_orig_src_dir : bool
-  ; rpc : Dune_rpc_impl.Server.t Lazy.t
+  ; rpc : [ `Allow of Dune_rpc_impl.Server.t Lazy.t | `Forbid_builds ]
   ; default_target : Arg.Dep.t (* For build & runtest only *)
   ; watch : Watch_mode_config.t
   ; print_metrics : bool
@@ -75,7 +75,12 @@ let default_target t = t.default_target
 
 let prefix_target t s = t.root.reach_from_root_prefix ^ s
 
-let rpc t = Lazy.force t.rpc
+let rpc t =
+  match t.rpc with
+  | `Forbid_builds -> `Forbid_builds
+  | `Allow rpc -> `Allow (Lazy.force rpc)
+
+let forbid_builds t = { t with rpc = `Forbid_builds }
 
 let stats t = t.stats
 
@@ -1012,7 +1017,16 @@ let term ~default_root_is_cwd =
         at_exit (fun () -> Dune_stats.close stats);
         stats)
   in
-  let rpc = lazy (Dune_rpc_impl.Server.create ~root:root.dir stats) in
+  let rpc =
+    `Allow
+      (lazy
+        (let registry =
+           match watch with
+           | Yes _ -> `Add
+           | No -> `Skip
+         in
+         Dune_rpc_impl.Server.create ~registry ~root:root.dir stats))
+  in
   if store_digest_preimage then Dune_engine.Reversible_digest.enable ();
   if print_metrics then (
     Memo.Perf_counters.enable ();

Diff for: bin/common.mli

@@ -4,7 +4,15 @@ val capture_outputs : t -> bool
 
 val root : t -> Workspace_root.t
 
-val rpc : t -> Dune_rpc_impl.Server.t
+val rpc :
+     t
+  -> [ `Allow of Dune_rpc_impl.Server.t
+       (** Will run rpc if in watch mode and acquire the build lock *)
+     | `Forbid_builds
+       (** Promise not to build anything. For now, this isn't checked *)
+     ]
+
+val forbid_builds : t -> t
 
 val stats : t -> Dune_stats.t option
 

Diff for: bin/import.ml

@@ -139,27 +139,47 @@ module Scheduler = struct
         (Constant
            (Pp.seq message (Pp.verbatim ", waiting for filesystem changes...")))
 
+  let rpc server =
+    { Dune_engine.Rpc.run = Dune_rpc_impl.Server.run server
+    ; stop = Dune_rpc_impl.Server.stop server
+    ; ready = Dune_rpc_impl.Server.ready server
+    }
+
   let go ~(common : Common.t) ~config:dune_config f =
     let stats = Common.stats common in
     let config =
       let insignificant_changes = Common.insignificant_changes common in
       Dune_config.for_scheduler dune_config stats ~insignificant_changes
         ~signal_watcher:`Yes
     in
+    let f =
+      match Common.rpc common with
+      | `Allow server ->
+        fun () -> Dune_engine.Rpc.with_background_rpc (rpc server) f
+      | `Forbid_builds -> f
+    in
     Run.go config ~on_event:(on_event dune_config) f
 
   let go_with_rpc_server_and_console_status_reporting ~(common : Common.t)
       ~config:dune_config run =
+    let server =
+      match Common.rpc common with
+      | `Allow server -> rpc server
+      | `Forbid_builds ->
+        Code_error.raise "rpc must be enabled in polling mode" []
+    in
     let stats = Common.stats common in
     let config =
       let insignificant_changes = Common.insignificant_changes common in
       Dune_config.for_scheduler dune_config stats ~insignificant_changes
         ~signal_watcher:`Yes
     in
     let file_watcher = Common.file_watcher common in
-    let rpc = Common.rpc common in
     let run () =
-      Fiber.fork_and_join_unit (fun () -> Dune_rpc_impl.Server.run rpc) run
+      let open Fiber.O in
+      Dune_engine.Rpc.with_background_rpc server @@ fun () ->
+      let* () = Dune_engine.Rpc.ensure_ready () in
+      run ()
     in
     Run.go config ~file_watcher ~on_event:(on_event dune_config) run
 end

Diff for: bin/ocaml_merlin.ml

@@ -207,7 +207,9 @@ let term =
              debugging purposes only and should not be considered as a stable \
              output.")
   in
-  let common = Common.set_print_directory common false in
+  let common =
+    Common.set_print_directory common false |> Common.forbid_builds
+  in
   let config = Common.init common ~log_file:No_log_file in
   Scheduler.go ~common ~config (fun () ->
       match dump_config with

Diff for: src/csexp_rpc/csexp_rpc.ml

@@ -227,9 +227,6 @@ module Server = struct
       }
 
     let create fd sockaddr ~backlog =
-      Unix.set_nonblock fd;
-      Unix.setsockopt fd Unix.SO_REUSEADDR true;
-      Socket.bind fd sockaddr;
       Unix.listen fd backlog;
       let r_interrupt_accept, w_interrupt_accept = Unix.pipe ~cloexec:true () in
       Unix.set_nonblock r_interrupt_accept;
@@ -268,6 +265,7 @@ module Server = struct
         [ `Init of Unix.file_descr | `Running of Transport.t | `Closed ]
     ; backlog : int
     ; sockaddr : Unix.sockaddr
+    ; ready : unit Fiber.Ivar.t
     }
 
   let create sockaddr ~backlog =
@@ -276,19 +274,27 @@ module Server = struct
         (Unix.domain_of_sockaddr sockaddr)
         Unix.SOCK_STREAM 0
     in
-    { sockaddr; backlog; state = `Init fd }
+    Unix.set_nonblock fd;
+    Unix.setsockopt fd Unix.SO_REUSEADDR true;
+    match Socket.bind fd sockaddr with
+    | exception Unix.Unix_error (EADDRINUSE, _, _) -> Error `Already_in_use
+    | () ->
+      Ok { sockaddr; backlog; state = `Init fd; ready = Fiber.Ivar.create () }
+
+  let ready t = Fiber.Ivar.read t.ready
 
   let serve (t : t) =
     let* async = Worker.create () in
     match t.state with
     | `Closed -> Code_error.raise "already closed" []
     | `Running _ -> Code_error.raise "already running" []
     | `Init fd ->
-      let+ transport =
+      let* transport =
         Worker.task_exn async ~f:(fun () ->
             Transport.create fd t.sockaddr ~backlog:t.backlog)
       in
       t.state <- `Running transport;
+      let+ () = Fiber.Ivar.fill t.ready () in
       let accept () =
         Worker.task async ~f:(fun () ->
             Transport.accept transport

Diff for: src/csexp_rpc/csexp_rpc.mli

@@ -48,7 +48,11 @@ module Server : sig
   (** RPC Server *)
   type t
 
-  val create : Unix.sockaddr -> backlog:int -> t
+  val create : Unix.sockaddr -> backlog:int -> (t, [ `Already_in_use ]) result
+
+  (** [ready t] returns a fiber that completes when clients can start connecting
+      to the server *)
+  val ready : t -> unit Fiber.t
 
   val stop : t -> unit
 

Diff for: src/dune_engine/dune_engine.ml

@@ -47,3 +47,4 @@ module Report_errors_config = Report_errors_config
 module Compound_user_error = Compound_user_error
 module Reflection = Reflection
 module No_io = No_io
+module Rpc = Rpc

Diff for: src/dune_engine/rpc.ml

@@ -0,0 +1,32 @@
+open Fiber.O
+
+type server =
+  { run : unit Fiber.t
+  ; stop : unit Fiber.t
+  ; ready : unit Fiber.t
+  }
+
+type t =
+  { server : server
+  ; pool : Fiber.Pool.t
+  }
+
+let t = Fiber.Var.create ()
+
+let with_background_rpc server f =
+  let pool = Fiber.Pool.create () in
+  Fiber.Var.set t { server; pool } (fun () ->
+      Fiber.fork_and_join_unit
+        (fun () -> Fiber.Pool.run pool)
+        (fun () -> Fiber.finalize f ~finally:(fun () -> Fiber.Pool.stop pool)))
+
+let ensure_ready () =
+  let* { server; pool } = Fiber.Var.get_exn t in
+  let* () = Fiber.Pool.task pool ~f:(fun () -> server.run) in
+  server.ready
+
+let stop () =
+  let* { server; pool } = Fiber.Var.get_exn t in
+  Fiber.fork_and_join_unit
+    (fun () -> Fiber.Pool.stop pool)
+    (fun () -> server.stop)

Diff for: src/dune_engine/rpc.mli

@@ -0,0 +1,11 @@
+type server =
+  { run : unit Fiber.t
+  ; stop : unit Fiber.t
+  ; ready : unit Fiber.t
+  }
+
+val with_background_rpc : server -> (unit -> 'a Fiber.t) -> 'a Fiber.t
+
+val ensure_ready : unit -> unit Fiber.t
+
+val stop : unit -> unit Fiber.t

Diff for: src/dune_rpc_impl/server.ml

@@ -35,39 +35,29 @@ module Config = struct
   type t =
     { handler : Dune_rpc_server.t
     ; pool : Fiber.Pool.t
-    ; backlog : int
     ; root : string
     ; where : Dune_rpc.Where.t
+    ; server : Csexp_rpc.Server.t
+    ; stats : Dune_stats.t option
+    ; registry : [ `Add | `Skip ]
     }
 end
 
 module Run = struct
   module Registry = Dune_rpc_private.Registry
   module Server = Dune_rpc_server.Make (Csexp_rpc.Session)
-
-  type t =
-    { server : Csexp_rpc.Server.t
-    ; handler : Dune_rpc_server.t
-    ; pool : Fiber.Pool.t
-    ; where : Dune_rpc.Where.t
-    ; stats : Dune_stats.t option
-    ; root : string
-    }
+  open Config
 
   let t_var : t Fiber.Var.t = Fiber.Var.create ()
 
-  let of_config { Config.handler; backlog; pool; root; where } stats =
-    let () =
-      let socket_file = Where.rpc_socket_file () in
-      Path.mkdir_p (Path.build (Path.Build.parent_exn socket_file));
-      at_exit (fun () -> Path.Build.unlink_no_err socket_file)
-    in
-    let server = Csexp_rpc.Server.create (Where.to_socket where) ~backlog in
-    { server; handler; stats; pool; root; where }
-
   let run t =
     Fiber.Var.set t_var t @@ fun () ->
     let cleanup_registry = ref None in
+    let with_registry f =
+      match t.registry with
+      | `Skip -> ()
+      | `Add -> f ()
+    in
     let run_cleanup_registry () =
       match !cleanup_registry with
       | None -> ()
@@ -86,6 +76,7 @@ module Run = struct
         (fun () ->
           let* sessions = Csexp_rpc.Server.serve t.server in
           let () =
+            with_registry @@ fun () ->
             let (`Caller_should_write { Registry.File.path; contents }) =
               let registry_config =
                 Registry.Config.create (Lazy.force Dune_util.xdg)
@@ -117,21 +108,10 @@ module Run = struct
         (fun () -> Fiber.Pool.run t.pool)
     in
     Fiber.finalize (with_print_errors run) ~finally:(fun () ->
-        run_cleanup_registry ();
+        with_registry run_cleanup_registry;
         Fiber.return ())
-
-  let stop () =
-    let open Fiber.O in
-    let* t = Fiber.Var.get t_var in
-    match t with
-    | None -> Code_error.raise "rpc not running" []
-    | Some s ->
-      Csexp_rpc.Server.stop s.server;
-      Fiber.return ()
 end
 
-let stop = Run.stop
-
 type pending_build_action =
   | Build of Dep_conf.t list * Build_outcome.t Fiber.Ivar.t
 
@@ -232,9 +212,16 @@ type t =
   ; pending_build_jobs :
       (Dep_conf.t list * Build_outcome.t Fiber.Ivar.t) Job_queue.t
   ; mutable clients : Clients.t
-  ; stats : Dune_stats.t option
   }
 
+let ready (t : t) =
+  let* () = Fiber.return () in
+  Csexp_rpc.Server.ready t.config.server
+
+let stop (t : t) =
+  let+ () = Fiber.return () in
+  Csexp_rpc.Server.stop t.config.server
+
 let handler (t : t Fdecl.t) : 'a Dune_rpc_server.Handler.t =
   let on_init session (_ : Initialize.Request.t) =
     let t = Fdecl.get t in
@@ -339,7 +326,9 @@ let handler (t : t Fdecl.t) : 'a Dune_rpc_server.Handler.t =
               ~f:(fun (_, entry) -> Session.Stage1.request_close entry.session))
       in
       let shutdown () =
-        Fiber.fork_and_join_unit Dune_engine.Scheduler.shutdown Run.stop
+        Fiber.fork_and_join_unit Dune_engine.Scheduler.shutdown (fun () ->
+            Csexp_rpc.Server.stop t.config.server;
+            Fiber.return ())
       in
       Fiber.fork_and_join_unit terminate_sessions shutdown
     in
@@ -411,24 +400,39 @@ let handler (t : t Fdecl.t) : 'a Dune_rpc_server.Handler.t =
   in
   rpc
 
-let create ~root stats =
+let create ~registry ~root stats =
   let t = Fdecl.create Dyn.opaque in
   let pending_build_jobs = Job_queue.create () in
   let handler = Dune_rpc_server.make (handler t) in
   let pool = Fiber.Pool.create () in
   let where = Where.default () in
-  let config = { Config.handler; backlog = 10; pool; root; where } in
-  let res = { config; pending_build_jobs; clients = Clients.empty; stats } in
+  let server =
+    let socket_file = Where.rpc_socket_file () in
+    Path.mkdir_p (Path.build (Path.Build.parent_exn socket_file));
+    match Csexp_rpc.Server.create (Where.to_socket where) ~backlog:10 with
+    | Ok s ->
+      at_exit (fun () -> Path.Build.unlink_no_err socket_file);
+      s
+    | Error `Already_in_use ->
+      User_error.raise
+        [ Pp.textf
+            "Dune already running in this build directory. If this is not the \
+             case, delete %s"
+            (Path.Build.to_string (Where.rpc_socket_file ()))
+        ]
+  in
+  let config = { Config.handler; pool; root; where; stats; server; registry } in
+  let res = { config; pending_build_jobs; clients = Clients.empty } in
   Fdecl.set t res;
   res
 
 let listening_address t = t.config.where
 
 let run t =
   let* () = Fiber.return () in
-  Run.run (Run.of_config t.config t.stats)
+  Run.run t.config
 
-let stats (t : t) = t.stats
+let stats (t : t) = t.config.stats
 
 let pending_build_action t =
   Job_queue.read t.pending_build_jobs