Lockfree bag based on the hash table (#173)

Lockfree bag based on the hash table

README.md

@@ -39,6 +39,7 @@ You can learn more about the **motivation** behind `Saturn` through the implemen
     - [Lock-free Multiple Producers Single Consumer Queue](#lock-free-multiple-producers-single-consumer-queue)
     - [Lock-free Skip List](#lock-free-skip-list)
     - [Lock-free Hash Table](#lock-free-hash-table)
+    - [Lock-free Bag](#lock-free-bag)
 - [About the Unsafe Data Structures](#about-the-unsafe-data-structures)
 - [Usage](#usage)
     - [Data Structures with Domain Roles](#data-structures-with-domain-roles)
@@ -141,6 +142,10 @@ opam install saturn
 - **Description**: A resizable lock-free hash table with a snapshot mechanism.
 - **Recommendation**: Contains useful high-level operations designed to work as building blocks of non-blocking algorithms.
 
+### Lock-free Bag
+
+- **Module**: [Bag](https://ocaml-multicore.github.io/saturn/Saturn/Bag/index.html)
+- **Description**: A resizable lock-free bag based on the hash table. The `pop` functions returns a random value contained on the bag. 
 
 # About the Unsafe Data Structures
 

src/bag.ml

@@ -0,0 +1,36 @@
+module Key = struct
+  type t = int
+
+  let equal = Int.equal
+  let hash = Fun.id
+end
+
+type 'v t = (int, 'v) Htbl.t
+
+let create () = Htbl.create ~hashed_type:(module Key) ()
+
+(* *)
+
+let rec push t value =
+  let key = Int64.to_int (Random.bits64 ()) in
+  if not (Htbl.try_add t key value) then push t value
+
+(* *)
+
+exception Empty
+
+type ('a, _) poly = Option : ('a, 'a option) poly | Value : ('a, 'a) poly
+
+let rec pop_as : type a r. a t -> (a, r) poly -> r =
+ fun t poly ->
+  match Htbl.find_random_exn t with
+  | key -> begin
+      match Htbl.remove_exn t key with
+      | value -> ( match poly with Option -> Some value | Value -> value)
+      | exception Not_found -> pop_as t poly
+    end
+  | exception Not_found -> (
+      match poly with Option -> None | Value -> raise Empty)
+
+let pop_exn t = pop_as t Value
+let pop_opt t = pop_as t Option

src/bag.mli

@@ -0,0 +1,62 @@
+(** Randomized lock-free bag *)
+
+(** {1 API} *)
+
+type !'v t
+(** Represents a lock-free bag of elements of type 'v *)
+
+val create : unit -> 'v t
+(** [create ()] creates a new empty lock-free bag. *)
+
+val push : 'v t -> 'v -> unit
+(** [push bag elt] adds [elt] to the [bag]. *)
+
+exception Empty
+(** Raised when {!pop_exn} is applied to an empty bag. *)
+
+val pop_exn : 'v t -> 'v
+(** [pop_exn bag] removes and returns a random element of the [bag].
+
+    @raise Empty if the [bag] is empty. *)
+
+val pop_opt : 'v t -> 'v option
+(** [pop_opt bag] removes and returns [Some] of a random element of the [bag] 
+and [None] if the [bag] is empty. *)
+
+(** {1 Example} 
+
+{[
+    # Random.init 0
+    - : unit = ()
+    # module Bag = Saturn.Bag
+    module Bag = Saturn.Bag
+    # let t : string Bag.t = Bag.create ()
+    val t : string Bag.t = <abstr>
+
+    # let planets = ["Mercury"; "Venus"; "Earth"; "Mars"; "Jupiter"; "Saturn"; "Uranus"; "Neptune"]
+    val planets : string list =
+      ["Mercury"; "Venus"; "Earth"; "Mars"; "Jupiter"; "Saturn"; "Uranus";
+       "Neptune"]
+    # List.iter (Bag.push t) planets
+    - : unit = ()
+    # Bag.pop_exn t
+    - : string = "Neptune"
+    # Bag.pop_opt t
+    - : string option = Some "Saturn"
+    # Bag.pop_exn t
+    - : string = "Mercury"
+    # Bag.pop_exn t
+    - : string = "Mars"
+    # Bag.pop_exn t
+    - : string = "Earth"
+    # Bag.pop_exn t
+    - : string = "Venus"
+    # Bag.pop_exn t
+    - : string = "Uranus"
+    # Bag.pop_exn t
+    - : string = "Jupiter"
+    # Bag.pop_exn t
+    Exception: Saturn__Bag.Empty.
+]} 
+
+*)

src/dune

@@ -42,5 +42,5 @@ let () =
    (<> %{os_type} Win32)
    (>= %{ocaml_version} 5.0.0)))
  (libraries saturn)
- (files treiber_stack.mli bounded_stack.mli ws_deque.mli mpsc_queue.mli skiplist.mli))
+ (files treiber_stack.mli bounded_stack.mli ws_deque.mli mpsc_queue.mli skiplist.mli bag.mli))
 |}

src/htbl/htbl_intf.mli

@@ -168,5 +168,8 @@ module type HTBL = sig
 
       # Htbl.remove_all t |> List.of_seq
       - : (int * string) list = [(101, "Basics"); (42, "The answer")]
-    ]} *)
+    ]} 
+    
+    The lockfree bag (see {!Saturn.Bag}) is implemented using this hash table.
+*)
 end

