Source file dune_rpc_private.ml

open Import
module Conv = Conv
module Versioned = Versioned
module Menu = Menu
module Procedures = Procedures
module Where = Where
module Registry = Registry
include Types
include Exported_types
module Version_error = Versioned.Version_error
module Decl = Decl
module Sub = Sub

module type Fiber = Fiber_intf.S

module Public = struct
  module Request = struct
    type ('a, 'b) t = ('a, 'b) Decl.Request.witness

    let ping = Procedures.Public.ping.decl

    let diagnostics = Procedures.Public.diagnostics.decl

    let format_dune_file = Procedures.Public.format_dune_file.decl

    let promote = Procedures.Public.promote.decl

    let build_dir = Procedures.Public.build_dir.decl
  end

  module Notification = struct
    type 'a t = 'a Decl.Notification.witness

    let shutdown = Procedures.Public.shutdown.decl
  end

  module Sub = struct
    type 'a t = 'a Sub.t

    let diagnostic = Sub.of_procedure Procedures.Poll.diagnostic

    let progress = Sub.of_procedure Procedures.Poll.progress

    let running_jobs = Sub.of_procedure Procedures.Poll.running_jobs
  end
end

module Server_notifications = struct
  let abort = Procedures.Server_side.abort.decl

  let log = Procedures.Server_side.log.decl
end

module Client = struct
  module type S = sig
    type t

    type 'a fiber

    type chan

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

      type 'a notification = 'a Versioned.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

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

        val write : t -> Sexp.t list option -> 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
        ; write : Sexp.t list option -> unit Fiber.t
        ; closed_read : bool
        ; mutable closed_write : bool
        ; disconnected : unit Fiber.Ivar.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
        { read
        ; write = (fun s -> Chan.write c s)
        ; closed_read = false
        ; closed_write = false
        ; disconnected
        }

      let write t s =
        let* () = Fiber.return () in
        match s with
        | Some _ -> t.write s
        | None ->
          if t.closed_write then Fiber.return ()
          else (
            t.closed_write <- true;
            t.write None)

      let read t =
        let* () = Fiber.return () in
        if t.closed_read then Fiber.return None else t.read ()
    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)

    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.write t.chan None)
          (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 send conn (packet : Packet.t list option) =
      let sexps =
        Option.map packet ~f:(List.map ~f:(Conv.to_sexp Packet.sexp))
      in
      Chan.write conn.chan sexps

    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* () = send conn (Some [ 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.Staged.request

      type 'a notification = 'a Versioned.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"
            ~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 ->
          send t (Some [ 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 <- [];
        send t.client (Some pending)
    end

    let read_packets t packets =
      let* () =
        Fiber.parallel_iter packets ~f:(function
          | Packet.Notification n -> (
            if
              String.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 = V.Handler.handle_request handler () (id, req) in
            send t (Some [ 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_dune_file;
      Builder.declare_request t Procedures.Public.promote;
      Builder.declare_request t Procedures.Public.build_dir;
      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 [ Dyn.String 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.write chan None)
      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
end