Cache MST blocks within applyWrites

mist/lib/mst.ml

···
 
   val create_empty : Store.t -> (t, exn) Lwt_result.t
 
-  val get_cid : t -> string -> Cid.t option Lwt.t
-
   val of_assoc : Store.t -> (string * Cid.t) list -> t Lwt.t
 
   val add : t -> string -> Cid.t -> t Lwt.t
···
   val leaves_of_node : node -> (string * Cid.t) list Lwt.t
 
   val leaves_of_root : t -> (string * Cid.t) list Lwt.t
-
-  val null_diff : t -> data_diff Lwt.t
 
   val equal : t -> t -> bool Lwt.t
 end
···
     Lwt_result.bind (Store.put_block blockstore cid encoded) (fun _ ->
         Lwt.return_ok {blockstore; root= cid} )
 
-  (* returns the cid for a given key, if it exists *)
-  let get_cid t key : Cid.t option Lwt.t =
-    let rec get_in_node (n : node) : Cid.t option Lwt.t =
-      let sorted_entries =
-        List.sort
-          (fun (a : entry) (b : entry) -> String.compare a.key b.key)
-          n.entries
-      in
-      let rec scan (prev : entry option) (entries : entry list) :
-          Cid.t option Lwt.t =
-        match entries with
-        | [] -> (
-          match prev with
-          | Some p -> (
-              p.right
-              >>? function Some r -> get_in_node r | None -> Lwt.return_none )
-          | None -> (
-              n.left
-              >>? function Some l -> get_in_node l | None -> Lwt.return_none ) )
-        | e :: rest ->
-            if key = e.key then Lwt.return_some e.value
-            else if key < e.key then
-              match prev with
-              | Some p -> (
-                  p.right
-                  >>? function
-                  | Some r ->
-                      get_in_node r
-                  | None ->
-                      Lwt.return_none )
-              | None -> (
-                  n.left
-                  >>? function
-                  | Some l ->
-                      get_in_node l
-                  | None ->
-                      Lwt.return_none )
-            else scan (Some e) rest
-      in
-      scan None sorted_entries
-    in
-    match%lwt retrieve_node t t.root with
-    | None ->
-        Lwt.fail (Invalid_argument "root cid not found in repo store")
-    | Some root ->
-        get_in_node root
-
   (* builds and persists a canonical mst from sorted leaves *)
   let of_assoc blockstore assoc : t Lwt.t =
     let open Lwt.Infix in
···
     let%lwt _ = Store.put_block blockstore cid encoded in
     Lwt.return cid
 
-  (* returns the layer a raw node belongs to *)
-  let rec get_layer_raw (t : t) (raw : node_raw) : int Lwt.t =
-    match (raw.l, raw.e) with
-    | None, [] ->
-        Lwt.return 0
-    | Some left_cid, [] -> (
-      match%lwt retrieve_node_raw t left_cid with
-      | Some left_raw ->
-          let%lwt left_layer = get_layer_raw t left_raw in
-          Lwt.return (left_layer + 1)
-      | None ->
-          failwith ("couldn't find node " ^ Cid.to_string left_cid) )
-    | _, e :: _ -> (
-      match e.p with
-      | 0 ->
-          Lwt.return (Util.leading_zeros_on_hash (Bytes.to_string e.k))
-      | _ ->
-          failwith "first node entry has nonzero p value" )
-
   (* decompress entry keys from a raw node *)
   let decompress_keys (raw : node_raw) : string list =
     let last_key = ref "" in
···
         let t' = {blockstore; root= result.root} in
         match%lwt retrieve_node_raw t' result.root with
         | Some raw ->
-            let%lwt layer = get_layer_raw t' raw in
+            let%lwt layer = get_node_height t' raw in
             Lwt.return (Some result.root, layer)
         | None ->
             Lwt.return (Some result.root, 0) )
···
     | None ->
         failwith ("couldn't find node " ^ Cid.to_string root_cid)
     | Some raw ->
-        let%lwt root_layer = get_layer_raw t raw in
+        let%lwt root_layer = get_node_height t raw in
         if key_layer > root_layer then
           add_above_root t root_cid root_layer key value key_layer
         else if key_layer = root_layer then
···
           let%lwt new_left_layer =
             match new_left_raw_opt with
             | Some r ->
