upload_queue.ml

(* Copyright (C) 2014, Thomas Leonard
 * See the README file for details. *)

open Lwt

type item = {
  size : int64;
  data : string Lwt_stream.t;
}

type add_error = [`Wrong_size of int64 | `Disk_full | `Unknown of exn]

let (++) = Int64.add
let (--) = Int64.sub
let (//) = Int64.div

cstruct index_entry {
  uint64 start_sector;
  uint64 length;
} as little_endian
let _ = hexdump_index_entry

module Make (B : V1_LWT.BLOCK) = struct
  module ReadBuffer = Block_buffer.Make(B)

  type file = {
    start_sector : int64;
    file_length : int64;
  }

  type t = {
    b : B.t;
    sector_size : int64;
    total_sectors : int64;
    mutable next_free_sector : int64;
    notify : unit Lwt_condition.t;
    queue : file Queue.t;
    max_queue_length : int;
    mutable n_writers : int;
  }

  (** Abort on error *)
  let (>>|=) x f =
    x >>= function
    | `Error _ -> failwith "Block device error"
    | `Ok v -> f v

  let sectors_needed q size =
    (size ++ q.sector_size -- 1L) // q.sector_size

  let page_size = Io_page.round_to_page_size 1
  let pages_needed size =
    (size + page_size - 1) / page_size

  let get_sector_cstruct q =
    let buffer = Io_page.get (pages_needed (Int64.to_int q.sector_size)) |> Io_page.to_cstruct in
    Cstruct.sub buffer 0 (Int64.to_int q.sector_size)

  let persist_queue q =
    let l = Queue.length q.queue in
    assert (l <= q.max_queue_length);
    let buffer = get_sector_cstruct q in

    let off = ref 0 in
    let add item =
      let entry = Cstruct.sub buffer !off sizeof_index_entry in
      set_index_entry_start_sector entry item.start_sector;
      set_index_entry_length entry item.file_length;
      off := !off + sizeof_index_entry in

    q.queue |> Queue.iter add;
    let end_of_list = {start_sector = 0L; file_length = 0L} in
    for _i = l to q.max_queue_length - 1 do
      add end_of_list;
    done;

    B.write q.b 0L [buffer] >>|= return

  let restore q =
    let buffer = get_sector_cstruct q in
    B.read q.b 0L [buffer] >>|= fun () ->

    let rec load i =
      if i = q.max_queue_length then return ()
      else (
        let off = i * sizeof_index_entry in
        let entry = Cstruct.sub buffer off sizeof_index_entry in
        let item = {
          start_sector = get_index_entry_start_sector entry;
          file_length = get_index_entry_length entry;
        } in
        if item.start_sector = 0L then return ()
        else if (item.start_sector < 1L ||
                 item.file_length < 0L ||
                 item.start_sector ++ sectors_needed q item.file_length > q.total_sectors) then (
          Log.info "Discarding invalid item at %Ld, length %Ld" item.start_sector item.file_length
        ) else (
          Log.info "Restored item at %Ld, length %Ld" item.start_sector item.file_length >>= fun () ->
          let next_free = item.start_sector ++ (sectors_needed q item.file_length) in
          if next_free > q.next_free_sector then q.next_free_sector <- next_free;
          Queue.add item q.queue;
          load (i + 1)
        )
      ) in
    load 0

  let create b =
    B.get_info b >>= fun info ->
    let max_queue_length = info.B.sector_size / sizeof_index_entry in
    let q = {
      b;
      sector_size = Int64.of_int info.B.sector_size;
      total_sectors = info.B.size_sectors;
      next_free_sector = 1L;
      notify = Lwt_condition.create ();
      queue = Queue.create ();
      max_queue_length;
      n_writers = 0;
    } in
    restore q >>= fun () ->
    return q

  let maybe_reset q =
    if Queue.is_empty q.queue && q.n_writers = 0 then (
      q.next_free_sector <- 1L;
      Log.info "Queue is empty and idle; resetting"
    ) else return ()

  module Upload = struct
    exception Wrong_size of int64
    exception Disk_full

    let add_as q {size; data} =
      let start_sector = q.next_free_sector in
      let next_file = start_sector ++ sectors_needed q size in
      if next_file > q.total_sectors then raise Disk_full;
      q.next_free_sector <- next_file;

      (* Stream data to file. *)
      let file_offset = ref 0L in

      let raw_buffer = Io_page.get 256 in
      let pages = Io_page.to_pages raw_buffer |> List.map Io_page.to_cstruct |> Array.of_list in
      let page_buffer = Io_page.to_cstruct raw_buffer in
      let page_buffer_used = ref 0 in

      let flush_page_buffer () =
        Log.info "flushing %d bytes to disk" !page_buffer_used >>= fun () ->
        let buffered_data = Array.sub pages 0 (pages_needed !page_buffer_used) |> Array.to_list in
        let next_offset = !file_offset ++ Int64.of_int !page_buffer_used in
        if next_offset > size then raise (Wrong_size next_offset);
        if Int64.rem !file_offset q.sector_size <> 0L then failwith "Non-aligned write!";
        let sector = start_sector ++ (!file_offset // q.sector_size) in
        B.write q.b sector buffered_data >>|= fun () ->
        file_offset := next_offset;
        page_buffer_used := 0;
        return () in

      let rec add_data src i =
        let src_remaining = String.length src - i in
        if src_remaining = 0 then return ()
        else (
          let page_buffer_free = Cstruct.len page_buffer - !page_buffer_used in
          let chunk_size = min page_buffer_free src_remaining in
          Cstruct.blit_from_string src i page_buffer !page_buffer_used chunk_size;
          page_buffer_used := !page_buffer_used + chunk_size;
          lwt () = if page_buffer_free = chunk_size then flush_page_buffer () else return () in
          add_data src (i + chunk_size)
        ) in

      data |> Lwt_stream.iter_s (fun data -> add_data data 0) >>=
      flush_page_buffer >>= fun () ->

      (* Check size is correct and flag file as complete. *)
      let actual_size = !file_offset in
      if actual_size = size then (
        Log.info "added file of size %Ld bytes" size >>= fun () ->
        if Queue.length q.queue >= q.max_queue_length then raise Disk_full;
        Queue.add {start_sector; file_length = size} q.queue;
        persist_queue q >>= fun () ->
        Lwt_condition.broadcast q.notify ();
        return (`Ok ())
      ) else (
        raise (Wrong_size actual_size)
      )

    let add q item =
      try_lwt
        q.n_writers <- q.n_writers + 1;
        Lwt.finalize
          (fun () -> add_as q item)
          (fun () -> q.n_writers <- q.n_writers - 1; maybe_reset q)
      with
      | Disk_full ->
          Log.info "disk full" >>= fun () ->
          return (`Error `Disk_full)
      | Wrong_size actual ->
          Log.info "wrong size: expected %Ld, got %Ld" item.size actual >>= fun () ->
          return (`Error (`Wrong_size actual))
      | ex ->
          Log.warn "failed to store upload: %s" (Printexc.to_string ex) >>= fun () ->
          return (`Error (`Unknown ex))
  end

  module Download = struct
    (** This interface is not safe with multiple clients, but add some minimal protection anyway. *)
    let mutex = Lwt_mutex.create ()

    let being_processed = ref None

    let send q item =
      lwt rb = ReadBuffer.create q.b
        ~qlen:5
        ~sectors_per_read:64
        ~start:item.start_sector
        ~len:(sectors_needed q item.file_length) in

      let next () =
        try_lwt
          if Some item <> !being_processed then failwith "Item deleted during read!";
          match ReadBuffer.read rb with
          | None -> return None
          | Some buf_t ->
              lwt buf = buf_t in
              let data = Cstruct.to_string buf in
              (* Log.info "Sending chunk '%s'" data >>= fun () -> *)
              return (Some data)
        with ex ->
          Log.info "read failed: %s" (Printexc.to_string ex) >>= fun () ->
          return None in
      return {
        size = item.file_length;
        data = Lwt_stream.from next
      }

    let peek q =
      Lwt_mutex.with_lock mutex (fun () ->
        match !being_processed with
        | Some item ->
            Log.info "resending head item to client" >>= fun () ->
            send q item      (* Maybe the client didn't get it; send it again *)
        | None ->
            Log.info "taking item from queue..." >>= fun () ->
            let rec loop () =
              let item =
                try Some (Queue.peek q.queue)
                with Queue.Empty -> None in
              match item with
              | None -> Lwt_condition.wait q.notify >>= loop
              | Some item ->
                  Log.info "returning item" >>= fun () ->
                  being_processed := Some item;
                  send q item
            in
            loop ()
      )

    let delete q =
      if !being_processed = None then (
        (* If the client wasn't sure we got their delete, they may retry.
         * Do this outside the mutex so we don't wait for uploads. *)
        Log.info "DELETE called but nothing pending"
      ) else if (Lwt_mutex.is_locked mutex) then (
        failwith "DELETE during another download operation!"
      ) else (
        Lwt_mutex.with_lock mutex (fun () ->
          Log.info "deleting head item" >>= fun () ->
          Queue.pop q.queue |> ignore;
          persist_queue q >>= fun () ->
          being_processed := None;
          maybe_reset q
        )
      )
  end
end