Source file xs_client_lwt.ml

(*
 * Copyright (C) Citrix Systems Inc.
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as published
 * by the Free Software Foundation; version 2.1 only. with the special
 * exception on linking described in file LICENSE.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Lesser General Public License for more details.
 *)

(** A multiplexing xenstore protocol client over a byte-level transport *)

open Lwt
open Xs_protocol

module type IO = sig
  type 'a t = 'a Lwt.t

  val return : 'a -> 'a t
  val ( >>= ) : 'a t -> ('a -> 'b t) -> 'b t

  type backend = [ `xen | `unix ]

  val backend : backend

  type channel

  val create : unit -> channel t
  val destroy : channel -> unit t
  val read : channel -> bytes -> int -> int -> int t
  val write : channel -> bytes -> int -> int -> unit t
end

module type S = sig
  type client

  val make : unit -> client Lwt.t
  val suspend : client -> unit Lwt.t
  val resume : client -> unit Lwt.t

  type handle

  val immediate : client -> (handle -> 'a Lwt.t) -> 'a Lwt.t
  val transaction : client -> (handle -> 'a Lwt.t) -> 'a Lwt.t
  val wait : client -> (handle -> 'a Lwt.t) -> 'a Lwt.t
  val directory : handle -> string -> string list Lwt.t
  val read : handle -> string -> string Lwt.t
  val write : handle -> string -> string -> unit Lwt.t
  val rm : handle -> string -> unit Lwt.t
  val mkdir : handle -> string -> unit Lwt.t
  val setperms : handle -> string -> Xs_protocol.ACL.t -> unit Lwt.t
  val debug : handle -> string list -> string list Lwt.t
  val getdomainpath : handle -> int -> string Lwt.t
  val watch : handle -> string -> Xs_protocol.Token.t -> unit Lwt.t
  val unwatch : handle -> string -> Xs_protocol.Token.t -> unit Lwt.t
  val introduce : handle -> int -> nativeint -> int -> unit Lwt.t
  val set_target : handle -> int -> int -> unit Lwt.t
end

let finally f g =
  Lwt.catch
    (fun () ->
      f () >>= fun result ->
      g () >>= fun () -> Lwt.return result)
    (fun e -> g () >>= fun () -> Lwt.fail e)

module StringSet = Xs_handle.StringSet

module Watcher = struct
  type t = {
      mutable paths : StringSet.t
          (* we never care about events or ordering, only paths *)
    ; mutable cancelling : bool (* we need to stop watching and clean up *)
    ; c : unit Lwt_condition.t
    ; m : Lwt_mutex.t
  }
  (** Someone who is watching paths is represented by one of these: *)

  let make () =
    {
      paths = StringSet.empty
    ; cancelling = false
    ; c = Lwt_condition.create ()
    ; m = Lwt_mutex.create ()
    }

  (** Register that a watched path has been changed *)
  let put (x : t) path =
    Lwt_mutex.with_lock x.m (fun () ->
        x.paths <- StringSet.add path x.paths;
        Lwt_condition.signal x.c ();
        return ())

  (** Return a set of modified paths, or an empty set if we're cancelling *)
  let get (x : t) =
    Lwt_mutex.with_lock x.m (fun () ->
        let rec loop () =
          if x.paths = StringSet.empty && not x.cancelling then
            Lwt_condition.wait ~mutex:x.m x.c >>= fun () -> loop ()
          else Lwt.return ()
        in
        loop () >>= fun () ->
        let results = x.paths in
        x.paths <- StringSet.empty;
        return results)

  (** Called to shutdown the watcher and trigger an orderly cleanup *)
  let cancel (x : t) =
    let (_ : unit Lwt.t) =
      Lwt_mutex.with_lock x.m (fun () ->
          x.cancelling <- true;
          Lwt_condition.signal x.c ();
          return ())
    in
    ()
end

exception Malformed_watch_event
exception Unexpected_rid of int32
exception Dispatcher_failed

module Client =
functor
  (IO : IO with type 'a t = 'a Lwt.t)
  ->
  struct
    module PS = PacketStream (IO)

    (* Represents a single acive connection to a server *)
    type client = {
        mutable transport : IO.channel
      ; ps : PS.stream
      ; rid_to_wakeup : (int32, Xs_protocol.t Lwt.u) Hashtbl.t
      ; mutable dispatcher_thread : unit Lwt.t
      ; mutable dispatcher_shutting_down : bool
      ; watchevents : (Token.t, Watcher.t) Hashtbl.t
      ; mutable suspended : bool
      ; suspended_m : Lwt_mutex.t
      ; suspended_c : unit Lwt_condition.t
    }

    (* The following values are only used if IO.backend = `xen. *)

    let client_cache = ref None
    (* The whole application must only use one xenstore client, which will
       multiplex all requests onto the same ring. *)

    let client_cache_m = Lwt_mutex.create ()
    (* Multiple threads will call 'make' in parallel. We must ensure only
       one client is created. *)

    let recv_one t =
      PS.recv t.ps >>= function Ok x -> return x | Exception e -> Lwt.fail e

    let send_one t = PS.send t.ps

    let handle_exn t e =
      Printf.fprintf stderr "Caught: %s\n%!" (Printexc.to_string e);
      (match e with
      | Xs_protocol.Response_parser_failed _ ->
          (* Lwt_io.hexdump Lwt_io.stderr x *)
          return ()
      | _ -> return ())
      >>= fun () ->
      t.dispatcher_shutting_down <- true;
      (* no more hashtable entries after this *)
      (* all blocking threads are failed with our exception *)
      Lwt_mutex.with_lock t.suspended_m (fun () ->
          Printf.fprintf stderr "Propagating exception to %d threads\n%!"
            (Hashtbl.length t.rid_to_wakeup);
          Hashtbl.iter (fun _ u -> Lwt.wakeup_later_exn u e) t.rid_to_wakeup;
          return ())
      >>= fun () -> Lwt.fail e

    let rec dispatcher t =
      Lwt.catch (fun () -> recv_one t) (handle_exn t) >>= fun pkt ->
      match get_ty pkt with
      | Op.Watchevent ->
          (match Unmarshal.list pkt with
          | Some [ path; token ] ->
              let token = Token.of_string token in
              (* We may get old watches: silently drop these *)
              if Hashtbl.mem t.watchevents token then
                Watcher.put (Hashtbl.find t.watchevents token) path
                >>= fun () -> dispatcher t
              else dispatcher t
          | _ -> handle_exn t Malformed_watch_event)
          >>= fun () -> dispatcher t
      | _ -> (
          let rid = get_rid pkt in
          Lwt_mutex.with_lock t.suspended_m (fun () ->
              if Hashtbl.mem t.rid_to_wakeup rid then
                return (Some (Hashtbl.find t.rid_to_wakeup rid))
              else return None)
          >>= function
          | None -> handle_exn t (Unexpected_rid rid)
          | Some thread ->
              Lwt.wakeup_later thread pkt;
              dispatcher t)

    let make_unsafe () =
      IO.create () >>= fun transport ->
      let t =
        {
          transport
        ; ps = PS.make transport
        ; rid_to_wakeup = Hashtbl.create 10
        ; dispatcher_thread = return ()
        ; dispatcher_shutting_down = false
        ; watchevents = Hashtbl.create 10
        ; suspended = false
        ; suspended_m = Lwt_mutex.create ()
        ; suspended_c = Lwt_condition.create ()
        }
      in
      t.dispatcher_thread <- dispatcher t;
      return t

    let make () =
      match IO.backend with
      | `unix -> make_unsafe ()
      | `xen ->
          Lwt_mutex.with_lock client_cache_m (fun () ->
              match !client_cache with
              | Some c -> return c
              | None ->
                  make_unsafe () >>= fun c ->
                  client_cache := Some c;
                  return c)

    let suspend t =
      Lwt_mutex.with_lock t.suspended_m (fun () ->
          t.suspended <- true;
          let rec loop () =
            if Hashtbl.length t.rid_to_wakeup > 0 then
              Lwt_condition.wait ~mutex:t.suspended_m t.suspended_c
              >>= fun () -> loop ()
            else Lwt.return ()
          in
          loop ())
      >>= fun () ->
      Hashtbl.iter (fun _ watcher -> Watcher.cancel watcher) t.watchevents;
      Lwt.cancel t.dispatcher_thread;
      return ()

    let resume_unsafe t =
      Lwt_mutex.with_lock t.suspended_m (fun () ->
          t.suspended <- false;
          t.dispatcher_shutting_down <- false;
          Lwt_condition.broadcast t.suspended_c ();
          return ())
      >>= fun () ->
      t.dispatcher_thread <- dispatcher t;
      return ()

    let resume t =
      match IO.backend with
      | `unix -> resume_unsafe t
      | `xen -> (
          match !client_cache with
          | None -> Lwt.return ()
          | Some c ->
              IO.create () >>= fun transport ->
              c.transport <- transport;
              resume_unsafe t)

    type handle = client Xs_handle.t

    let make_rid =
      let counter = ref 0l in
      fun () ->
        let result = !counter in
        counter := Int32.succ !counter;
        result

    let rpc hint h payload unmarshal =
      let open Xs_handle in
      let rid = make_rid () in
      let request = Request.print payload (get_tid h) rid in
      let t, u = wait () in
      let c = get_client h in
      if c.dispatcher_shutting_down then Lwt.fail Dispatcher_failed
      else
        Lwt_mutex.with_lock c.suspended_m (fun () ->
            let rec loop () =
              if c.suspended then
                Lwt_condition.wait ~mutex:c.suspended_m c.suspended_c
                >>= fun () -> loop ()
              else Lwt.return ()
            in
            loop () >>= fun () ->
            Hashtbl.add c.rid_to_wakeup rid u;
            send_one c request)
        >>= fun () ->
        t >>= fun res ->
        Lwt_mutex.with_lock c.suspended_m (fun () ->
            Hashtbl.remove c.rid_to_wakeup rid;
            Lwt_condition.broadcast c.suspended_c ();
            return ())
        >>= fun () -> return (response hint request res unmarshal)

    let directory h path =
      rpc "directory"
        (Xs_handle.accessed_path h path)
        Request.(PathOp (path, Directory))
        Unmarshal.list

    let read h path =
      rpc "read"
        (Xs_handle.accessed_path h path)
        Request.(PathOp (path, Read))
        Unmarshal.string

    let write h path data =
      rpc "write"
        (Xs_handle.accessed_path h path)
        Request.(PathOp (path, Write data))
        Unmarshal.ok

    let rm h path =
      rpc "rm"
        (Xs_handle.accessed_path h path)
        Request.(PathOp (path, Rm))
        Unmarshal.ok

    let mkdir h path =
      rpc "mkdir"
        (Xs_handle.accessed_path h path)
        Request.(PathOp (path, Mkdir))
        Unmarshal.ok

    let setperms h path acl =
      rpc "setperms"
        (Xs_handle.accessed_path h path)
        Request.(PathOp (path, Setperms acl))
        Unmarshal.ok

    let debug h cmd_args = rpc "debug" h (Request.Debug cmd_args) Unmarshal.list

    let getdomainpath h domid =
      rpc "getdomainpath" h (Request.Getdomainpath domid) Unmarshal.string

    let watch h path token =
      rpc "watch" (Xs_handle.watch h path)
        (Request.Watch (path, Token.to_string token))
        Unmarshal.ok

    let unwatch h path token =
      rpc "unwatch" (Xs_handle.unwatch h path)
        (Request.Unwatch (path, Token.to_string token))
        Unmarshal.ok

    let introduce h domid store_mfn store_port =
      rpc "introduce" h
        (Request.Introduce (domid, store_mfn, store_port))
        Unmarshal.ok

    let set_target h stubdom_domid domid =
      rpc "set_target" h
        (Request.Set_target (stubdom_domid, domid))
        Unmarshal.ok

    let immediate client f = f (Xs_handle.no_transaction client)
    let counter = ref 0l

    let rec add_watches h token = function
      | [] -> Lwt.return_unit
      | p :: ps ->
          Lwt.try_bind
            (fun () -> watch h p token)
            (fun () -> add_watches h token ps)
            (fun ex ->
              (* If we fail to add all the watches (e.g. because we exceeded our watch quota)
                 then mark the remaining paths as unwatched so we don't try to release them. *)
              p :: ps
              |> List.iter (fun p ->
                     let _ : _ Xs_handle.t = Xs_handle.unwatch h p in
                     ());
              Lwt.fail ex)

    let wait client f =
      let open StringSet in
      counter := Int32.succ !counter;
      let token =
        Token.of_string (Printf.sprintf "%ld:xs_client.wait" !counter)
      in
      (* When we register the 'watcher', the dispatcher thread will signal us when
         watches arrive. *)
      let watcher = Watcher.make () in
      Hashtbl.add client.watchevents token watcher;

      (* We signal the caller via this cancellable task: *)
      let result, wakener = Lwt.task () in
      on_cancel result (fun () ->
          (* Trigger an orderly cleanup in the background: *)
          Watcher.cancel watcher);
      let h = Xs_handle.watching client in
      (* Adjust the paths we're watching (if necessary) and block (if possible) *)
      let adjust_paths () =
        let current_paths = Xs_handle.get_watched_paths h in
        (* Paths which weren't read don't need to be watched: *)
        let old_paths = diff current_paths (Xs_handle.get_accessed_paths h) in
        Lwt_list.iter_s (fun p -> unwatch h p token) (elements old_paths)
        >>= fun () ->
        (* Paths which were read do need to be watched: *)
        let new_paths = diff (Xs_handle.get_accessed_paths h) current_paths in
        add_watches h token (elements new_paths) >>= fun () ->
        (* If we're watching the correct set of paths already then just block *)
        if old_paths = empty && new_paths = empty then
          Watcher.get watcher >>= fun results ->
          (* an empty results set means we've been cancelled: trigger cleanup *)
          if results = empty then fail (Failure "goodnight") else return ()
        else return ()
      in
      (* Main client loop: *)
      let rec loop () =
        Lwt.catch
          (fun () ->
            f h >>= fun result ->
            wakeup wakener result;
            return true)
          (function
            | Eagain -> return false
            | ex ->
                wakeup_exn wakener ex;
                return true)
        >>= function
        | true -> return ()
        | false ->
            Lwt.try_bind adjust_paths loop (fun ex ->
                wakeup_exn wakener ex;
                Lwt.return_unit)
      in
      Lwt.async (fun () ->
          finally loop (fun () ->
              let current_paths = Xs_handle.get_watched_paths h in
              Lwt_list.iter_s
                (fun p -> unwatch h p token)
                (elements current_paths)
              >>= fun () ->
              Hashtbl.remove client.watchevents token;
              return ()));
      result

    let rec transaction client f =
      rpc "transaction_start"
        (Xs_handle.no_transaction client)
        Request.Transaction_start Unmarshal.int32
      >>= fun tid ->
      let h = Xs_handle.transaction client tid in
      f h >>= fun result ->
      Lwt.catch
        (fun () ->
          rpc "transaction_end" h (Request.Transaction_end true)
            Unmarshal.string
          >>= fun res' ->
          if res' = "OK" then return result
          else
            Lwt.fail
              (Error (Printf.sprintf "Unexpected transaction result: %s" res')))
        (function Eagain -> transaction client f | e -> Lwt.fail e)
  end