Naїve Command Write-Ahead Log implementation

wal.ml

open Unix
open Printf

exception EOF
exception WriteError

module BE = EndianString.BigEndian
module M = Map.Make(struct type t = int let compare = compare end)

type entry = Todo of (int * string) | Done of int

type t = {
  fn : string; (* file name *)
  fd : file_descr option;
  id : int; (* last stored string id *)
  em : string M.t;
}

let get_fd fn = openfile fn [O_CREAT; O_APPEND; O_RDWR; O_SYNC; O_DSYNC] 0o600

let rec read_func fd buf off len =
  let res = Unix.read fd buf off len in
  match res with
  | 0 -> raise EOF
  | -1 -> exit 1
  | i when i = len -> ()
  | i when i < len -> read_func fd buf (off + i) (len - i)
  | i when i < -1 ->
    let s =
      sprintf
        "read_func readed negative number of bytes: %d\n%!" i in
    raise (Failure s)
  | _ ->
    let s =
      sprintf
        "read_func readed more, than requested: res = %d, off = %d, len = %d\n%!"
        res off len
    in raise (Failure s)

let read_entry fd buf =
  read_func fd buf 0 5;
  let id = Int32.to_int @@ BE.get_int32 buf 0
  and todo = 0 = BE.get_int8 buf 4 in
  if todo then begin
    read_func fd buf 0 4;
    let len = Int32.to_int @@ BE.get_int32 buf 0 in
    let str = Bytes.create len in
    read_func fd str 0 len;
    Todo (id,str)
  end else Done id

let rec read_file ?(buf = Bytes.create 10) ?(res = []) fd =
  try
    let e = read_entry fd buf in
    read_file ~buf ~res:(e :: res) fd
  with EOF -> res

let write_entry fd = function
  | Todo (id, str) ->
    let len = Bytes.length str in
    let buf = Bytes.create 9 in
    BE.set_int32 buf 0 @@ Int32.of_int id;
    BE.set_int8  buf 4 0; (* XXX TODO value = 0 *)
    BE.set_int32 buf 5 @@ Int32.of_int len;
    let cnt = write fd buf 0 9 in
    if (len + 9) <> (cnt + write fd str 0 len)
    then raise WriteError
    else ()
  | Done id ->
    let buf = Bytes.create 5 in
    BE.set_int32 buf 0 @@ Int32.of_int id;
    BE.set_int8  buf 4 1; (* XXX: DONE value = 1 *)
    if 5 <> write fd buf 0 5
    then raise WriteError
    else ()

let init fn =
  let fd = get_fd fn in
  let em = List.fold_left
      (fun m e ->
         match e with
         | Todo (i,s) -> M.add i s m
         | Done i -> try M.remove i m with _ -> m
      ) M.empty @@ read_file fd in
  let fd,id =
    if M.is_empty em then begin
      close fd;
      unlink fn;
      None,0
    end else begin
      let id,_ = M.max_binding em in
      Some fd,(id+1)
    end
  in
  { fn; fd; id; em }

let todo t s =
  let fd = match t.fd with
    | Some fd -> fd
    | None -> get_fd t.fn
  in
  let t = { t with fd = Some fd } in
  let id = t.id in
  let e = Todo (id,s) in
  let _ = write_entry fd e
  and em = M.add id s t.em in
  { t with em; id = t.id + 1 }, id

let finished t id =
  let e = Done id in
  let fd =
    match t.fd with
    | Some fd -> fd
    | None -> get_fd t.fn
  in
  write_entry fd e;
  let em = M.remove id t.em in
  let fd,id =
    if M.is_empty em then begin
      close fd;
      unlink t.fn;
      None,0
    end else begin
      let id,_ = M.max_binding em in
      Some fd,(id+1)
    end
  in
  { t with id; fd; em }