Use GADTs to do the same simplification to shutdown as to diagnostics

Signed-off-by: Ambre Austen Suhamy <[email protected]>

Author: ElectreAAS

bin/build.ml

@@ -199,12 +199,12 @@ let build =
       Scheduler.go_without_rpc_server ~common ~config (fun () ->
         let open Fiber.O in
         let+ build_outcome =
-          Rpc.Common.fire_request
+          Rpc.Common.fire_message
             ~name:"build"
             ~wait:true
             ~lock_held_by
             builder
-            Dune_rpc_impl.Decl.build
+            (Rpc.Common.Request Dune_rpc_impl.Decl.build)
             targets
         in
         Rpc.Common.wrap_build_outcome_exn ~print_on_success:true build_outcome)

bin/diagnostics.ml

@@ -10,11 +10,11 @@ let term =
   Rpc.Common.client_term builder (fun () ->
     let open Fiber.O in
     let+ errors =
-      Rpc.Common.fire_request
+      Rpc.Common.fire_message
         ~name:"diagnostics_cmd"
         ~wait:false
         builder
-        Dune_rpc_private.Procedures.Public.diagnostics
+        (Rpc.Common.Request Dune_rpc_private.Procedures.Public.diagnostics)
         ()
     in
     List.iter errors ~f:(fun err ->

bin/exec.ml

@@ -228,12 +228,12 @@ let build_prog_via_rpc_if_necessary ~dir ~no_rebuild builder lock_held_by prog =
         let targets = Rpc.Group.Build.prepare_targets [ target ] in
         let open Fiber.O in
         let+ build_outcome =
-          Rpc.Common.fire_request
+          Rpc.Common.fire_message
             ~name:"build"
             ~wait:true
             ~lock_held_by
             builder
-            Dune_rpc_impl.Decl.build
+            (Rpc.Common.Request Dune_rpc_impl.Decl.build)
             targets
         in
         Rpc.Common.wrap_build_outcome_exn ~print_on_success:false build_outcome)

bin/fmt.ml

@@ -48,13 +48,13 @@ let run_fmt_command ~common ~config ~preview builder =
     Scheduler.go_without_rpc_server
       ~common
       ~config
-      (Rpc.Common.fire_request
+      (Rpc.Common.fire_message
          ~name:"format"
          ~wait:true
          ~warn_forwarding:false
          ~lock_held_by
          builder
-         Dune_rpc.Procedures.Public.format)
+         (Rpc.Common.Request Dune_rpc.Procedures.Public.format))
 ;;
 
 let command =

bin/promotion.ml

@@ -65,12 +65,12 @@ module Apply = struct
       Scheduler.go_without_rpc_server ~common ~config (fun () ->
         let open Fiber.O in
         let+ build_outcome =
-          Rpc.Common.fire_request
+          Rpc.Common.fire_message
             ~name:"promote_many"
             ~wait:true
             ~lock_held_by
             builder
-            Dune_rpc_private.Procedures.Public.promote_many
+            (Rpc.Common.Request Dune_rpc_private.Procedures.Public.promote_many)
             files_to_promote
         in
         Rpc.Common.wrap_build_outcome_exn ~print_on_success:true build_outcome)

bin/rpc/rpc_build.ml

@@ -15,7 +15,12 @@ let term =
   @@ fun () ->
   let open Fiber.O in
   let+ response =
-    Rpc_common.fire_request ~name:"build" ~wait builder Dune_rpc_impl.Decl.build targets
+    Rpc_common.fire_message
+      ~name:"build"
+      ~wait
+      builder
+      (Rpc_common.Request Dune_rpc_impl.Decl.build)
+      targets
   in
   match response with
   | Success -> print_endline "Success"

bin/rpc/rpc_common.ml

@@ -25,17 +25,40 @@ let raise_rpc_error (e : Rpc_error.t) =
     ]
 ;;
 
+type ('a, 'b) message_kind =
+  | Request : ('a, 'b) Dune_rpc.Decl.request -> ('a, 'b) message_kind
+  | Notification : 'a Dune_rpc.Decl.notification -> ('a, unit) message_kind
+
 let request_exn client request arg =
   let* decl =
     Client.Versioned.prepare_request client (Dune_rpc.Decl.Request.witness request)
   in
   match decl with
-  | Ok decl ->
-    let+ res = Client.request client decl arg in
+  | Ok decl -> Client.request client decl arg
+  | Error e -> raise (Dune_rpc.Version_error.E e)
+;;
+
+let notify_exn client notification arg =
+  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 prepare_message_and_send : type b. Client.t -> ('a, b) message_kind -> 'a -> b Fiber.t
+  =
+  fun client message arg ->
+  match message with
+  | Notification witness -> notify_exn client witness arg
+  | Request witness ->
+    let+ res = request_exn client witness arg in
     (match res with
-     | Ok response -> response
+     | Ok (result : b) -> result
      | Error e -> raise_rpc_error e)
-  | Error e -> raise (Dune_rpc.Version_error.E e)
 ;;
 
 let client_term builder f =
@@ -86,13 +109,13 @@ let warn_ignore_arguments lock_held_by =
     ]
 ;;
 
-let fire_request
+let fire_message
       ~name
       ~wait
       ?(warn_forwarding = true)
       ?(lock_held_by = Dune_util.Global_lock.Lock_held_by.Unknown)
       builder
-      request
+      message
       arg
   =
   let* connection = establish_client_session ~wait in
@@ -101,7 +124,7 @@ let fire_request
   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)
+    ~f:(fun client -> prepare_message_and_send client message arg)
 ;;
 
 let wrap_build_outcome_exn ~print_on_success build_outcome =

bin/rpc/rpc_common.mli

@@ -7,32 +7,28 @@ 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
 
 (** Cmdliner argument for a wait flag. *)
 val wait_term : bool Cmdliner.Term.t
 
+type ('a, 'b) message_kind =
+  | Request : ('a, 'b) Dune_rpc.Decl.request -> ('a, 'b) message_kind
+  | Notification : 'a Dune_rpc.Decl.notification -> ('a, unit) message_kind
+
 (** Send a request 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_request
+val fire_message
   :  name:string
   -> wait:bool
   -> ?warn_forwarding:bool
   -> ?lock_held_by:Dune_util.Global_lock.Lock_held_by.t
   -> Common.Builder.t
-  -> ('a, 'b) Dune_rpc.Decl.request
+  -> ('a, 'b) message_kind
   -> 'a
   -> 'b Fiber.t
 

bin/rpc/rpc_ping.ml

@@ -16,11 +16,11 @@ let term =
   @@ fun () ->
   let open Fiber.O in
   let+ () =
-    Rpc_common.fire_request
+    Rpc_common.fire_message
       ~name:"ping_cmd"
       ~wait
       builder
-      Dune_rpc_private.Procedures.Public.ping
+      (Rpc_common.Request Dune_rpc_private.Procedures.Public.ping)
       ()
   in
   Console.print [ Pp.text "Server appears to be responding normally" ]

bin/runtest.ml

@@ -51,12 +51,12 @@ let runtest_term =
     Scheduler.go_without_rpc_server ~common ~config (fun () ->
       let open Fiber.O in
       let+ build_outcome =
-        Rpc.Common.fire_request
+        Rpc.Common.fire_message
           ~name:"runtest"
           ~wait:false
           ~lock_held_by
           builder
-          Dune_rpc.Procedures.Public.runtest
+          (Rpc.Common.Request Dune_rpc.Procedures.Public.runtest)
           dir_or_cram_test_paths
       in
       Rpc.Common.wrap_build_outcome_exn ~print_on_success:true build_outcome)