Source file storage.ml

(*****************************************************************************)
(*                                                                           *)
(* Open Source License                                                       *)
(* Copyright (c) 2019,2020 DaiLambda, Inc. <contact@dailambda.jp>            *)
(*                                                                           *)
(* Permission is hereby granted, free of charge, to any person obtaining a   *)
(* copy of this software and associated documentation files (the "Software"),*)
(* to deal in the Software without restriction, including without limitation *)
(* the rights to use, copy, modify, merge, publish, distribute, sublicense,  *)
(* and/or sell copies of the Software, and to permit persons to whom the     *)
(* Software is furnished to do so, subject to the following conditions:      *)
(*                                                                           *)
(* The above copyright notice and this permission notice shall be included   *)
(* in all copies or substantial portions of the Software.                    *)
(*                                                                           *)
(* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR*)
(* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,  *)
(* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL   *)
(* THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER*)
(* LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING   *)
(* FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER       *)
(* DEALINGS IN THE SOFTWARE.                                                 *)
(*                                                                           *)
(*****************************************************************************)

(*
    +0
    |0                   19|20         23|24        27|28        31|
    |HEADER STRING         |<bytes/cell->|<-max idx ->|<-version ->|

    Header for disk sync #1

    +32
    |0               11|12             23|24        27|28        31|
    |<-     hash     ->|<-      0      ->|<- i root ->|<- i next ->|

    Header for disk sync #2

    +64
    |0               11|12             23|24        27|28        31|
    |<-     hash     ->|<-      0      ->|<- i root ->|<- i next ->|

    Header for process sync #1

    +96
    |0               11|12             23|24        27|28        31|
    |<-     hash     ->|<- hash w/ key ->|<- i root ->|<- i next ->|

    Header for process sync #2

    +128
    |0               11|12             23|24        27|28        31|
    |<-     hash     ->|<- hash w/ key ->|<- i root ->|<- i next ->|

  Readers have 2 ways to get the latest header:

   Header for disk sync:
     This is updated less frequently, only when [flush] is called by the writer.
     Trustable any time (unless a non writer program destroys it).

   Header for process sync:
     This is updated regularly when [commit] is called by the writer.
     Trustable only when readers are sure that the writer updates it.
*)

open Utils
open Lwt.Syntax
open Result_lwt.Syntax

module Int64 = Stdint.Int64

module B = Mmap.Buffer

type mode = Reader | Writer

type config =
  { head_string : string;
    version : int;
    bytes_per_cell : int;
    max_index : Index.t;
  }

let check_head_string head_string =
  if String.length head_string <> 20 then
    invalid_arg (Printf.sprintf "invalid head_string size: %S" head_string)

let check_config config =
  check_head_string config.head_string;
  if config.bytes_per_cell < 4 then invalid_arg "too small bytes_pre_cell";
  if config.bytes_per_cell >= 0x10000 then invalid_arg "too large bytes_pre_cell";
  if not (0 <= config.version
          && config.version <= Stdint.Uint32.(to_int max_int)) then
    invalid_arg "invalid version";
  if Index.to_int config.max_index <= 0 then invalid_arg "invalid max_index"

type writer_key = string

type storage = {
  mutable array : Mmap.t ;
  (* mmaped array where the nodes are written and indexed. *)

  start_index : Index.t ;

  mutable current_length : Index.t ;
  (* Current length of the node table.
     The next index number to be used.

     current_length < mapped_length
  *)

  mutable mapped_length : Index.t ;

  mutable last_root_index : Index.t option ;
  (* The last commit entry position *)

  fd : Lwt_unix.file_descr ;
  (* File descriptor to the mapped file *)

  pos : int ;
  (* Position of the first cell in the file, in bytes
     It is always 0
  *)

  mode : mode ;

  version : int ;
  head_string : string ;
  bytes_per_cell : int ;
  max_index : Index.t ;

  (* How much space allocated for each resize, in cells *)
  resize_step : Index.t ;

  fn : string ;

  (* [true] if [fd] is already closed. *)
  mutable closed : bool ;

  (* The cells before this index (exclusive) are msynced *)
  mutable last_msynced_next_index : Index.t ;

  (* Random key string for process sync header.
     The writer process writes the process sync headers with this key.
     A reader requires the writer_key of the writer to read the process
     sync headers
  *)
  mutable writer_key : writer_key option
}