-                get_layer_raw t r
+                get_node_height t r
             | None ->
                 Lwt.return 0
           in
···
             let%lwt new_right_layer =
               match new_right_raw_opt with
               | Some r ->
-                  get_layer_raw t r
+                  get_node_height t r
               | None ->
                   Lwt.return 0
             in
···
     | None ->
         Lwt.return_none
     | Some raw ->
-        let%lwt root_layer = get_layer_raw t raw in
+        let%lwt root_layer = get_node_height t raw in
         if key_layer > root_layer then
           (* key can't exist above root *)
           Lwt.return_some (root_cid, root_layer)
···
               let%lwt result_layer =
                 match%lwt retrieve_node_raw t result.root with
                 | Some r ->
-                    get_layer_raw t r
+                    get_node_height t r
                 | None ->
                     Lwt.return 0
               in
···
     | Some (new_root, _layer) ->
         Lwt.return {t with root= new_root}
 
-  (* produces a diff from an empty mst to the current one *)
-  let null_diff curr : data_diff Lwt.t =
-    let%lwt curr_nodes, _, curr_leaf_set = collect_nodes_and_leaves curr in
-    let%lwt curr_leaves = leaves_of_root curr in
-    let adds = List.map (fun (key, cid) : diff_add -> {key; cid}) curr_leaves in
-    Lwt.return
-      { adds
-      ; updates= []
-      ; deletes= []
-      ; new_mst_blocks= curr_nodes
-      ; new_leaf_cids= curr_leaf_set
-      ; removed_cids= Cid.Set.empty }
-
   (* checks that two msts are identical by recursively comparing their entries *)
   let equal (t1 : t) (t2 : t) : bool Lwt.t =
     let rec nodes_equal (n1 : node) (n2 : node) : bool Lwt.t =
···
         Lwt.return false
 end
 
-module Differ (Prev : Intf) (Curr : Intf) = struct
-  let diff ~(t_curr : Curr.t) ~(t_prev : Prev.t) : data_diff Lwt.t =
-    let%lwt curr_nodes, curr_node_set, curr_leaf_set =
-      Curr.collect_nodes_and_leaves t_curr
-    in
-    let%lwt _, prev_node_set, prev_leaf_set =
-      Prev.collect_nodes_and_leaves t_prev
-    in
-    (* just convenient to have these functions *)
-    let in_prev_nodes cid = Cid.Set.mem cid prev_node_set in
-    let in_curr_nodes cid = Cid.Set.mem cid curr_node_set in
-    let in_prev_leaves cid = Cid.Set.mem cid prev_leaf_set in
-    let in_curr_leaves cid = Cid.Set.mem cid curr_leaf_set in
-    (* new mst blocks are curr nodes that are not in prev *)
-    let new_mst_blocks =
-      List.filter (fun (cid, _) -> not (in_prev_nodes cid)) curr_nodes
-    in
-    (* removed cids are prev nodes not in curr plus prev leaves not in curr *)
-    let removed_node_cids =
-      Cid.Set.fold
-        (fun cid acc ->
-          if not (in_curr_nodes cid) then Cid.Set.add cid acc else acc )
-        prev_node_set Cid.Set.empty
-    in
-    let removed_leaf_cids =
-      Cid.Set.fold
-        (fun cid acc ->
-          if not (in_curr_leaves cid) then Cid.Set.add cid acc else acc )
-        prev_leaf_set Cid.Set.empty
-    in
-    let removed_cids = Cid.Set.union removed_node_cids removed_leaf_cids in
-    (* new leaf cids are curr leaves not in prev *)
-    let new_leaf_cids =
-      Cid.Set.fold
-        (fun cid acc ->
-          if not (in_prev_leaves cid) then Cid.Set.add cid acc else acc )
-        curr_leaf_set Cid.Set.empty
-    in
-    (* compute adds/updates/deletes by merging sorted leaves *)
-    let%lwt curr_leaves = Curr.leaves_of_root t_curr in
-    let%lwt prev_leaves = Prev.leaves_of_root t_prev in
-    let rec merge (pl : (string * Cid.t) list) (cl : (string * Cid.t) list)
-        (adds : diff_add list) (updates : diff_update list)
-        (deletes : diff_delete list) =
-      match (pl, cl) with
-      | [], [] ->
-          (* we prepend for speed, then reverse at the end *)
-          (List.rev adds, List.rev updates, List.rev deletes)
-      | [], (k, c) :: cr ->
-          (* more curr than prev, goes in adds *)
-          merge [] cr ({key= k; cid= c} :: adds) updates deletes
-      | (k, c) :: pr, [] ->
-          (* more prev than curr, goes in deletes *)
-          merge pr [] adds updates ({key= k; cid= c} :: deletes)
-      | (k1, c1) :: pr, (k2, c2) :: cr ->
-          if k1 = k2 then (* if key & value are the same, keep going *)
-            if Cid.equal c1 c2 then merge pr cr adds updates deletes
-            else (* same key, different value; update *)
-              merge pr cr adds ({key= k1; prev= c1; cid= c2} :: updates) deletes
-          else if k1 < k2 then
-            merge pr ((k2, c2) :: cr) adds updates
-              ({key= k1; cid= c1} :: deletes)
-          else
-            merge ((k1, c1) :: pr) cr
-              ({key= k2; cid= c2} :: adds)
-              updates deletes
-    in
-    let adds, updates, deletes = merge prev_leaves curr_leaves [] [] [] in
-    Lwt.return
-      {adds; updates; deletes; new_mst_blocks; new_leaf_cids; removed_cids}
-end
-
 module Inductive (M : Intf) = struct
   module Cache_bs = Cache_blockstore (Memory_blockstore)
   module Mem_mst = Make (Cache_bs)
···
           (String_map.bindings map)
       in
       (* save this now so we can read blocks from it later *)
-      let block_map = mem_mst.blockstore.bs.blocks in
+      let blockstore = mem_mst.blockstore in
       (* apply inverse of operations in reverse order,
          check that mst root matches prev_root *)
       let%lwt inverted_mst, added_cids =
···
              (Cid.to_string prev_root)
              (Cid.to_string inverted_mst.root) ) ;
       let proof_cids =
