Merge pull request #12675 from ElectreAAS/push-puoskmvorrun

refactor: add fire_notification to rpc_common and use in shutdown

Author: Alizter

bin/rpc/rpc_common.ml

@@ -39,6 +39,18 @@ let request_exn client request arg =
   | Error e -> raise (Dune_rpc.Version_error.E e)
 ;;
 
+let notify_exn client notification arg =
+  let open Fiber.O in
+  let* res =
+    Client.Versioned.prepare_notification
+      client
+      (Dune_rpc.Decl.Notification.witness notification)
+  in
+  match res with
+  | Ok decl -> Client.notification client decl arg
+  | Error e -> raise (Dune_rpc.Version_error.E e)
+;;
+
 let client_term builder f =
   let builder = Common.Builder.forbid_builds builder in
   let builder = Common.Builder.disable_log_file builder in
@@ -97,6 +109,17 @@ let warn_ignore_arguments lock_held_by =
     ]
 ;;
 
+let should_warn ~warn_forwarding builder =
+  warn_forwarding && not (Common.Builder.equal builder Common.Builder.default)
+;;
+
+let send_request ~f connection name =
+  Dune_rpc_impl.Client.client
+    connection
+    (Dune_rpc.Initialize.Request.create ~id:(Dune_rpc.Id.make (Sexp.Atom name)))
+    ~f
+;;
+
 let fire_request
       ~name
       ~wait
@@ -108,12 +131,23 @@ let fire_request
   =
   let open Fiber.O in
   let* connection = establish_client_session ~wait in
-  if warn_forwarding && not (Common.Builder.equal builder Common.Builder.default)
-  then warn_ignore_arguments lock_held_by;
-  Dune_rpc_impl.Client.client
-    connection
-    (Dune_rpc.Initialize.Request.create ~id:(Dune_rpc.Id.make (Sexp.Atom name)))
-    ~f:(fun client -> request_exn client request arg)
+  if should_warn ~warn_forwarding builder then warn_ignore_arguments lock_held_by;
+  send_request connection name ~f:(fun client -> request_exn client request arg)
+;;
+
+let fire_notification
+      ~name
+      ~wait
+      ?(warn_forwarding = true)
+      ?(lock_held_by = Dune_util.Global_lock.Lock_held_by.Unknown)
+      builder
+      notification
+      arg
+  =
+  let open Fiber.O in
+  let* connection = establish_client_session ~wait in
+  if should_warn ~warn_forwarding builder then warn_ignore_arguments lock_held_by;
+  send_request connection name ~f:(fun client -> notify_exn client notification arg)
 ;;
 
 let wrap_build_outcome_exn ~print_on_success build_outcome =

bin/rpc/rpc_common.mli

@@ -7,14 +7,6 @@ val active_server_exn : unit -> Dune_rpc.Where.t
 (** Raise an RPC response error. *)
 val raise_rpc_error : Dune_rpc.Response.Error.t -> 'a
 
-(** Make a request and raise an exception if the preparation for the request
-    fails in any way. Returns an [Error] if the response errors. *)
-val request_exn
-  :  Dune_rpc_client.Client.t
-  -> ('a, 'b) Dune_rpc.Decl.request
-  -> 'a
-  -> 'b Fiber.t
-
 (** Cmdliner term for a generic RPC client. *)
 val client_term : Common.Builder.t -> (unit -> 'a Fiber.t) -> 'a
 
@@ -40,6 +32,21 @@ val fire_request
   -> 'a
   -> 'b Fiber.t
 
+(** Send a notification to the RPC server. If [wait], it will poll forever until a server is listening.
+    Should be scheduled by a scheduler that does not come with a RPC server on its own.
+
+    [warn_forwarding] defaults to true, warns the user that since a RPC server is running, some arguments are ignored.
+    [lock_held_by] defaults to [Unknown], is only used to allow error messages to print the PID. *)
+val fire_notification
+  :  name:string
+  -> wait:bool
+  -> ?warn_forwarding:bool
+  -> ?lock_held_by:Dune_util.Global_lock.Lock_held_by.t
+  -> Common.Builder.t
+  -> 'a Dune_rpc.Decl.notification
+  -> 'a
+  -> unit Fiber.t
+
 val wrap_build_outcome_exn
   :  print_on_success:bool
   -> Dune_rpc.Build_outcome_with_diagnostics.t

bin/shutdown.ml

@@ -1,28 +1,4 @@
 open Import
-module Client = Dune_rpc_client.Client
-
-let send_shutdown cli =
-  let open Fiber.O in
-  let* decl =
-    Client.Versioned.prepare_notification
-      cli
-      Dune_rpc_private.Public.Notification.shutdown
-  in
-  match decl with
-  | Ok decl -> Client.notification cli decl ()
-  | Error e -> raise (Dune_rpc_private.Version_error.E e)
-;;
-
-let exec () =
-  let open Fiber.O in
-  let where = Rpc.Rpc_common.active_server_exn () in
-  let* conn = Client.Connection.connect_exn where in
-  Dune_rpc_impl.Client.client
-    conn
-    ~f:send_shutdown
-    (Dune_rpc_private.Initialize.Request.create
-       ~id:(Dune_rpc_private.Id.make (Sexp.Atom "shutdown_cmd")))
-;;
 
 let info =
   let doc = "Cancel and shutdown any builds in the current workspace." in
@@ -31,7 +7,13 @@ let info =
 
 let term =
   let+ builder = Common.Builder.term in
-  Rpc.Rpc_common.client_term builder exec
+  Rpc.Rpc_common.client_term
+    builder
+    (Rpc.Rpc_common.fire_notification
+       ~name:"shutdown_cmd"
+       ~wait:false
+       builder
+       Dune_rpc_private.Procedures.Public.shutdown)
 ;;
 
 let command = Cmd.v info term