type t = storage

let set_last_root_index t x  = t.last_root_index <- x
let get_last_root_index t    = t.last_root_index

let set_current_length t x = t.current_length <- x
let get_current_length t = t.current_length

let size t = Stdint.Int64.(Index.to_int64 (get_current_length t ) * of_int t.bytes_per_cell)

let filename t = t.fn

let mode t = t.mode

let version t = t.version

let start_index t = t.start_index

let writer_key t = t.writer_key

(* msync takes long time.  The order of the flushing is not deterministic.
*)
let msync t =
  if t.mode = Writer then Mmap.msync_lwt t.array else Lwt.return_unit

let msync2 t ~off ~len =
  if t.mode = Writer then begin
    if len = 0 then Lwt.return_unit
    else
      let* () = Log.lwt_info "msync off:%a len:%d@." Index.pp off len in
      let off = Index.to_int off * t.bytes_per_cell + t.pos in
      let len = len * t.bytes_per_cell in
      Mmap.msync2_lwt t.array ~off ~len
  end
  else Lwt.return_unit

let do_madvise_read_ahead =
  let b = Envconf.madvise_read_ahead in
  if b then Log.fatal "Configuration: madvise(READ_AHEAD)";
  b

let make_array ~bytes_per_cell fd ~pos mapped_length =
  let res = Mmap.make fd ~pos ~shared:true (bytes_per_cell * mapped_length) in
  if not do_madvise_read_ahead then begin
    (* This resets all the mmap area. Inefficient *)
    let (), t = with_time (fun () -> Mmap.madvise_random res) in
    Log.debug "madvise len=%d %a" mapped_length Mtime.Span.pp t
  end;
  res

type Error.t += Index_overflow of t

let () = Error.register_printer (function
    | Index_overflow t -> Some ("Plebeia index overflow for " ^ t.fn)
    | _ -> None)

let resize required t =
  let open Index in
  assert (t.mode = Writer);
  assert (required > t.mapped_length);
  let new_mapped_length =
    ((required - t.mapped_length) / t.resize_step + Index.one) * t.resize_step  + t.mapped_length
  in
  if new_mapped_length < t.mapped_length || t.max_index < new_mapped_length
  then begin
    Log.notice "Storage: size overflow!";
    Error (Index_overflow t)
  end else begin
    Log.notice "Storage: resizing to %Ld" (Index.to_int64 new_mapped_length);
    let (), sec = with_time (fun () ->
        let array =
          make_array ~bytes_per_cell:t.bytes_per_cell
            (Lwt_unix.unix_file_descr t.fd) ~pos:t.pos
            (Index.to_int new_mapped_length)
        in
        t.array <- array;
        t.mapped_length <- new_mapped_length)
    in
    Log.notice "Storage: resized to %Ld in %a" (Index.to_int64 new_mapped_length) Mtime.Span.pp sec;
    Ok ()
  end

let may_resize required t =
  if t.mapped_length < required then
    resize required t
  else Ok ()

(* Access *)

let get_bytes t i n =
  let i = Index.to_int i in
  Mmap.get_buffer ~off:(i * t.bytes_per_cell) ~len:n t.array

let get_cell t i = get_bytes t i t.bytes_per_cell
let get_cell2 t i = get_bytes t i (t.bytes_per_cell * 2)

let new_index c =
  let i = c.current_length in
  let i' = Index.succ i in
  if i' < i || c.max_index < i' then Error (Index_overflow c)
  else begin
    c.current_length <- i';
    match may_resize i' c with
    | Ok () -> Ok i
    | Error _ as e ->
        (* We must recover the origianl current_length *)
        c.current_length <- i;
        e
  end

