lockfree map, used to make Cache.Count domain safe

Makefile

@@ -1,5 +1,5 @@
 
-.PHONY: build lib doc clean install uninstall test gen gen_ragel gen_metaocaml archive
+.PHONY: build lib doc clean install uninstall test benchs gen gen_ragel gen_metaocaml archive
 
 OCAMLBUILD=ocamlbuild -use-ocamlfind -no-links -j 0
 
@@ -27,6 +27,9 @@ top:
 test:
 		dune runtest $(DUNEFLAGS)
 
+benchs:
+		RUN_BENCHS=1 dune exec --profile bench $(DUNEFLAGS) bench/bench_sharded_hash_trie.exe
+
 doc:
 		dune build $(DUNEFLAGS) @doc
 

bench/bench_sharded_hash_trie.ml

@@ -0,0 +1,159 @@
+(* Compare ExtThread.ShardedHashTrie (64 and 256 shards) with Saturn.Htbl, a flat bucketed assoc list,
+   Hashtbl + Mutex, and (on one domain) a plain Hashtbl.
+
+   dune exec --release bench/bench_sharded_hash_trie.exe *)
+
+open Devkit
+
+type ('k, 'v) ops = { find : 'k -> 'v; get_or_create : 'k -> f:('k -> 'v) -> 'v }
+type 'k impl = { name : string; make : unit -> ('k, int Atomic.t) ops }
+
+module Impls (K : Hashtbl.HashedType) = struct
+  let devkit ?(shard_bits = 6) () = { name = Printf.sprintf "devkit/s%d" shard_bits; make = fun () ->
+    let t = ExtThread.ShardedHashTrie.create ~shard_bits ~hashable:{ equal = K.equal; hash = K.hash } () in
+    { find = ExtThread.ShardedHashTrie.find t; get_or_create = (fun k ~f -> ExtThread.ShardedHashTrie.get_or_create t k ~f) } }
+
+  let saturn = { name = "saturn"; make = fun () ->
+    let t = Saturn.Htbl.create ~hashed_type:(module K) () in
+    let rec get_or_create k ~f =
+      match Saturn.Htbl.find_exn t k with
+      | v -> v
+      | exception Not_found -> let v = f k in if Saturn.Htbl.try_add t k v then v else get_or_create k ~f
+    in
+    { find = Saturn.Htbl.find_exn t; get_or_create } }
+
+  (* fixed buckets of assoc lists, the simplest lock-free option *)
+  let flat = { name = "flat128"; make = fun () ->
+    let b = Array.init 128 (fun _ -> Atomic.make []) in
+    let bucket k = b.(K.hash k land 127) in
+    let rec assoc k = function [] -> raise_notrace Not_found | (k', v) :: tl -> if K.equal k k' then v else assoc k tl in
+    let find k = assoc k (Atomic.get (bucket k)) in
+    let rec get_or_create k ~f =
+      let b = bucket k in
+      let l = Atomic.get b in
+      match assoc k l with
+      | v -> v
+      | exception Not_found -> let v = f k in if Atomic.compare_and_set b l ((k, v) :: l) then v else get_or_create k ~f
+    in
+    { find; get_or_create } }
+
+  let mutex = { name = "hashtbl+mutex"; make = fun () ->
+    let module T = Hashtbl.Make (K) in
+    let t = T.create 16 and m = Mutex.create () in
+    { find = (fun k -> Mutex.protect m (fun () -> T.find t k));
+      get_or_create = (fun k ~f -> Mutex.protect m (fun () ->
+        match T.find t k with v -> v | exception Not_found -> let v = f k in T.add t k v; v)) } }
+
+  (* not domain-safe: baseline for 1-domain runs only *)
+  let plain = { name = "hashtbl (unsafe)"; make = fun () ->
+    let module T = Hashtbl.Make (K) in
+    let t = T.create 16 in
+    { find = T.find t;
+      get_or_create = (fun k ~f -> match T.find t k with v -> v | exception Not_found -> let v = f k in T.add t k v; v) } }
+
+  let tries = [ devkit (); devkit ~shard_bits:8 () ]
+  let all = tries @ [ saturn; flat; mutex ]
+  let growing = tries @ [ saturn; mutex ] (* flat128 degrades to long lists *)
+end
+
+module S = Impls (struct include String let hash = Hashtbl.hash end)
+module I = Impls (struct type t = int let equal = Int.equal let hash = Hashtbl.hash end)
+
+let counter _ = Atomic.make 0
+
+let on_domains n f = List.init n (fun d -> Domain.spawn (fun () -> f d)) |> List.iter Domain.join
+
+let report ~title ~ops_per_call samples =
+  Printf.printf "\n== %s\n" title;
+  List.iter (fun (name, ts) ->
+    let rates = List.map (fun (t : Benchmark.t) -> Int64.to_float t.iters *. float ops_per_call /. t.wall /. 1e6) ts in
+    let rates = List.sort compare rates in
+    Printf.printf "  %-16s %8.1f Mops/s (median of %d)\n%!" name (List.nth rates (List.length rates / 2)) (List.length rates))
+    samples
+
+(* the unsynchronized Hashtbl only makes sense on one domain *)
+let with_plain domains impls plain = if domains = 1 then impls @ [ plain ] else impls
+
+let bench ~title ~ops_per_call impls run =
+  let samples = Benchmark.throughputN ~style:Benchmark.Nil ~repeat:5 1
+      (List.map (fun impl -> impl.name, run, impl) impls) in
+  report ~title ~ops_per_call samples
+
+(* Cache.Count: a few constant keys, hit on every call *)
+let count_like domains =
+  let keys = Array.init 20 (Printf.sprintf "metric_name_%d") in
+  let per_domain = 1_000_000 in
+  bench ~title:(Printf.sprintf "count-like: 20 string keys, hits + incr, %d domain(s)" domains)
+    ~ops_per_call:(per_domain * domains) (with_plain domains S.all S.plain)
+    (fun impl ->
+       let t = impl.make () in
+       Array.iter (fun k -> ignore (t.get_or_create k ~f:counter)) keys;
+       on_domains domains (fun d ->
+         for i = 0 to per_domain - 1 do
+           Atomic.incr (t.get_or_create keys.((i + d) mod 20) ~f:counter)
+         done))
+
+(* growing table: every call is a miss and an insert *)
+let inserts domains =
+  let n = 200_000 in
+  bench ~title:(Printf.sprintf "inserts: %d fresh int keys, %d domain(s)" n domains)
+    ~ops_per_call:n (with_plain domains I.all I.plain)
+    (fun impl ->
+       let t = impl.make () in
+       on_domains domains (fun d ->
+         let i = ref d in
+         while !i < n do ignore (t.get_or_create !i ~f:counter); i := !i + domains done))
+
+(* read-mostly: 10k warm keys, one insert of a new key every [insert_every] ops, finds otherwise.
+   The table grows by [per_domain / insert_every] keys per domain and call. *)
+let mixed ~label ~insert_every domains =
+  let warm = 10_000 and per_domain = 2_000_000 in
+  bench ~title:(Printf.sprintf "%s: %d warm int keys, 1 insert per %d ops, %d domain(s)" label warm insert_every domains)
+    ~ops_per_call:(per_domain * domains) (with_plain domains I.growing I.plain)
+    (fun impl ->
+       let t = impl.make () in
+       for k = 0 to warm - 1 do ignore (t.get_or_create k ~f:counter) done;
+       on_domains domains (fun d ->
+         let fresh = ref (warm + d) in
+         for i = 0 to per_domain - 1 do
+           if i mod insert_every = 0 then (ignore (t.get_or_create !fresh ~f:counter); fresh := !fresh + domains)
+           else ignore (t.find ((i * 7919) mod warm))
+         done))
+
+let mixes = [ "mixed90", 10; "mixed99", 100; "mixed99.9", 1000 ]
+let all_mixed () = List.iter (fun (label, insert_every) -> List.iter (mixed ~label ~insert_every) [ 1; 4; 8 ]) mixes
+
+(* sanity check every implementation before timing it *)
+let self_check () =
+  List.iter (fun impl ->
+    let t = impl.make () in
+    let n = 100_000 in
+    on_domains 4 (fun d -> let i = ref d in while !i < n do ignore (t.get_or_create !i ~f:(fun k -> Atomic.make k)); i := !i + 4 done);
+    for k = 0 to n - 1 do if Atomic.get (t.find k) <> k then failwith (impl.name ^ ": self-check failed") done)
+    I.all
+
+let alloc_per_hit () =
+  Printf.printf "\n== minor words per hit (20 string keys, 1 domain)\n";
+  let keys = Array.init 20 (Printf.sprintf "metric_name_%d") in
+  List.iter (fun impl ->
+    let t = impl.make () in
+    Array.iter (fun k -> ignore (t.get_or_create k ~f:counter)) keys;
+    let n = 1_000_000 in
+    let before = Gc.minor_words () in
+    for i = 0 to n - 1 do ignore (t.get_or_create keys.(i mod 20) ~f:counter) done;
+    Printf.printf "  %-16s %6.2f\n%!" impl.name ((Gc.minor_words () -. before) /. float n))
+    S.all
+
+let () =
+  self_check ();
+  match Sys.argv with
+  | [| _; "mixed" |] -> all_mixed ()
+  | [| _; "single" |] ->
+    count_like 1;
+    inserts 1;
+    List.iter (fun (label, insert_every) -> mixed ~label ~insert_every 1) mixes
+  | _ ->
+    alloc_per_hit ();
+    List.iter count_like [ 1; 4; 8 ];
+    List.iter inserts [ 1; 4 ];
+    all_mixed ()

bench/dune

@@ -0,0 +1,6 @@
+; optional: needs saturn and benchmark. Build with make benchs.
+(executable
+ (name bench_sharded_hash_trie)
+ (optional)
+ (enabled_if (= %{profile} bench))
+ (libraries devkit saturn benchmark unix))

cache.ml

@@ -60,37 +60,39 @@ module TimeLimited2(E: Set.OrderedType)
 end
 
 module Count = struct
-  open Hashtbl
-  type 'a t = ('a, int ref) Hashtbl.t
-  let create () : 'a t = create 16
-  let clear = Hashtbl.clear
-  let entry t x = match find t x with r -> r | exception Not_found -> let r = ref 0 in Hashtbl.add t x r; r
-  let plus t x n = entry t x += n
-  let minus t x n = entry t x -= n
+  module H = ExtThread.ShardedHashTrie
+  (* counters are atomic, the table only grows (or is cleared) *)
+  type 'a t = ('a, int Atomic.t) H.t
+  let create () : 'a t = H.create ~shard_bits:6 ()
+  let clear = H.clear
+  let entry t x = H.get_or_create t x ~f:(fun _ -> Atomic.make 0)
+  let plus t x n = ignore (Atomic.fetch_and_add (entry t x) n : int)
+  let minus t x n = plus t x (-n)
   let of_enum e = let h = create () in Enum.iter (fun (k,n) -> plus h k n) e; h
   let of_list l = of_enum @@ List.enum l
   let add t x = plus t x 1
   let del t x = minus t x 1
