Source file client.ml

open Import
open Types
open Exported_types

module type S = sig
  type t
  type 'a fiber
  type chan

  module Versioned : sig
    type ('a, 'b) request = ('a, 'b) Versioned_intf.Staged.request
    type 'a notification = 'a Versioned_intf.Staged.notification

    val prepare_request
      :  t
      -> ('a, 'b) Decl.Request.witness
      -> (('a, 'b) request, Version_error.t) result fiber

    val prepare_notification
      :  t
      -> 'a Decl.Notification.witness
      -> ('a notification, Version_error.t) result fiber
  end

  val request
    :  ?id:Id.t
    -> t
    -> ('a, 'b) Versioned.request
    -> 'a
    -> ('b, Response.Error.t) result fiber

  val notification : t -> 'a Versioned.notification -> 'a -> unit fiber
  val disconnected : t -> unit fiber

  module Stream : sig
    type 'a t

    val cancel : _ t -> unit fiber
    val next : 'a t -> 'a option fiber
  end

  val poll : ?id:Id.t -> t -> 'a Sub.t -> ('a Stream.t, Version_error.t) result fiber

  module Batch : sig
    type client := t
    type t

    val create : client -> t

    val request
      :  ?id:Id.t
      -> t
      -> ('a, 'b) Versioned.request
      -> 'a
      -> ('b, Response.Error.t) result fiber

    val notification : t -> 'a Versioned.notification -> 'a -> unit
    val submit : t -> unit fiber
  end

  module Handler : sig
    type t

    val create
      :  ?log:(Message.t -> unit fiber)
      -> ?abort:(Message.t -> unit fiber)
      -> unit
      -> t
  end

  type proc =
    | Request : ('a, 'b) Decl.request -> proc
    | Notification : 'a Decl.notification -> proc
    | Poll : 'a Procedures.Poll.t -> proc
    | Handle_request : ('a, 'b) Decl.request * ('a -> 'b fiber) -> proc

  val connect_with_menu
    :  ?handler:Handler.t
    -> private_menu:proc list
    -> chan
    -> Initialize.Request.t
    -> f:(t -> 'a fiber)
    -> 'a fiber

  val connect
    :  ?handler:Handler.t
    -> chan
    -> Initialize.Request.t
    -> f:(t -> 'a fiber)
    -> 'a fiber
end

type abort =
  | Invalid_session of Conv.error
  | Server_aborted of Message.t

exception Abort of abort

let () =
  Printexc.register_printer (function
    | Abort error ->
      let dyn =
        match error with
        | Invalid_session e -> Dyn.variant "Invalid_session" [ Conv.dyn_of_error e ]
        | Server_aborted e ->
          Dyn.variant "Server_aborted" [ Sexp.to_dyn (Message.to_sexp_unversioned e) ]
      in
      Some (Dyn.to_string dyn)
    | _ -> None)
;;

module Make
    (Fiber : Fiber_intf.S)
    (Chan : sig
       type t

       val write : t -> Sexp.t list -> unit Fiber.t
       val close : t -> unit Fiber.t
       val read : t -> Sexp.t option Fiber.t
     end) =
struct
  open Fiber.O
  module V = Versioned.Make (Fiber)

  module Chan = struct
    type t =
      { read : unit -> Sexp.t option Fiber.t
      ; closed_read : bool
      ; mutable closed_write : bool
      ; disconnected : unit Fiber.Ivar.t
      ; chan : Chan.t
      }

    let of_chan c =
      let disconnected = Fiber.Ivar.create () in
      let read () =
        let* result = Chan.read c in
        match result with
        | None ->
          let+ () = Fiber.Ivar.fill disconnected () in
          None
        | _ -> Fiber.return result
      in
      { chan = c; read; closed_read = false; closed_write = false; disconnected }
    ;;

    let close t =
      let* () = Fiber.return () in
      if t.closed_write
      then Fiber.return ()
      else (
        t.closed_write <- true;
        Chan.close t.chan)
    ;;

    let write (t : t) packet =
      List.map ~f:(Conv.to_sexp Packet.sexp) packet |> Chan.write t.chan
    ;;

    let read t =
      let* () = Fiber.return () in
      if t.closed_read then Fiber.return None else Chan.read t.chan
    ;;
  end

  type t =
    { chan : Chan.t
    ; requests :
        ( Id.t
          , [ `Cancelled
            | `Pending of
                [ `Completed of Response.t | `Connection_dead | `Cancelled ] Fiber.Ivar.t
            ] )
          Table.t
    ; initialize : Initialize.Request.t
    ; mutable next_id : int
    ; mutable running : bool
    ; mutable handler_initialized : bool
    ; (* We need this field to be an Ivar to ensure that any typed
         communications are correctly versioned. The contract of the [Fiber]
         interface ensures that this will be filled before any user code is
         run. *)
      handler : unit V.Handler.t Fiber.t
    ; on_preemptive_abort : Message.t -> unit Fiber.t
    }

  (* When the client is terminated via this function, the session is
     considered to be dead without a way to recover. *)
  let terminate t =
    let* () = Fiber.return () in
    match t.running with
    | false -> Fiber.return ()
    | true ->
      t.running <- false;
      let ivars = ref [] in
      Table.filteri_inplace t.requests ~f:(fun ~key:_ ~data:ivar ->
        ivars := ivar :: !ivars;
        false);
      let ivars () =
        Fiber.return
          (match !ivars with
           | [] -> None
           | x :: xs ->
             ivars := xs;
             Some x)
      in
      Fiber.fork_and_join_unit
        (fun () -> Chan.close t.chan)
        (fun () ->
           Fiber.parallel_iter ivars ~f:(fun status ->
             match status with
             | `Cancelled -> Fiber.return ()
             | `Pending ivar -> Fiber.Ivar.fill ivar `Connection_dead))
  ;;

  let terminate_with_error t message info =
    Fiber.fork_and_join_unit
      (fun () -> terminate t)
      (fun () ->
         (* TODO stop using code error here. If [terminate_with_error] is
            called, it's because the other side is doing something unexpected,
            not because we have a bug *)
         Code_error.raise message info)
  ;;

  let create ~chan ~initialize ~handler ~on_preemptive_abort =
    let requests = Table.create (module Id) 16 in
    { chan
    ; requests
    ; next_id = 0
    ; initialize
    ; running = true
    ; handler_initialized = false
    ; handler
    ; on_preemptive_abort
    }
  ;;

  let prepare_request' conn (id, req) =
    match conn.running with
    | false ->
      let err =
        let payload =
          Sexp.record
            [ "id", Id.to_sexp id; "req", Conv.to_sexp (Conv.record Call.fields) req ]
        in
        Response.Error.create
          ~payload
          ~message:"request sent while connection is dead"
          ~kind:Connection_dead
          ()
      in
      Error err
    | true ->
      let ivar = Fiber.Ivar.create () in
      (match Table.add conn.requests id (`Pending ivar) with
       | Ok () -> ()
       | Error _ -> Code_error.raise "duplicate id" [ "id", Id.to_dyn id ]);
      Ok ivar
  ;;

  let request_untyped conn (id, req) =
    let* () = Fiber.return () in
    match prepare_request' conn (id, req) with
    | Error e -> Fiber.return (`Completed (Error e))
    | Ok ivar ->
      let* () = Chan.write conn.chan [ Request (id, req) ] in
      Fiber.Ivar.read ivar
  ;;

  let parse_response t decode = function
    | Error e -> Fiber.return (Error e)
    | Ok res ->
      (match decode res with
       | Ok s -> Fiber.return (Ok s)
       | Error e ->
         terminate_with_error
           t
           "response not matched by decl"
           [ "e", Response.Error.to_dyn e ])
  ;;

  let gen_id t = function
    | Some id -> id
    | None ->
      let id = Sexp.List [ Atom "auto"; Atom (Int.to_string t.next_id) ] in
      t.next_id <- t.next_id + 1;
      Id.make id
  ;;

  module Versioned = struct
    type ('a, 'b) request = ('a, 'b) Versioned_intf.Staged.request
    type 'a notification = 'a Versioned_intf.Staged.notification

    let prepare_request t (decl : _ Decl.Request.witness) =
      let+ handler = t.handler in
      V.Handler.prepare_request handler decl
    ;;

    let prepare_notification (type a) t (decl : a Decl.Notification.witness) =
      let+ handler = t.handler in
      V.Handler.prepare_notification handler decl
    ;;
  end

  let request t id ({ encode_req; decode_resp } : _ Versioned.request) req =
    let req = encode_req req in
    let* res = request_untyped t (id, req) in
    match res with
    | `Connection_dead -> Fiber.return `Connection_dead
    | `Cancelled -> Fiber.return `Cancelled
    | `Completed res ->
      let+ res = parse_response t decode_resp res in
      `Completed res
  ;;

  let cancel t id =
    match Table.find t.requests id with
    | None | Some `Cancelled -> Fiber.return ()
    | Some (`Pending ivar) ->
      Table.remove t.requests id;
      Fiber.Ivar.fill ivar `Cancelled
  ;;

  let make_notification
        (type a)
        t
        ({ encode } : a Versioned.notification)
        (n : a)
        (k : Call.t -> 'a)
    : 'a
    =
    let call = encode n in
    match t.running with
    | true -> k call
    | false ->
      let err =
        let payload = Conv.to_sexp (Conv.record Call.fields) call in
        Response.Error.create
          ~payload
          ~message:"notification sent while connection is dead"
            (* CR-soon rgrinberg: Should be Connection_dead *)
          ~kind:Code_error
          ()
      in
      raise (Response.Error.E err)
  ;;

  let notification (type a) t (stg : a Versioned.notification) (n : a) =
    let* () = Fiber.return () in
    make_notification t stg n (fun call -> Chan.write t.chan [ Notification call ])
  ;;

  let disconnected t = Fiber.Ivar.read t.chan.disconnected

  module Stream = struct
    type nonrec 'a t =
      { poll : (Id.t, 'a option) Versioned.request
      ; cancel : Id.t Versioned.notification
      ; client : t
      ; id : Id.t
      ; mutable pending_request_id : Id.t option
      ; counter : int
      ; mutable active : bool
      }

    let create sub client id =
      let+ handler = client.handler in
      let open Result.O in
      let+ poll = V.Handler.prepare_request handler (Sub.poll sub)
      and+ cancel = V.Handler.prepare_notification handler (Sub.poll_cancel sub) in
      { poll; cancel; client; id; pending_request_id = None; counter = 0; active = true }
    ;;

    let check_active t =
      if not t.active then Code_error.raise "polling is inactive" [ "id", Id.to_dyn t.id ]
    ;;

    let next t =
      let* () = Fiber.return () in
      check_active t;
      (match t.pending_request_id with
       | Some _ ->
         Code_error.raise "Poll.next: previous Poll.next did not terminate yet" []
       | None -> ());
      let id =
        Sexp.record [ "poll", Id.to_sexp t.id; "i", Sexp.Atom (string_of_int t.counter) ]
        |> Id.make
      in
      t.pending_request_id <- Some id;
      let+ res = request t.client id t.poll t.id in
      t.pending_request_id <- None;
      match res with
      | `Connection_dead | `Cancelled -> None
      | `Completed (Ok res) -> res
      | `Completed (Error e) ->
        (* cwong: Should this really be a raise? *)
        raise (Response.Error.E e)
    ;;

    let cancel t =
      let* () = Fiber.return () in
      check_active t;
      t.active <- false;
      (* XXX should we add a pool to stop waiting for the notification to
           reach the server? *)
      let notify () = notification t.client t.cancel t.id in
      match t.pending_request_id with
      | None -> notify ()
      | Some id -> Fiber.fork_and_join_unit (fun () -> cancel t.client id) notify
    ;;
  end

  let no_cancel_raise_connection_dead id = function
    | `Cancelled -> assert false
    | `Completed s -> s
    | `Connection_dead ->
      let payload = Sexp.record [ "id", Id.to_sexp id ] in
      let error =
        Response.Error.create
          ~kind:Connection_dead
          ~payload
          ~message:"Connection terminated. This request will never receive a response."
          ()
      in
      Error error
  ;;

  let request ?id t spec req =
    let id = gen_id t id in
    let+ res = request t id spec req in
    no_cancel_raise_connection_dead id res
  ;;

  let poll ?id client sub =
    let* () = Fiber.return () in
    let id = gen_id client id in
    Stream.create sub client id
  ;;

  module Batch = struct
    type nonrec t =
      { client : t
      ; mutable pending : Packet.t list
      }

    let create client = { client; pending = [] }

    let notification t n a =
      make_notification t.client n a (fun call ->
        t.pending <- Notification call :: t.pending)
    ;;

    let request
          (type a b)
          ?id
          t
          ({ encode_req; decode_resp } : (a, b) Versioned.request)
          (req : a)
      : (b, _) result Fiber.t
      =
      let* () = Fiber.return () in
      let id = gen_id t.client id in
      let call = encode_req req in
      let ivar = prepare_request' t.client (id, call) in
      match ivar with
      | Error e -> Fiber.return (Error e)
      | Ok ivar ->
        t.pending <- Packet.Request (id, call) :: t.pending;
        let* res = Fiber.Ivar.read ivar in
        (* currently impossible because there's no batching for polling and
             cancellation is only available for polled requests *)
        let res = no_cancel_raise_connection_dead id res in
        parse_response t.client decode_resp res
    ;;

    let submit t =
      let* () = Fiber.return () in
      let pending = List.rev t.pending in
      t.pending <- [];
      Chan.write t.client.chan pending
    ;;
  end

  let read_packets t packets =
    let* () =
      Fiber.parallel_iter packets ~f:(function
        | Packet.Notification n ->
          if
            Method.Name.equal n.method_ Procedures.Server_side.abort.decl.method_
            && not t.handler_initialized
          then (
            match
              Conv.of_sexp ~version:t.initialize.dune_version Message.sexp n.params
            with
            | Ok msg -> t.on_preemptive_abort msg
            | Error error ->
              terminate_with_error
                t
                "fatal: server aborted connection, but couldn't parse reason"
                [ "reason", Sexp.to_dyn n.params; "error", Conv.dyn_of_error error ])
          else
            let* handler = t.handler in
            let* result = V.Handler.handle_notification handler () n in
            (match result with
             | Error e ->
               terminate_with_error
                 t
                 "received bad notification from server"
                 [ "error", Response.Error.to_dyn e; "notification", Call.to_dyn n ]
             | Ok () -> Fiber.return ())
        | Request (id, req) ->
          let* handler = t.handler in
          let* result =
            Fiber.collect_errors (fun () -> V.Handler.handle_request handler () (id, req))
          in
          let result =
            match result with
            | Ok r -> r
            | Error [ Response.Error.E t ] -> Error t
            | Error exns ->
              let payload =
                Sexp.List
                  (List.map exns ~f:(fun exn -> Sexp.Atom (Printexc.to_string exn)))
              in
              Error
                (Response.Error.create
                   ~payload
                   ~kind:Code_error
                   ~message:"unexpected error"
                   ())
          in
          Chan.write t.chan [ Response (id, result) ]
        | Response (id, response) ->
          (match Table.find t.requests id with
           | Some status ->
             Table.remove t.requests id;
             (match status with
              | `Pending ivar -> Fiber.Ivar.fill ivar (`Completed response)
              | `Cancelled -> Fiber.return ())
           | None ->
             terminate_with_error
               t
               "unexpected response"
               [ "id", Id.to_dyn id; "response", Response.to_dyn response ]))
    in
    terminate t
  ;;

  module Handler = struct
    type nonrec t =
      { log : Message.t -> unit Fiber.t
      ; abort : Message.t -> unit Fiber.t
      }

    let log { Message.payload; message } =
      let+ () = Fiber.return () in
      match payload with
      | None -> Format.eprintf "%s@." message
      | Some payload -> Format.eprintf "%s: %s@." message (Sexp.to_string payload)
    ;;

    let abort m = raise (Abort (Server_aborted m))
    let default = { log; abort }

    let create ?log ?abort () =
      let t =
        let t = default in
        match log with
        | None -> t
        | Some log -> { t with log }
      in
      let t =
        match abort with
        | None -> t
        | Some abort -> { t with abort }
      in
      t
    ;;
  end

  type proc =
    | Request : ('a, 'b) Decl.request -> proc
    | Notification : 'a Decl.notification -> proc
    | Poll : 'a Procedures.Poll.t -> proc
    | Handle_request : ('a, 'b) Decl.request * ('a -> 'b Fiber.t) -> proc

  let setup_versioning ~private_menu ~(handler : Handler.t) =
    let module Builder = V.Builder in
    let t : unit Builder.t = Builder.create () in
    (* CR-soon cwong: It is a *huge* footgun that you have to remember to
         declare a request here, or via [private_menu], and there is no
         mechanism to warn you if you forget. The closest thing is either seeing
         that [dune rpc status] does not report the new procedure, or need to
         deal with the [Notification_error.t], which contains some good context,
         but very little to indicate this specific problem. *)
    Builder.declare_request t Procedures.Public.ping;
    Builder.declare_request t Procedures.Public.diagnostics;
    Builder.declare_request t Procedures.Poll.(poll running_jobs);
    Builder.declare_notification t Procedures.Public.shutdown;
    Builder.declare_request t Procedures.Public.format;
    Builder.declare_request t Procedures.Public.format_dune_file;
    Builder.declare_request t Procedures.Public.promote;
    Builder.declare_request t Procedures.Public.promote_many;
    Builder.declare_request t Procedures.Public.build_dir;
    Builder.declare_request t Procedures.Public.runtest;
    Builder.implement_notification t Procedures.Server_side.abort (fun () ->
      handler.abort);
    Builder.implement_notification t Procedures.Server_side.log (fun () -> handler.log);
    Builder.declare_request t Procedures.Poll.(poll diagnostic);
    Builder.declare_request t Procedures.Poll.(poll progress);
    Builder.declare_notification t Procedures.Poll.(cancel running_jobs);
    Builder.declare_notification t Procedures.Poll.(cancel diagnostic);
    Builder.declare_notification t Procedures.Poll.(cancel progress);
    List.iter private_menu ~f:(function
      | Handle_request (r, h) -> Builder.implement_request t r (fun () -> h)
      | Request r -> Builder.declare_request t r
      | Notification n -> Builder.declare_notification t n
      | Poll p ->
        Builder.declare_request t (Procedures.Poll.poll p);
        Builder.declare_notification t (Procedures.Poll.cancel p));
    t
  ;;

  let connect_raw
        chan
        (initialize : Initialize.Request.t)
        ~(private_menu : proc list)
        ~(handler : Handler.t)
        ~f
    =
    let packets () =
      let+ read = Chan.read chan in
      Option.map read ~f:(fun sexp ->
        match Conv.of_sexp Packet.sexp ~version:initialize.dune_version sexp with
        | Error e -> raise (Abort (Invalid_session e))
        | Ok message -> message)
    in
    let builder = setup_versioning ~handler ~private_menu in
    let handler_var = Fiber.Ivar.create () in
    let client =
      let on_preemptive_abort = handler.abort in
      let handler = Fiber.Ivar.read handler_var in
      create ~initialize ~chan ~handler ~on_preemptive_abort
    in
    let run () =
      let* init =
        let id = Id.make (List [ Atom "initialize" ]) in
        let initialize = Initialize.Request.to_call initialize in
        let+ res = request_untyped client (id, initialize) in
        no_cancel_raise_connection_dead id res
      in
      match init with
      | Error e -> raise (Response.Error.E e)
      | Ok csexp ->
        let* menu =
          match
            Conv.of_sexp ~version:initialize.dune_version Initialize.Response.sexp csexp
          with
          | Error e -> raise (Abort (Invalid_session e))
          | Ok _resp ->
            let id = Id.make (List [ Atom "version menu" ]) in
            let supported_versions =
              let request =
                Version_negotiation.Request.create
                  (V.Builder.registered_procedures builder)
              in
              Version_negotiation.Request.to_call request
            in
            let* resp = request_untyped client (id, supported_versions) in
            (* we don't allow cancelling negotiation *)
            (match no_cancel_raise_connection_dead id resp with
             | Error e -> raise (Response.Error.E e)
             | Ok sexp ->
               (match
                  Conv.of_sexp
                    ~version:initialize.dune_version
                    Version_negotiation.Response.sexp
                    sexp
                with
                | Error e -> raise (Abort (Invalid_session e))
                | Ok (Selected methods) ->
                  (match Menu.of_list methods with
                   | Ok m -> Fiber.return m
                   | Error (method_, a, b) ->
                     terminate_with_error
                       client
                       "server responded with invalid version menu"
                       [ ( "duplicated"
                         , Dyn.Tuple [ Method.Name.to_dyn method_; Dyn.Int a; Dyn.Int b ]
                         )
                       ])))
        in
        let handler =
          V.Builder.to_handler builder ~menu ~session_version:(fun () ->
            client.initialize.dune_version)
        in
        client.handler_initialized <- true;
        let* () = Fiber.Ivar.fill handler_var handler in
        Fiber.finalize (fun () -> f client) ~finally:(fun () -> Chan.close chan)
    in
    Fiber.fork_and_join_unit (fun () -> read_packets client packets) run
  ;;

  let connect_with_menu ?(handler = Handler.default) ~private_menu chan init ~f =
    connect_raw (Chan.of_chan chan) init ~handler ~private_menu ~f
  ;;

  let connect = connect_with_menu ~private_menu:[]
end