let new_indices =
  (* need to kick out too large [n] for the correctness of
     the overflow check below

     The maximum possible indices requested at one [new_indices] call
     is (65536 + 6) / bytes_per_cell + 1, the largest chunk.
     For bytes_per_cell = 32, it is 2049.
  *)
  fun c n ->
    let maximum_new_indices = 65536 + 6 / c.bytes_per_cell + 1 in
    if not (0 < n && n <= maximum_new_indices) then invalid_arg "Plebeia: new_indices: invalid size";
    let i = c.current_length in
    let i' = Index.(i + Unsafe.of_int n) in
    if i' < i || c.max_index < i' then Error (Index_overflow c)
    else begin
      c.current_length <- i';
      match may_resize i' c with
      | Ok () -> Ok i
      | Error _ as e ->
          (* We must recover the origianl current_length *)
          c.current_length <- i;
          e
    end

module Header = struct
  type Error.t += Header_broken

  let () =
    Error.register_printer (function
        | Header_broken -> Some "header broken"
        | _ -> None)

  type t =
    { last_next_index : Index.t
    ; last_root_index : Index.t option
    }

  let pp ppf t =
    Format.fprintf ppf "{ last_next_index= %a; last_root_index= %a }"
      Index.pp t.last_next_index
      (Format.option Index.pp) t.last_root_index

  module Hashfunc_12 = Hashfunc.Blake2B(struct let bytes = 12 end)

  let hash ~key string_24_31 =
    match (key : writer_key option) with
    | None -> Hashfunc_12.hash_string string_24_31
    | Some key -> Hashfunc_12.hash_strings [(key :> string); string_24_31]

  (* If Writer modifies the data about being read, it should be detected
     by a hash disagreement.
  *)
  let raw_read ~key array i =
    let cstr = Mmap.get_buffer ~off:(i * 32) ~len:32 array in
    let last_next_index = B.get_index cstr 28 in
    let last_root_index = Index.zero_then_none @@ B.get_index cstr 24 in
    let h = B.copy cstr 0 12 in
    let string_24_31 = B.copy cstr 24 8 in
    let h' = hash ~key:None string_24_31 in
    if h <> h' then Error (B.to_string cstr)
    else
      match key with
      | None -> Ok (`NoKey, { last_next_index ; last_root_index })
      | Some _ ->
          let h2 = B.copy cstr 12 12 in
          let h2' = hash ~key string_24_31 in
          if h2 <> h2' then Ok (`BadKey, { last_next_index ; last_root_index })
          else Ok (`GoodKey, { last_next_index ; last_root_index })

  let with_lock fd m f =
    let* _ = Lwt_unix.lseek fd 0 SEEK_SET in
    let* () = Lwt_unix.lockf fd m 256 in
    match f () with
    | exception e ->
        let* () = Lwt_unix.lockf fd F_ULOCK 256 in
        Lwt.fail e
    | res ->
        let+ () = Lwt_unix.lockf fd F_ULOCK 256 in
        res

  (* If the Writer has crashed during the double writes to the disk,
     one of the header must be valid.  The Readers can load a valid header.

     The header is lock protected.  When the Writer is modifying the header
     the Readers cannot access it, and vice versa.
  *)
  let read_array ~key i fd array =
    with_lock fd F_RLOCK @@ fun () ->
    match raw_read ~key array i with
    | Ok x -> Some x
    | Error _ -> (* something wrong in +32 *)
        match raw_read ~key array (i+1) with
        | Ok x -> Some x
        | Error _ -> None (* Header is corrupted. *)

  let read_array_disk_sync = read_array ~key:None 1
  let read_array_process_sync ~key = read_array ~key:(Some key) 3

  let read ~key off t = read_array ~key off t.fd t.array

  let read_disk_sync t =
    let+ res = read 1 ~key:None t in
    Option.map snd res

  let read_process_sync t = read 3 ~key:t.writer_key t

  let check_writer_key t key =
    let+ res = read 3 ~key:(Some key) t in
    match res with
    | None -> Error Header_broken
    | Some (`NoKey, _) -> assert false
    | Some (`GoodKey, _) -> Ok `GoodKey
    | Some (`BadKey, _) -> Ok `BadKey

  let raw_write t i src =
    let cstr = Mmap.get_buffer ~off:(i * 32) ~len:32 t.array in
    (* This is not atomic, therefore the header can be corrupted
       if the writer is killed during this write. *)
    Mmap.Buffer.blit_from_string src 0 cstr 0 32

  let write ~key i t h =
    if t.mode = Reader then invalid_arg "Reader cannot write";
    let cstr = Mmap.Buffer.create 32 in
    B.set_index cstr 28 h.last_next_index;
    B.set_index cstr 24 (Option.default Index.zero h.last_root_index);
    let string_24_31 = B.copy cstr 24 8 in
    let h = hash ~key:None string_24_31 in
    B.write_string h cstr 0;
    let h' = hash ~key string_24_31 in
    B.write_string h' cstr 12;
    let s = B.to_string cstr in
    (* One of the double writes must be valid on the memory
       even if the program crashes. *)
    with_lock t.fd F_LOCK @@ fun () ->
      raw_write t i s;
      raw_write t (i+1) s

  let write_disk_sync = write 1 ~key:None

  let write_process_sync t =
    assert (t.writer_key <> None);
    write 3 ~key:t.writer_key t

  (* only update the header for process sync *)
  let commit t =
    let cp = { last_next_index = t.current_length (* XXX inconsistent names ... *)
             ; last_root_index = t.last_root_index }
    in
    write_process_sync t cp

  let flush t =
    let* () = commit t in
    let cp = { last_next_index = t.current_length (* XXX inconsistent names ... *)
             ; last_root_index = t.last_root_index }
    in
    let msync () =
      let off = t.last_msynced_next_index in
      let len = Index.(to_int @@ t.current_length - t.last_msynced_next_index) in
      let off, len =
        if len < 0 then begin
          Log.warning "msync2 NEGATIVE %a , %d . Modifying existing cells!!!@." Index.pp off len;
          Index.zero, Index.to_int t.current_length
        end else off, len
      in
      let* () = msync2 t ~off ~len in
      t.last_msynced_next_index <- t.current_length;
      Lwt.return_unit
    in
    let* () = msync () in
    let* () = write_disk_sync t cp in
    (* XXX should precompute the cells for the header *)
    msync2 t ~off:Index.zero ~len:((256 + t.bytes_per_cell - 1) / t.bytes_per_cell)

  let write_config t =
    let cstr = Mmap.get_buffer ~off:0 ~len:32 t.array in
    B.blit_from_string t.head_string 0 cstr 0 20;
    B.set_uint32 cstr 20 t.bytes_per_cell;
    B.set_index cstr 24 t.max_index;
    B.set_uint32 cstr 28 t.version

  let read_config array : config =
    let cstr = Mmap.get_buffer ~off:0 ~len:32 array in
    let head_string = B.copy cstr 0 20 in
    let bytes_per_cell = B.get_uint32 cstr 20 in
    let max_index = B.get_index cstr 24 in
    let version = B.get_uint32 cstr 28 in
    { head_string; bytes_per_cell; max_index; version }
end

let check_writer_key = Header.check_writer_key

let make_writer_key () =
  let rng = Random.State.make_self_init () in
  String.init 32 (fun _ -> Char.chr @@ Random.State.int rng 256)

let make_new_writer_key ts =
  let rec loop () =
    let key = make_writer_key () in
    let rec check = function
      | [] -> Lwt.return true
      | t::ts ->
          let* res = check_writer_key t key in
          match res with
          | Error _ -> check ts
          | Ok `BadKey -> check ts
          | Ok `GoodKey -> Lwt.return false
    in
    let* res = check ts in
    if res then Lwt.return key
    else loop ()
  in
  loop ()

let commit = Header.commit

let flush = Header.flush

let get_start_index bytes_per_cell =
  (* We use 256 bytes to store the header *)
  let cells_for_header =
    (256 + bytes_per_cell - 1) / bytes_per_cell
  in
  Index.Unsafe.of_int cells_for_header

let default_resize_step_bytes =
  try
    let n = int_of_string @@ Sys.getenv "PLEBEIA_RESIZE_STEP_BYTES" in
    Log.fatal "Configuration: resize_step_bytes=%d" n;
    n
  with
  | Not_found -> 1 lsl 30 (* 1GB *)
  | _ -> Stdlib.failwith "PLEBEIA_RESIZE_STEP_BYTES=<int>"

let create
    ?length
    ?(resize_step_bytes=default_resize_step_bytes)
    ~config
    ~key
    fn =
  let pos = 0 in
  let writer_key = key in
  check_config config;
  let* fd = Lwt_unix.openfile fn [O_CREAT; O_EXCL; O_RDWR] 0o644 in
  let resize_step = Index.Unsafe.of_int (resize_step_bytes / config.bytes_per_cell) in
  let mapped_length =
    match length with
    | None -> resize_step
    | Some i -> i
  in
  let array =
    make_array
      ~bytes_per_cell:config.bytes_per_cell
      (Lwt_unix.unix_file_descr fd)
      ~pos (Index.to_int mapped_length)
  in
  let start_index = get_start_index config.bytes_per_cell in
  let t =
    { array ;
      mapped_length ;
      start_index ;
      current_length = start_index ;
      last_root_index = None ;
      fd ;
      pos ;
      mode= Writer;
      fn ;
      closed = false ;
      version = config.version;
      head_string = config.head_string;
      bytes_per_cell = config.bytes_per_cell;
      max_index = config.max_index;
      resize_step;
      last_msynced_next_index = start_index ;
      writer_key = Some writer_key;
    }
  in
  Header.write_config t;
  let+ () = Header.flush t in
  t

let null =
  let fd = Lwt_unix.stdin in
  let array = Mmap.null in
  { array ;
    mapped_length = Index.zero ;
    start_index = Index.zero ;
    current_length = Index.zero ;
    last_root_index = None ;
    fd ;
    pos = 0;
    mode= Reader ;
    fn = "null";
    closed = true ;

    version = 0;
    head_string = "DUMMY";
    bytes_per_cell = 32;
    max_index = Index.max_int;
    resize_step = Index.one;
    last_msynced_next_index = Index.zero ;
    writer_key = None;
  }

let is_null t = t == null

let override_version t version =
  let t = { t with version } in
  Header.write_config t

let truncate ?length t =
  let* () = msync t in
  let* () = Lwt_unix.ftruncate t.fd t.pos in
  let mapped_length =
    match length with
    | None -> t.resize_step
    | Some i ->
        match Sys.int_size with
        | 31 | 32 -> Index.Unsafe.of_int i
        | 63 ->
            if i > Index.(Unsafe.to_int max_int) then failwithf "create: too large: %d@." i
            else Index.Unsafe.of_int i
        | _ -> assert false
  in
  let array =
    make_array ~bytes_per_cell:t.bytes_per_cell
      (Lwt_unix.unix_file_descr t.fd) ~pos:t.pos
      (Index.to_int mapped_length)
  in
  Header.write_config t;

  t.array <- array;
  t.mapped_length <- mapped_length;
  t.current_length <- t.start_index;
  t.last_root_index <- None ;
  Header.flush t

let load_config ~pos fn =
  let* fd = Lwt_unix.openfile fn [O_RDWR] 0o644 in
  let array =
    make_array ~bytes_per_cell:32 (Lwt_unix.unix_file_descr fd)
      ~pos:pos 1
  in
  let config = Header.read_config array in
  let+ () = Lwt_unix.close fd in
  config

type Error.t += No_such_file of string

let () = Error.register_printer (function
    | No_such_file s -> Some ("no such file: " ^ s)
    | _ -> None)

let open_existing_for_read fn =
  let mode = Reader in
  let writer_key = None in
  let pos = 0 in
  let* exists = Lwt_unix.file_exists fn in
  if not exists then Lwt.return_error (No_such_file fn)
  else
    let* config = load_config ~pos fn in
    check_config config;
    (* Even for Reader, fd must be writable for mmap *)
    let* fd = Lwt_unix.openfile fn [O_RDWR] 0o644 in
    let* st = Lwt_unix.LargeFile.fstat fd in
    let sz = Int64.sub st.Unix.LargeFile.st_size (Int64.of_int pos) in
    assert (Int64.rem sz (Int64.of_int config.bytes_per_cell) = 0L); (* XXX think about the garbage *)
    let cells = Int64.(sz / Int64.of_int config.bytes_per_cell) in
    if cells > Index.to_int64 config.max_index then assert false;
    let mapped_length = Index.of_int64 cells in
    let array =
      make_array
        ~bytes_per_cell:config.bytes_per_cell
        (Lwt_unix.unix_file_descr fd)
        ~pos (Index.to_int mapped_length)
    in
    let start_index = get_start_index config.bytes_per_cell in
    let resize_step =
      let resize_step_bytes = default_resize_step_bytes in
      Index.Unsafe.of_int (resize_step_bytes / config.bytes_per_cell)
    in
    let* res_disk = Header.read_array_disk_sync fd array in
    let*=? current_length, last_root_index =
      (* Readers ignore the process sync header unless [enable_process_sync]
         is called *)
      match res_disk with
      | None -> Lwt.return_error Header.Header_broken
      | Some (`NoKey, h_disk_sync) ->
          Lwt.return_ok (h_disk_sync.last_next_index, h_disk_sync.last_root_index)
      | Some _ -> assert false
    in
    let t =
      { array ;
        start_index ;
        mapped_length ;
        current_length;
        last_root_index;
        fd = fd ;
        pos ;
        mode ;
        fn ;
        version = config.version;
        head_string = config.head_string;
        bytes_per_cell = config.bytes_per_cell;
        max_index = config.max_index;
        resize_step ;
        closed = false ;
        last_msynced_next_index = current_length ;
        writer_key
      }
    in
    Lwt.return_ok (config, t)

let open_existing_for_write
    ?(resize_step_bytes=default_resize_step_bytes)
    ~key
    fn
  =
  let pos = 0 in
  let mode = Writer in
  let writer_key = key in
  let* exists = Lwt_unix.file_exists fn in
  if not exists then Lwt.return_none
  else
    let* config = load_config ~pos fn in
    check_config config;
    (* Even for Reader, fd must be writable for mmap *)
    let* fd = Lwt_unix.openfile fn [O_RDWR] 0o644 in
    let* st = Lwt_unix.LargeFile.fstat fd in
    let sz = Int64.sub st.Unix.LargeFile.st_size (Int64.of_int pos) in
    assert (Int64.rem sz (Int64.of_int config.bytes_per_cell) = 0L); (* XXX think about the garbage *)
    let cells = Int64.(sz / Int64.of_int config.bytes_per_cell) in
    if cells > Index.to_int64 config.max_index then assert false;
    let mapped_length = Index.of_int64 cells in
    let array =
      make_array
        ~bytes_per_cell:config.bytes_per_cell
        (Lwt_unix.unix_file_descr fd)
        ~pos (Index.to_int mapped_length)
    in
    let start_index = get_start_index config.bytes_per_cell in
    let resize_step = Index.Unsafe.of_int (resize_step_bytes / config.bytes_per_cell) in
    let* res_disk = Header.read_array_disk_sync fd array in
    let* writer_key, current_length, last_root_index =
      let+ res_process = Header.read_array_process_sync ~key:writer_key fd array in
      match res_disk, res_process with
      | None, _ ->
          (* No way to recover it, I am afraid *)
          failwithf "Failed to load header: the disk sync header is broken"
      | _, None ->
          (* We can still reset everything to res_disk. *)
          failwithf "Failed to load header: the process sync header is broken"
      | Some ((`BadKey | `GoodKey), _), _ -> assert false
      | Some (`NoKey, h_disk_sync), Some (_, h_process_sync) ->
          match
            compare h_disk_sync.Header.last_next_index h_process_sync.Header.last_next_index,
            compare h_disk_sync.Header.last_root_index h_process_sync.Header.last_root_index
          with
          | 0, 0 -> Some writer_key, h_disk_sync.last_next_index, h_disk_sync.last_root_index
          | 1, _ | _, 1 ->
              (* We can still reset everything to res_disk. *)
              failwithf "Header invariant is broken: the header for the process sync is older than the header for the disk sync"
          | _ ->
              let diff = Index.(-) h_process_sync.last_next_index h_disk_sync.last_next_index in
              Log.warning "A Plebeia writer crash was detected.  %a disk cells are skipped." Index.pp diff;
              Some writer_key, h_process_sync.last_next_index, h_disk_sync.last_root_index
    in
    let t =
      { array ;
        start_index ;
        mapped_length ;
        current_length;
        last_root_index;
        fd = fd ;
        pos ;
        mode ;
        fn ;
        version = config.version;
        head_string = config.head_string;
        bytes_per_cell = config.bytes_per_cell;
        max_index = config.max_index;
        resize_step ;
        closed = false ;
        last_msynced_next_index = current_length ;
        writer_key
      }
    in
    let+ () = Header.flush t in
    Some (config, t)

let open_for_write
    ?resize_step_bytes
    ~config
    ~key
    fn =
  let* res = open_existing_for_write ?resize_step_bytes ~key fn in
  match res with
  | Some (config', storage) ->
      if config <> config' then invalid_arg "Storage.open_for_write: config mismatch";
      Lwt.return storage
  | None ->
      create ?resize_step_bytes ~config ~key fn

(* XXX auto close when GC'ed *)
let close ({ fd ; mode ; closed ; _ } as t) =
  if closed then Lwt.return_unit
  else begin
    let* () =
      if mode <> Reader then Header.flush t
      else Lwt.return_unit
    in
    let+ () = Lwt_unix.close fd in
    t.closed <- true
  end

let reopen_with t h =
  (* We load the header first, before fstat the file.

     If we would do opposite, the following might happen:

     * Reader fstats
     * Writer extends the file, making the fstats obsolete
     * Write update the header, last_indices point out of the obsolete fstat
     * Reader reads the header
     * Reader fails to load the last_indices, since it is not mapped
  *)
  if h.Header.last_next_index = t.current_length
     && h.last_root_index = t.last_root_index then begin
    Lwt.return_unit
  end else
    let* st = Lwt_unix.LargeFile.fstat t.fd in
    let sz = Int64.sub st.Unix.LargeFile.st_size (Int64.of_int t.pos) in
    assert (Int64.rem sz (Int64.of_int t.bytes_per_cell) = 0L);  (* XXX think about the garbage *)
    let cells = Int64.(sz / (Int64.of_int t.bytes_per_cell)) in
    if cells > Index.to_int64 t.max_index then assert false;
    let mapped_length = Index.of_int64 cells in
    (* Prevent too many calls of [mmap] which quickly waste
       the virtual memory space *)
    if mapped_length <> t.mapped_length then begin
      t.array <- make_array
          ~bytes_per_cell:t.bytes_per_cell
          (Lwt_unix.unix_file_descr t.fd)
          ~pos:t.pos (Index.to_int mapped_length);
      t.mapped_length <- mapped_length;
    end;
    t.current_length   <- h.Header.last_next_index;
    t.last_root_index  <- h.Header.last_root_index;
    Lwt.return_unit

type Error.t += Invalid_writer_key

let () = Error.register_printer (function
    | Invalid_writer_key -> Some "invalid writer key"
    | _ -> None)

type Error.t += Not_reader

let () = Error.register_printer (function
    | Not_reader -> Some "not reader"
    | _ -> None)

let enable_process_sync t writer_key =
  match t.mode with
  | Reader ->
      t.writer_key <- Some writer_key;
      let+ res = Header.read_process_sync t in
      (match res with
       | Some (`GoodKey, _) -> Ok ()
       | _ ->
           t.writer_key <- None;
           Error Invalid_writer_key)
  | Writer -> Lwt.return_error Not_reader

let update_reader t =
  match t.mode with
  | Writer -> Lwt.return_unit
  | Reader ->
      match t.writer_key with
      | None ->
          let* res = Header.read_disk_sync t in
          begin match res with
          | None -> failwithf "Failed to load header"
          | Some h -> reopen_with t h
          end
      | Some _ ->
          let* res = Header.read_process_sync t in
          match res with
          | None -> failwithf "Failed to load header"
          | Some (_, h) ->
              (* `BadKey is ok here. It means another writer has started *)
              reopen_with t h

module Chunk = struct

  (* Store data bigger than [bytes_per_cell] bytes *)

  (* size information is uint16 *)
  let max_bytes_per_chunk = 65535

  (* How many cells are required to save the given size of bytes
     in one chunk *)
  let ncells bytes_per_cell size = (size + 6 + bytes_per_cell - 1) / bytes_per_cell

  let get_footer_fields storage last_index =
    let buf =
      Mmap.get_buffer
        ~off:((last_index + 1) * storage.bytes_per_cell - 6)
        ~len:6 storage.array
    in
    let cdr = B.get_index buf 2 in
    let size = B.get_uint16 buf 0 in
    (cdr, size)

  let get_chunk storage last_index =
    let cdr, size = get_footer_fields storage (Index.to_int last_index) in
    let ncells = ncells storage.bytes_per_cell size in
    let first_index = Index.(last_index - Unsafe.of_int ncells + one) in
    (get_bytes storage first_index size, size, cdr)

  let get_chunks storage last_index =
    let rec aux (bufs, size) last_index =
      let buf, bytes, cdr = get_chunk storage last_index in
      let bufs = buf :: bufs in
      let size = size + bytes in (* overflow in 32bit? *)
      if cdr = Index.zero then (bufs, size)
      else aux (bufs, size) cdr
    in
    aux ([], 0) last_index

  let string_of_cstructs bufs =
    String.concat "" @@ List.map B.to_string bufs

  let read t i = string_of_cstructs @@ fst @@ get_chunks t i

  let write_to_chunk storage cdr s off len =
    if mode storage = Reader then invalid_arg "Reader cannot write";
    assert (String.length s >= off + len);
    assert (len <= max_bytes_per_chunk);
    let bytes_per_cell = storage.bytes_per_cell in
    let ncells = ncells bytes_per_cell len in
    let cdr_pos = ncells * bytes_per_cell - 4 in
    let size_pos = cdr_pos - 2 in

    let+? i = new_indices storage ncells in

    let last_index = Index.(i + Unsafe.of_int ncells - one) in
    let chunk = get_bytes storage i (bytes_per_cell * ncells) in
    B.blit_from_string s off chunk 0 len;
    B.set_uint16 chunk size_pos len;
    B.set_index chunk cdr_pos cdr;
    last_index

  let write storage s =
    if mode storage = Reader then invalid_arg "Reader cannot write";
    let rec f off remain cdr  =
      let len = if remain > max_bytes_per_chunk then max_bytes_per_chunk else remain in
      let*? cdr' = write_to_chunk storage cdr s off len in
      let off' = off + len in
      let remain' = remain - len in
      if remain' > 0 then f off' remain' cdr'
      else Ok cdr'
    in
    f 0 (String.length s) Index.zero

  let test_write_read st storage =
    let size = Random.State.int st (65536 * 5) in (* 5 chunks at maximum *)
    let s = String.init size @@ fun i -> Char.chr (Char.code 'A' + i mod 20) in
    let i = from_Ok @@ write storage s in
    let s' = string_of_cstructs @@ fst @@ get_chunks storage i in
    if (s <> s') then begin prerr_endline s; prerr_endline s' end;
    assert (s = s')
end

module Internal = struct
  let msync = msync
  let set_current_length = set_current_length
  module Header = Header
end