-        Cid.Set.union added_cids mem_mst.blockstore.reads
+        Cid.Set.union added_cids (Cache_bs.get_reads blockstore)
         |> Cid.Set.remove prev_root |> Cid.Set.add new_root
       in
       let {blocks= proof_bm; _} : Block_map.with_missing =
-        Block_map.get_many (Cid.Set.elements proof_cids) block_map
+        Block_map.get_many
+          (Cid.Set.elements proof_cids)
+          (Cache_bs.get_cache blockstore)
       in
       Lwt.return_ok proof_bm
     with e -> Lwt.return_error e

mist/lib/storage/block_map.ml

···
 
 let get_many cids m =
   let blocks, missing =
-    List.fold_left
-      (fun (b, mis) cid ->
-        match get cid m with
-        | Some bytes ->
-            (Cid_map.add cid bytes b, mis)
-        | None ->
-            (b, mis @ [cid]) )
-      (Cid_map.empty, []) cids
+    List.partition_map
+      (fun cid ->
+        match get cid m with Some data -> Left (cid, data) | None -> Right cid )
+      cids
   in
-  {blocks; missing= List.rev missing}
+  {blocks= Cid_map.of_list blocks; missing}
 
 let has = Cid_map.mem
 

mist/lib/storage/cache_blockstore.ml

···
-type 'bs data = {mutable reads: Cid.Set.t; bs: 'bs}
+type 'bs data = {mutable reads: Cid.Set.t; mutable cache: Block_map.t; bs: 'bs}
 
 module Make
     (Bs : Blockstore.Writable) : sig
   include Blockstore.Writable
 
   val create : Bs.t -> t
+
+  val get_reads : t -> Cid.Set.t
+
+  val get_cache : t -> Block_map.t
 end
 with type t = Bs.t data = struct
   type t = Bs.t data
 
-  let create bs = {reads= Cid.Set.empty; bs}
+  let create bs = {reads= Cid.Set.empty; cache= Block_map.empty; bs}
+
+  let get_reads t = t.reads
+
+  let get_cache t = t.cache
 
   let get_bytes t cid =
-    match%lwt Bs.get_bytes t.bs cid with
-    | Some _ as res ->
+    match Block_map.get cid t.cache with
+    | Some _ as cached ->
         t.reads <- Cid.Set.add cid t.reads ;