-  let enum t = enum t |> Enum.map (fun (k,n) -> k, !n)
-  let iter t f = iter (fun k n -> f k !n) t
-  let fold t f acc = Hashtbl.fold (fun k n acc -> f k !n acc) t acc
-  let count t k = match Hashtbl.find t k with n -> !n | exception Not_found -> 0
-  let count_all t = Hashtbl.fold (fun _ n acc -> acc + !n) t 0
-  let size = Hashtbl.length
-  let show t ?(sep=" ") f = enum t |>
-    List.of_enum |> List.sort ~cmp:(Action.compare_by fst) |>
+  let iter t f = H.iter t (fun k n -> f k (Atomic.get n))
+  let fold t f acc = H.fold t (fun k n acc -> f k (Atomic.get n) acc) acc
+  let to_list t = fold t (fun k n acc -> (k,n) :: acc) []
+  let enum t = List.enum (to_list t)
+  let count t k = match H.find_opt t k with Some n -> Atomic.get n | None -> 0
+  let count_all t = fold t (fun _ n acc -> acc + n) 0
+  let size = H.length
+  let show t ?(sep=" ") f = to_list t |>
+    List.sort ~cmp:(Action.compare_by fst) |>
     List.map (fun (x,n) -> sprintf "%S: %u" (f x) n) |>
     String.concat sep