src/saturn.ml

@@ -40,3 +40,4 @@ module Size = Size
 module Skiplist = Skiplist
 module Htbl = Htbl
 module Htbl_unsafe = Htbl_unsafe
+module Bag = Bag

src/saturn.mli

@@ -44,3 +44,4 @@ module Skiplist = Skiplist
 module Size = Size
 module Htbl = Htbl
 module Htbl_unsafe = Htbl_unsafe
+module Bag = Bag

test/bag/dune

@@ -0,0 +1,18 @@
+(test
+ (package saturn)
+ (name qcheck_bag)
+ (modules qcheck_bag)
+ (libraries
+  saturn
+  barrier
+  qcheck
+  qcheck-core
+  qcheck-alcotest
+  alcotest
+  domain_shims))
+
+(test
+ (package saturn)
+ (name stm_bag)
+ (modules stm_bag)
+ (libraries saturn qcheck-core qcheck-stm.stm stm_run))

test/bag/qcheck_bag.ml

@@ -0,0 +1,68 @@
+module Bag = Saturn.Bag
+
+let pop_all bag =
+  let rec loop acc =
+    match Bag.pop_opt bag with Some x -> loop (x :: acc) | None -> acc
+  in
+  loop []
+
+let tests =
+  QCheck.
+    [
+      Test.make ~name:"sequential"
+        (pair small_nat (list int))
+        (fun (npop, lpush) ->
+          let bag = Bag.create () in
+          List.iter (Bag.push bag) lpush;
+          let popped = List.init npop (fun _ -> Bag.pop_opt bag) in
+          let popped =
+            List.filter Option.is_some popped |> List.map Option.get
+          in
+
+          let remaining = pop_all bag in
+
+          List.for_all (fun x -> List.mem x lpush) popped
+          && List.for_all (fun x -> List.mem x lpush) remaining
+          && List.sort Int.compare (popped @ remaining)
+             = List.sort Int.compare lpush);
+      Test.make ~name:"parallel"
+        (pair small_nat (list int))
+        (fun (npop, lpush) ->
+          let bag = Bag.create () in
+          let barrier = Barrier.create 2 in
+
+          let domain1 =
+            Domain.spawn @@ fun () ->
+            Barrier.await barrier;
+            List.iter
+              (fun elt ->
+                Bag.push bag elt;
+                Domain.cpu_relax ())
+              lpush
+          in
+          let domain2 =
+            Domain.spawn @@ fun () ->
+            Barrier.await barrier;
+            List.init npop (fun _ ->
+                Domain.cpu_relax ();
+                Bag.pop_opt bag)
+          in
+
+          let () = Domain.join domain1 in
+          let popped = Domain.join domain2 in
+
+          let popped =
+            List.filter Option.is_some popped |> List.map Option.get
+          in
+
+          let remaining = pop_all bag in
+
+          List.for_all (fun x -> List.mem x lpush) popped
+          && List.for_all (fun x -> List.mem x lpush) remaining
+          && List.sort Int.compare (popped @ remaining)
+             = List.sort Int.compare lpush);
+    ]
+
+let () =
+  let to_alcotest = List.map QCheck_alcotest.to_alcotest in
+  Alcotest.run "QCheck Bag" [ ("test_sequential", to_alcotest tests) ]

test/bag/stm_bag.ml

@@ -0,0 +1,42 @@
+open QCheck
+open STM
+module Bag = Saturn.Bag
+
+(* Only check that the size of the bag stays consistent. *)
+
+module Spec = struct
+  type cmd = Push | Pop
+
+  let show_cmd c = match c with Push -> "Push ()" | Pop -> "Pop"
+
+  module Sint = Set.Make (Int)
+
+  type state = int
+  type sut = unit Bag.t
+
+  let arb_cmd _s =
+    QCheck.make ~print:show_cmd (Gen.oneof [ Gen.return Push; Gen.return Pop ])
+
+  let init_state = 0
+  let init_sut () = Bag.create ()
+  let cleanup _ = ()
+
+  let next_state c s =
+    match c with Push -> s + 1 | Pop -> if s > 0 then s - 1 else s
+
+  let precond _ _ = true
+
+  let run c d =
+    match c with
+    | Push -> Res (unit, Bag.push d ())
+    | Pop -> Res (option unit, Bag.pop_opt d)
+
+  let postcond c (s : state) res =
+    match (c, res) with
+    | Push, Res ((Unit, _), _res) -> true
+    | Pop, Res ((Option Unit, _), res) ->
+        if s > 0 then res = Some () else res = None
+    | _, _ -> false
+end
+
+let () = Stm_run.run ~name:"Saturn.Bag" (module Spec) |> exit