-        Lwt.return res
-    | None ->
-        Lwt.return_none
+        Lwt.return cached
+    | None -> (
+      match%lwt Bs.get_bytes t.bs cid with
+      | Some data as res ->
+          t.cache <- Block_map.set cid data t.cache ;
+          t.reads <- Cid.Set.add cid t.reads ;
+          Lwt.return res
+      | None ->
+          Lwt.return_none )
 
-  let has t cid = Bs.has t.bs cid
+  let has t cid =
+    if Block_map.has cid t.cache then Lwt.return_true else Bs.has t.bs cid
 
   let get_blocks t cids =
-    let%lwt bm = Bs.get_blocks t.bs cids in
-    t.reads <-
-      Cid.Set.union t.reads (Cid.Set.of_list (Block_map.keys bm.blocks)) ;
-    Lwt.return bm
+    let {Block_map.blocks= cached; missing} = Block_map.get_many cids t.cache in
+    (* mark cached as read *)
+    Block_map.iter (fun cid _ -> t.reads <- Cid.Set.add cid t.reads) cached ;
+    (* fetch missing from underlying store *)
+    let%lwt fetched = Bs.get_blocks t.bs missing in
+    (* cache and mark as read *)
+    Block_map.iter
+      (fun cid data ->
+        t.cache <- Block_map.set cid data t.cache ;
+        t.reads <- Cid.Set.add cid t.reads )
+      fetched.blocks ;
+    (* combine results *)
+    let blocks =
+      List.fold_left
+        (fun acc (cid, data) -> Block_map.set cid data acc)
+        fetched.blocks (Block_map.entries cached)
+    in
+    Lwt.return {Block_map.blocks; missing= fetched.missing}
 
-  let put_block t cid bytes = Bs.put_block t.bs cid bytes
+  let put_block t cid bytes =
+    t.cache <- Block_map.set cid bytes t.cache ;
+    Bs.put_block t.bs cid bytes
 
-  let put_many t blocks = Bs.put_many t.bs blocks
+  let put_many t blocks =
+    Block_map.iter
+      (fun cid data -> t.cache <- Block_map.set cid data t.cache)
+      blocks ;
+    Bs.put_many t.bs blocks
 
-  let delete_block t cid = Bs.delete_block t.bs cid
+  let delete_block t cid =
+    t.cache <- Block_map.remove cid t.cache ;
+    Bs.delete_block t.bs cid
 
-  let delete_many t cids = Bs.delete_many t.bs cids
+  let delete_many t cids =
+    List.iter (fun cid -> t.cache <- Block_map.remove cid t.cache) cids ;
+    Bs.delete_many t.bs cids
 end

mist/test/test_mst.ml

···
 open Lwt.Infix
 open Lwt_result.Syntax
 module Mem_mst = Mst.Make (Storage.Memory_blockstore)
-module Mem_diff = Mst.Differ (Mem_mst) (Mem_mst)
 module String_map = Dag_cbor.String_map
