269 lines
8.3 KiB
OCaml
269 lines
8.3 KiB
OCaml
(**************************************************************************)
|
|
(* *)
|
|
(* Copyright (c) 2014 - 2018. *)
|
|
(* Dynamic Ledger Solutions, Inc. <contact@tezos.com> *)
|
|
(* *)
|
|
(* All rights reserved. No warranty, explicit or implicit, provided. *)
|
|
(* *)
|
|
(**************************************************************************)
|
|
|
|
open Rresult
|
|
|
|
type t = {
|
|
dir : Lmdb.t ;
|
|
parent : (Lmdb.rw Lmdb.txn * Lmdb.db * Lmdb.rw Lmdb.cursor) Lwt.key ;
|
|
}
|
|
|
|
type key = string list
|
|
type value = MBytes.t
|
|
|
|
type error += Unknown of string list
|
|
|
|
let () =
|
|
Error_monad.register_error_kind
|
|
`Permanent
|
|
~id:"raw_store.unknown"
|
|
~title:"Missing key in store"
|
|
~description:"Missing key in store"
|
|
~pp:(fun ppf keys ->
|
|
Format.fprintf ppf
|
|
"Missing key in store: %s"
|
|
(String.concat "/" keys))
|
|
Data_encoding.(obj1 (req "key" (list string)))
|
|
(function Unknown keys -> Some keys | _ -> None)
|
|
(fun keys -> Unknown keys)
|
|
|
|
let concat = String.concat "/"
|
|
let split = String.split_on_char '/'
|
|
|
|
let lwt_fail_error err =
|
|
Lwt.fail_with (Lmdb.string_of_error err)
|
|
|
|
let of_result = function
|
|
| Ok res -> Lwt.return res
|
|
| Error err -> lwt_fail_error err
|
|
|
|
let (>>=?) v f =
|
|
match v with
|
|
| Error err -> lwt_fail_error err
|
|
| Ok v -> f v
|
|
|
|
let init ?mapsize path =
|
|
if not (Sys.file_exists path) then Unix.mkdir path 0o755 ;
|
|
match Lmdb.opendir ?mapsize ~flags:[NoTLS] path 0o644 with
|
|
| Ok dir -> return { dir ; parent = Lwt.new_key () }
|
|
| Error err -> failwith "%a" Lmdb.pp_error err
|
|
|
|
let close { dir } = Lmdb.closedir dir
|
|
|
|
let known { dir ; parent } key =
|
|
begin match Lwt.get parent with
|
|
| Some (txn, db, _cursor) -> Lmdb.mem txn db (concat key)
|
|
| None ->
|
|
Lmdb.with_ro_db dir ~f:begin fun txn db ->
|
|
Lmdb.mem txn db (concat key)
|
|
end
|
|
end |> of_result
|
|
|
|
let read_opt { dir ; parent } key =
|
|
begin match Lwt.get parent with
|
|
| Some (txn, db, _cursor) -> Lmdb.get txn db (concat key) >>| MBytes.copy
|
|
| None ->
|
|
Lmdb.with_ro_db dir ~f:begin fun txn db ->
|
|
Lmdb.get txn db (concat key) >>| MBytes.copy
|
|
end
|
|
end |> function
|
|
| Ok v -> Lwt.return_some v
|
|
| Error KeyNotFound -> Lwt.return_none
|
|
| Error err -> lwt_fail_error err
|
|
|
|
let read { dir ; parent } key =
|
|
begin match Lwt.get parent with
|
|
| Some (txn, db, _cursor) -> Lmdb.get txn db (concat key) >>| MBytes.copy
|
|
| None ->
|
|
Lmdb.with_ro_db dir ~f:begin fun txn db ->
|
|
Lmdb.get txn db (concat key) >>| MBytes.copy
|
|
end
|
|
end |> function
|
|
| Ok v -> return v
|
|
| Error _err -> fail (Unknown key)
|
|
|
|
let read_exn { dir ; parent } key =
|
|
begin match Lwt.get parent with
|
|
| Some (txn, db, _cursor) -> Lmdb.get txn db (concat key) >>| MBytes.copy
|
|
| None ->
|
|
Lmdb.with_ro_db dir ~f:begin fun txn db ->
|
|
Lmdb.get txn db (concat key) >>| MBytes.copy
|
|
end
|
|
end |> of_result
|
|
|
|
let store { dir ; parent } k v =
|
|
begin match Lwt.get parent with
|
|
| Some (txn, db, _cursor) -> Lmdb.put txn db (concat k) v
|
|
| None ->
|
|
Lmdb.with_rw_db dir ~f:begin fun txn db ->
|
|
Lmdb.put txn db (concat k) v
|
|
end
|
|
end |> of_result
|
|
|
|
let remove { dir ; parent } k =
|
|
let remove txn db =
|
|
match Lmdb.del txn db (concat k) with
|
|
| Ok () -> Ok ()
|
|
| Error KeyNotFound -> Ok ()
|
|
| Error err -> Error err in
|
|
begin match Lwt.get parent with
|
|
| Some (txn, db, _cursor) -> remove txn db
|
|
| None -> Lmdb.with_rw_db dir ~f:remove
|
|
end |> of_result
|
|
|
|
let is_prefix s s' =
|
|
String.(length s <= length s' && compare s (sub s' 0 (length s)) = 0)
|
|
|
|
let known_dir { dir ; parent } k =
|
|
let k = concat k in
|
|
let cursor_fun cursor =
|
|
Lmdb.cursor_at cursor k >>= fun () ->
|
|
Lmdb.cursor_get cursor >>| fun (first_k, _v) ->
|
|
(is_prefix k (MBytes.to_string first_k))
|
|
in
|
|
begin match Lwt.get parent with
|
|
| Some (txn, db, _cursor) ->
|
|
Lmdb.with_cursor txn db ~f:cursor_fun
|
|
| None ->
|
|
Lmdb.with_ro_db dir ~f:begin fun txn db ->
|
|
Lmdb.with_cursor txn db ~f:cursor_fun
|
|
end
|
|
end |> of_result
|
|
|
|
let remove_dir { dir ; parent } k =
|
|
let k = concat k in
|
|
let cursor_fun cursor =
|
|
Lmdb.cursor_at cursor k >>= fun () ->
|
|
Lmdb.cursor_iter cursor ~f:begin fun (kk, _v) ->
|
|
let kk_string = MBytes.to_string kk in
|
|
if is_prefix k kk_string then begin
|
|
Lmdb.cursor_del cursor
|
|
end
|
|
else Error KeyNotFound
|
|
end in
|
|
begin match Lwt.get parent with
|
|
| Some (txn, db, _cursor) ->
|
|
Lmdb.with_cursor txn db ~f:cursor_fun
|
|
| None ->
|
|
Lmdb.with_rw_db dir ~f:begin fun txn db ->
|
|
Lmdb.with_cursor txn db ~f:cursor_fun
|
|
end
|
|
end |> function
|
|
| Error KeyNotFound
|
|
| Ok () -> Lwt.return_unit
|
|
| Error err -> lwt_fail_error err
|
|
|
|
let list_equal l1 l2 len =
|
|
if len < 0 || len > List.length l1 || len > List.length l2
|
|
then invalid_arg "list_compare: invalid len" ;
|
|
let rec inner l1 l2 len =
|
|
match len, l1, l2 with
|
|
| 0, _, _ -> true
|
|
| _, [], _
|
|
| _, _, [] -> false
|
|
| _, h1 :: t1, h2 :: t2 ->
|
|
if h1 <> h2 then false
|
|
else inner t1 t2 (pred len)
|
|
in
|
|
inner l1 l2 len
|
|
|
|
let is_child ~parent ~child =
|
|
let plen = List.length parent in
|
|
let clen = List.length child in
|
|
clen > plen && list_equal parent child plen
|
|
|
|
let list_sub l pos len =
|
|
if len < 0 || pos < 0 || pos + len > List.length l then
|
|
invalid_arg "list_sub" ;
|
|
let rec inner (acc, n) = function
|
|
| [] -> List.rev acc
|
|
| h :: t ->
|
|
if n = 0 then List.rev acc
|
|
else inner (h :: acc, pred n) t in
|
|
inner ([], len) l
|
|
|
|
let with_rw_cursor_lwt ?sync ?metasync ?flags ?name { dir ; parent } ~f =
|
|
let local_parent =
|
|
match Lwt.get parent with
|
|
| None -> None
|
|
| Some (txn, _db, _cursor) -> Some txn in
|
|
Lmdb.create_rw_txn
|
|
?sync ?metasync ?parent:local_parent dir >>=? fun txn ->
|
|
Lmdb.opendb ?flags ?name txn >>=? fun db ->
|
|
Lmdb.opencursor txn db >>=? fun cursor ->
|
|
Lwt.with_value parent (Some (txn, db, cursor)) begin fun () ->
|
|
Lwt.try_bind (fun () -> f cursor)
|
|
begin fun res ->
|
|
Lmdb.cursor_close cursor ;
|
|
Lmdb.commit_txn txn >>=? fun () ->
|
|
Lwt.return res
|
|
end
|
|
begin fun exn ->
|
|
Lmdb.cursor_close cursor ;
|
|
Lmdb.abort_txn txn ;
|
|
Lwt.fail exn
|
|
end
|
|
end
|
|
|
|
let cursor_next_lwt cursor acc f =
|
|
match Lmdb.cursor_next cursor with
|
|
| Error KeyNotFound -> acc
|
|
| Error err -> lwt_fail_error err
|
|
| Ok () -> Lwt.bind acc f
|
|
|
|
let fold t k ~init ~f =
|
|
let base_len = List.length k in
|
|
let rec inner ht cursor acc =
|
|
Lmdb.cursor_get cursor >>=? fun (kk, _v) ->
|
|
let kk = MBytes.to_string kk in
|
|
let kk_split = split kk in
|
|
match is_child ~child:kk_split ~parent:k with
|
|
| false -> Lwt.return acc
|
|
| true ->
|
|
let cur_len = List.length kk_split in
|
|
if cur_len = succ base_len then begin
|
|
cursor_next_lwt cursor (f (`Key kk_split) acc) (inner ht cursor)
|
|
end
|
|
else begin
|
|
let dir = list_sub kk_split 0 (succ base_len) in
|
|
if Hashtbl.mem ht dir then
|
|
cursor_next_lwt cursor (Lwt.return acc) (inner ht cursor)
|
|
else begin
|
|
Hashtbl.add ht dir () ;
|
|
cursor_next_lwt cursor (f (`Dir dir) acc) (inner ht cursor)
|
|
end
|
|
end in
|
|
with_rw_cursor_lwt t ~f:begin fun cursor ->
|
|
match Lmdb.cursor_at cursor (concat k) with
|
|
| Error KeyNotFound -> Lwt.return init
|
|
| Error err -> lwt_fail_error err
|
|
| Ok () ->
|
|
let ht = Hashtbl.create 31 in
|
|
inner ht cursor init
|
|
end
|
|
|
|
let fold_keys t k ~init ~f =
|
|
with_rw_cursor_lwt t ~f:begin fun cursor ->
|
|
match Lmdb.cursor_at cursor (concat k) with
|
|
| Error KeyNotFound -> Lwt.return init
|
|
| Error err -> lwt_fail_error err
|
|
| Ok () ->
|
|
let rec inner acc =
|
|
Lmdb.cursor_get cursor >>=? fun (kk, _v) ->
|
|
let kk = MBytes.to_string kk in
|
|
let kk_split = split kk in
|
|
match is_child ~child:kk_split ~parent:k with
|
|
| false -> Lwt.return acc
|
|
| true -> cursor_next_lwt cursor (f kk_split acc) inner
|
|
in inner init
|
|
end
|
|
|
|
let keys t =
|
|
fold_keys t ~init:[] ~f:(fun k acc -> Lwt.return (k :: acc))
|