-  let show_sorted t ?limit ?(sep="\n") f = enum t |>
-    List.of_enum |> List.sort ~cmp:(flip @@ Action.compare_by snd) |>
+  let show_sorted t ?limit ?(sep="\n") f = to_list t |>
+    List.sort ~cmp:(flip @@ Action.compare_by snd) |>
     (match limit with None -> id | Some n -> List.take n) |>
     List.map (fun (x,n) -> sprintf "%6d : %S" n (f x)) |>
     String.concat sep
   let stats t ?(cmp=compare) f =
-    if Hashtbl.length t = 0 then
+    let a = to_list t |> Array.of_list in
+    if Array.length a = 0 then
       "<empty>"
     else
-      let a = Array.of_enum (enum t) in
       let total = Array.fold_left (fun t (_,n) -> t + n) 0 a in
       let half = total / 2 in
       let cmp (x,_) (y,_) = cmp x y in
@@ -108,10 +110,10 @@ module Count = struct
       sprintf "total %d median %s min %s max %s"
         total (match !med with None -> "?" | Some x -> show x) (show mi) (show ma)
   let distrib t =
-    if Hashtbl.length t = 0 then
+    let a = to_list t |> Array.of_list in
+    if Array.length a = 0 then
       [||]
     else
-      let a = Array.of_enum (enum t) in
       let total = Array.fold_left (fun t (_,n) -> t + n) 0 a in
       let limits = Array.init 10 (fun i -> total * (i + 1) / 10) in
       let cmp (x,_) (y,_) = compare (x:float) y in