+
+module Differ (Prev : Mst.Intf) (Curr : Mst.Intf) = struct
+  let diff ~(t_curr : Curr.t) ~(t_prev : Prev.t) : Mst.data_diff Lwt.t =
+    let%lwt curr_nodes, curr_node_set, curr_leaf_set =
+      Curr.collect_nodes_and_leaves t_curr
+    in
+    let%lwt _, prev_node_set, prev_leaf_set =
+      Prev.collect_nodes_and_leaves t_prev
+    in
+    let in_prev_nodes cid = Cid.Set.mem cid prev_node_set in
+    let in_curr_nodes cid = Cid.Set.mem cid curr_node_set in
+    let in_prev_leaves cid = Cid.Set.mem cid prev_leaf_set in
+    let in_curr_leaves cid = Cid.Set.mem cid curr_leaf_set in
+    let new_mst_blocks =
+      List.filter (fun (cid, _) -> not (in_prev_nodes cid)) curr_nodes
+    in
+    let removed_node_cids =
+      Cid.Set.fold
+        (fun cid acc ->
+          if not (in_curr_nodes cid) then Cid.Set.add cid acc else acc )
+        prev_node_set Cid.Set.empty
+    in
+    let removed_leaf_cids =
+      Cid.Set.fold
+        (fun cid acc ->
+          if not (in_curr_leaves cid) then Cid.Set.add cid acc else acc )
+        prev_leaf_set Cid.Set.empty
+    in
+    let removed_cids = Cid.Set.union removed_node_cids removed_leaf_cids in
+    let new_leaf_cids =
+      Cid.Set.fold
+        (fun cid acc ->
+          if not (in_prev_leaves cid) then Cid.Set.add cid acc else acc )
+        curr_leaf_set Cid.Set.empty
+    in
+    let%lwt curr_leaves = Curr.leaves_of_root t_curr in
+    let%lwt prev_leaves = Prev.leaves_of_root t_prev in
+    let rec merge (pl : (string * Cid.t) list) (cl : (string * Cid.t) list)
+        (adds : Mst.diff_add list) (updates : Mst.diff_update list)
+        (deletes : Mst.diff_delete list) =
+      match (pl, cl) with
+      | [], [] ->
+          (List.rev adds, List.rev updates, List.rev deletes)
+      | [], (k, c) :: cr ->
+          merge [] cr ({key= k; cid= c} :: adds) updates deletes
+      | (k, c) :: pr, [] ->
+          merge pr [] adds updates ({key= k; cid= c} :: deletes)
+      | (k1, c1) :: pr, (k2, c2) :: cr ->
+          if k1 = k2 then
+            if Cid.equal c1 c2 then merge pr cr adds updates deletes
+            else
+              merge pr cr adds ({key= k1; prev= c1; cid= c2} :: updates) deletes
+          else if k1 < k2 then
+            merge pr ((k2, c2) :: cr) adds updates
+              ({key= k1; cid= c1} :: deletes)
+          else
+            merge ((k1, c1) :: pr) cr
+              ({key= k2; cid= c2} :: adds)
+              updates deletes
+    in
+    let adds, updates, deletes = merge prev_leaves curr_leaves [] [] [] in
+    Lwt.return
+      {Mst.adds; updates; deletes; new_mst_blocks; new_leaf_cids; removed_cids}
+end
+
+module Mem_diff = Differ (Mem_mst) (Mem_mst)
 
 let cid_of_string_exn s =
   match Cid.of_string s with Ok c -> c | Error msg -> failwith msg
···
   let%lwt mst' =
     Lwt_list.fold_left_s (fun t (k, v) -> Mem_mst.add t k v) mst shuffled
   in
-  let%lwt () =
-    Lwt_list.iter_s
-      (fun (k, v) ->
-        let%lwt got = Mem_mst.get_cid mst' k in
-        Alcotest.(check bool)
-          "added records retrievable" true
-          (Option.value
-             (Option.map (fun x -> Cid.equal v x) got)
-             ~default:false )
-        |> Lwt.return )
-      shuffled
-  in
+  let%lwt result_map = Mem_mst.build_map mst' in
+  List.iter
+    (fun (k, v) ->
+      let got = String_map.find_opt k result_map in
+      Alcotest.(check bool)
+        "added records retrievable" true
+        (Option.value (Option.map (fun x -> Cid.equal v x) got) ~default:false) )
+    shuffled ;
   let%lwt total = Mem_mst.leaf_count mst' in
   Alcotest.(check int) "leaf count after adds" 1000 total ;
   Lwt.return_ok ()
···
       (mst, []) to_edit
   in
   let edited = List.rev edited in
-  let%lwt () =
-    Lwt_list.iter_s
-      (fun (k, v) ->
-        let%lwt got = Mem_mst.get_cid edited_mst k in
-        Alcotest.(check bool)
-          "updated records retrievable" true
-          (Option.value
-             (Option.map (fun x -> Cid.equal v x) got)
-             ~default:false )
-        |> Lwt.return )
-      edited
-  in
+  let%lwt result_map = Mem_mst.build_map edited_mst in
+  List.iter
+    (fun (k, v) ->
+      let got = String_map.find_opt k result_map in
+      Alcotest.(check bool)
+        "updated records retrievable" true
+        (Option.value (Option.map (fun x -> Cid.equal v x) got) ~default:false) )
+    edited ;
   let%lwt total = Mem_mst.leaf_count edited_mst in
   Alcotest.(check int) "leaf count stable after edits" 1000 total ;
   Lwt.return_ok ()
