Refactor call and rpc client handlers

This makes sure that the rpc can send requests before having received
anything.

Author: Andrew Jeffery

lib/grpc-lwt/client.ml

@@ -21,17 +21,19 @@ let make_request ~scheme ~service ~rpc =
 
 let call ~service ~rpc ?(scheme = "https") ~handler ~do_request () =
   let request = make_request ~service ~rpc ~scheme in
-  let write_body, write_body_notify = Lwt.wait () in
-  let out, notify_out = Lwt.wait () in
+  let read_body, read_body_notify = Lwt.wait () in
+  let handler_res, handler_res_notify = Lwt.wait () in
+  let out, out_notify = Lwt.wait () in
   let response_handler (response : H2.Response.t) body =
+    Lwt.wakeup_later read_body_notify body;
     Lwt.async (fun () ->
         if response.status <> `OK then (
-          Lwt.wakeup_later notify_out
+          Lwt.wakeup_later out_notify
             (Error (Grpc.Status.v Grpc.Status.Unknown));
           Lwt.return_unit )
         else
-          let+ out = handler write_body body in
-          Lwt.wakeup_later notify_out (Ok out))
+          let+ handler_res = handler_res in
+          Lwt.wakeup_later out_notify (Ok handler_res))
   in
   let status, status_notify = Lwt.wait () in
   let trailers_handler headers =
@@ -50,23 +52,26 @@ let call ~service ~rpc ?(scheme = "https") ~handler ~do_request () =
         let status = Grpc.Status.v ?message code in
         Lwt.wakeup_later status_notify status
   in
-  let body =
+  let write_body =
     do_request ?trailers_handler:(Some trailers_handler) request
       ~response_handler
   in
-  Lwt.wakeup_later write_body_notify body;
+  Lwt.async (fun () ->
+      let+ handler_res = handler write_body read_body in
+      Lwt.wakeup_later handler_res_notify handler_res);
   let* out = out in
   let+ status = status in
   match out with Error _ as e -> e | Ok out -> Ok (out, status)
 
 module Rpc = struct
   type 'a handler =
-    [ `write ] H2.Body.t Lwt.t -> [ `read ] H2.Body.t -> 'a Lwt.t
+    [ `write ] H2.Body.t -> [ `read ] H2.Body.t Lwt.t -> 'a Lwt.t
 
   let bidirectional_streaming ~f write_body read_body =
-    let* write_body = write_body in
     let decoder_stream, decoder_push = Lwt_stream.create () in
-    Connection.grpc_recv_streaming read_body decoder_push;
+    Lwt.async (fun () ->
+        let+ read_body = read_body in
+        Connection.grpc_recv_streaming read_body decoder_push);
     let encoder_stream, encoder_push = Lwt_stream.create () in
     Lwt.async (fun () ->
         Connection.grpc_send_streaming_client write_body encoder_stream);

lib/grpc-lwt/client.mli

@@ -1,6 +1,6 @@
 module Rpc : sig
   type 'a handler =
-    [ `write ] H2.Body.t Lwt.t -> [ `read ] H2.Body.t -> 'a Lwt.t
+    [ `write ] H2.Body.t -> [ `read ] H2.Body.t Lwt.t -> 'a Lwt.t
 
   val bidirectional_streaming :
     f:((string -> unit) -> Grpc.Buffer.t Lwt_stream.t -> 'a Lwt.t) -> 'a handler