@@ -133,7 +135,7 @@ module Count = struct
     let data = show_sorted t ?limit ~sep f in
     let stats = stats t ?cmp f in
     stats^sep^data
-  let names (t : 'a t) = List.of_enum @@ Hashtbl.keys t
+  let names (t : 'a t) = H.fold t (fun k _ acc -> k :: acc) []
 end
 
 

cache.mli

@@ -38,7 +38,12 @@ module LRU(K : Hashtbl.HashedType) : sig
   val lfu_free : 'v t -> int
 end
 
-(** Count elements *)
+(** Count elements. Domain safe.
+
+   Collection-wide operations such as [iter], [fold], [size],
+   [clear], etc. can race against individual modifications (ie it's
+   possible that a call to [clear t] races against [plus t "x" 42]
+   and an empty structure is never observable) *)
 module Count : sig
   type 'a t
   val create : unit -> 'a t

devkit.opam

@@ -12,9 +12,11 @@ build: [
 ]
 depends: [
   "ocaml" {>= "5.0"}
-  "dune" {>= "2.0"}
+  "dune" {>= "2.3"}
   ("extlib" {>= "1.7.1"} | "extlib-compat" {>= "1.7.1"})
   "ounit2"
+  "qcheck-core" {with-test}
+  "benchmark" {with-dev-setup}
   "camlzip"
   "libevent" {>= "0.8.0"}
   "curl" {>= "0.10.0"}

dune-project

@@ -1,3 +1,3 @@
-(lang dune 2.0)
+(lang dune 2.3)
 (name devkit)
 (implicit_transitive_deps false)