···
   in
   let%lwt total = Mem_mst.leaf_count deleted_mst in
   Alcotest.(check int) "leaf count after deletes" 900 total ;
-  let%lwt () =
-    Lwt_list.iter_s
-      (fun (k, _) ->
-        let%lwt got = Mem_mst.get_cid deleted_mst k in
-        Alcotest.(check bool) "deleted record missing" true (got = None)
-        |> Lwt.return )
-      to_delete
-  in
-  let%lwt () =
-    Lwt_list.iter_s
-      (fun (k, v) ->
-        let%lwt got = Mem_mst.get_cid deleted_mst k in
-        Alcotest.(check bool)
-          "remaining records intact" true
-          (Option.value
-             (Option.map (fun x -> Cid.equal v x) got)
-             ~default:false )
-        |> Lwt.return )
-      the_rest
-  in
+  let%lwt result_map = Mem_mst.build_map deleted_mst in
+  List.iter
+    (fun (k, _) ->
+      let got = String_map.find_opt k result_map in
+      Alcotest.(check bool) "deleted record missing" true (got = None) )
+    to_delete ;
+  List.iter
+    (fun (k, v) ->
+      let got = String_map.find_opt k result_map in
+      Alcotest.(check bool)
+        "remaining records intact" true
+        (Option.value (Option.map (fun x -> Cid.equal v x) got) ~default:false) )
+    the_rest ;
   Lwt.return_ok ()
 
 let test_order_independent () =
···
   in
   let%lwt mst = Mem_mst.add mst "com.example/key1" cid1 in
   let%lwt mst = Mem_mst.add mst "com.example/key1" cid2 in
-  let%lwt got = Mem_mst.get_cid mst "com.example/key1" in
-  ( match got with
+  let%lwt result_map = Mem_mst.build_map mst in
+  ( match String_map.find_opt "com.example/key1" result_map with
   | Some cid ->
       Alcotest.(check bool) "update replaces value" true (Cid.equal cid cid2)
   | None ->

pegasus/lib/repository.ml

···
 module Block_map = User_store.Block_map
 module Lex = Mist.Lex
 module Mst = Mist.Mst.Make (User_store)
+module Cached_store = Mist.Storage.Cache_blockstore (User_store)
+module Cached_mst = Mist.Mst.Make (Cached_store)
 module Mem_mst = Mist.Mst.Make (Mist.Storage.Memory_blockstore)
 module String_map = Lex.String_map
 module Tid = Mist.Tid
···
          (Cid.to_string (Option.get swap_commit))
          (match t.commit with Some (c, _) -> Cid.to_string c | None -> "null") ) ;
   let%lwt block_map = Lwt.map ref (get_map t) in
-  let mst : Mst.t ref = ref (Mst.create t.db prev_commit.data) in
+  let cached_store = Cached_store.create t.db in
+  let mst : Cached_mst.t ref =
+    ref (Cached_mst.create cached_store prev_commit.data)
+  in
   (* ops to emit, built in loop because prev_data (previous cid) is otherwise inaccessible *)
   let commit_ops : commit_evt_op list ref = ref [] in
   let added_leaves = ref Block_map.empty in
···
             added_leaves := Block_map.set cid block !added_leaves ;
             commit_ops :=
               !commit_ops @ [{action= `Create; path; cid= Some cid; prev= None}] ;
-            let%lwt new_mst = Mst.add !mst path cid in
+            let%lwt new_mst = Cached_mst.add !mst path cid in
             mst := new_mst ;
             let refs =
               Util.find_blob_refs value
···
             commit_ops :=
               !commit_ops
               @ [{action= `Update; path; cid= Some new_cid; prev= old_cid}] ;
-            let%lwt new_mst = Mst.add !mst path new_cid in
+            let%lwt new_mst = Cached_mst.add !mst path new_cid in
             mst := new_mst ;
             let refs =
               Util.find_blob_refs value
···
             block_map := String_map.remove path !block_map ;
             commit_ops :=
               !commit_ops @ [{action= `Delete; path; cid= None; prev= cid}] ;
-            let%lwt new_mst = Mst.delete !mst path in
+            let%lwt new_mst = Cached_mst.delete !mst path in
             mst := new_mst ;
             Lwt.return
               (Delete {type'= "com.atproto.repo.applyWrites#deleteResult"}) )