WIP: Add blocking

polytypic · polytypic · commit 9e6ea54d03ec · 2023-03-05T13:25:51.000+02:00
diff --git a/README.md b/README.md
@@ -52,6 +52,7 @@ is distributed under the [ISC license](LICENSE.md).
   - [Race to cooperate](#race-to-cooperate)
     - [A three-stack lock-free queue](#a-three-stack-lock-free-queue)
     - [A rehashable lock-free hash table](#a-rehashable-lock-free-hash-table)
+  - [Blocking](#blocking)
 - [Development](#development)
 
 ## A quick tour
@@ -515,25 +516,6 @@ Or transfer elements between different transactional data structures:
 The ability to compose transactions allows algorithms and data-structures to be
 used for a wider variety of purposes.
 
-#### About transactions
-
-The transaction mechanism provided by **kcas** is quite intentionally designed
-to be very simple and efficient. This also means that it cannot provide certain
-features, because adding such features would either add significant dependencies
-or overheads to the otherwise simple and efficient implementation. In
-particular, the transactions provided by **kcas** do not directly provide
-blocking or the ability to wait for changes to shared memory locations before
-retrying a transaction. The way
-[`commit`](https://ocaml-multicore.github.io/kcas/doc/kcas/Kcas/Tx/index.html#val-commit)
-works is that it simply retries the transaction in case it failed. To avoid
-contention, a
-[`backoff`](https://ocaml-multicore.github.io/kcas/doc/kcas/Kcas/Backoff/index.html)
-mechanism is used, but otherwise
-[`commit`](https://ocaml-multicore.github.io/kcas/doc/kcas/Kcas/Tx/index.html#val-commit)
-will essentially perform a
-[busy-wait](https://en.wikipedia.org/wiki/Busy_waiting), which should usually be
-avoided.
-
 ### Programming with explicit transaction log passing
 
 The [`Xt`](https://ocaml-multicore.github.io/kcas/doc/kcas/Kcas/Xt/index.html)
@@ -1230,6 +1212,89 @@ What we have here is a lock-free hash table with rehashing that should not be
 highly prone to starvation. In other respects this is a fairly naive hash table
 implementation. You might want to think about various ways to improve upon it.
 
+### Blocking
+
+Consider the following approach to unconditionally pop an element from a stack:
+
+```ocaml
+# let pop ~xt stack =
+    match Xt.update ~xt stack tl_safe with
+    | [] -> raise Exit
+    | elem :: _ -> elem
+val pop : xt:'a Xt.t -> 'b list Loc.t -> 'b = <fun>
+```
+
+A transaction that attempts to use the above on an empty stack will be keep on
+retrying until some other domain pushes an element to the stack. This sort of
+[busy-wait](https://en.wikipedia.org/wiki/Busy_waiting) should generally be
+avoided.
+
+```ocaml
+# let stack () : _ stack = Loc.make ~awaitable:true []
+val stack : unit -> 'a stack = <fun>
+```
+
+```ocaml
+# let scheduler : Xt.scheduler =
+    let open struct
+      type t = {
+        mutex : Mutex.t;
+        condition : Condition.t;
+        mutable signaled : bool;
+      }
+    end in
+    let key = Domain.DLS.new_key @@ fun () -> {
+      mutex = Mutex.create ();
+      condition = Condition.create ();
+      signaled = false;
+    } in
+    fun suspend ->
+      let t = Domain.DLS.get key in
+      suspend (fun () ->
+        Mutex.lock t.mutex;
+        t.signaled <- true;
+        Mutex.unlock t.mutex;
+        Condition.broadcast t.condition);
+      Mutex.lock t.mutex;
+      while not t.signaled do
+        Condition.wait t.condition t.mutex
+      done;
+      t.signaled <- false;
+      Mutex.unlock t.mutex
+val scheduler : Xt.scheduler = <fun>
+```
+
+```ocaml
+# let push ~xt stack elem = Xt.of_tx ~xt (push stack elem)
+val push : xt:'a Xt.t -> 'b list Loc.t -> 'b -> unit = <fun>
+```
+
+```ocaml
+# let input_stack : int stack = stack ()
+val input_stack : int stack = <abstr>
+# let output_stack : int stack = stack ()
+val output_stack : int stack = <abstr>
+```
+
+```ocaml
+# let other = Domain.spawn @@ fun () ->
+    let x = Xt.commit ~scheduler { tx = pop input_stack } in
+    Xt.commit { tx = push output_stack x }
+val other : unit Domain.t = <abstr>
+```
+
+```ocaml
+# Xt.commit { tx = push input_stack 42 }
+- : unit = ()
+# Xt.commit ~scheduler { tx = pop output_stack }
+- : int = 42
+```
+
+```ocaml
+# Domain.join other
+- : unit = ()
+```
+
 ## Development
 
 ### Formatting
diff --git a/src/kcas.ml b/src/kcas.ml
@@ -5,14 +5,23 @@
 
 module Backoff = Backoff
 
+module AwaitableId = struct
+  let id = Atomic.make (Int.min_int asr 1)
+  let get_unique () = Atomic.fetch_and_add id 1
+  let is id = id < 0
+  let to_awaiters id = Int.min_int - 1 - id
+end
+
 module Id = struct
   let id = Atomic.make 1
   let get_unique () = Atomic.fetch_and_add id 1
 end
 
+type awaiter = [ `Init | `Resumed | `Waiting of unit -> unit ] Atomic.t
 type determined = [ `After | `Before ]
 
-type 'a loc = { state : 'a state Atomic.t; id : int }
+type 'a loc = { state : 'a state Atomic.t; id : int; awaiters : awaiters }
+and awaiters = awaiter list loc
 and 'a state = { mutable before : 'a; mutable after : 'a; mutable casn : casn }
 and cass = CASN : 'a loc * 'a state * cass * cass -> cass | NIL : cass
 and casn = status Atomic.t
@@ -190,11 +199,26 @@ let cas loc before state =
   && Atomic.compare_and_set loc.state state' state
   [@@inline]
 
+let make_awaiters id =
+  let rec awaiters =
+    {
+      state = Atomic.make @@ new_state [];
+      id = AwaitableId.to_awaiters id;
+      awaiters;
+    }
+  in
+  awaiters
+
 module Loc = struct
   type 'a t = 'a loc
 
-  let make after =
-    { state = Atomic.make @@ new_state after; id = Id.get_unique () }
+  let make ?(awaitable = false) after =
+    let state = Atomic.make @@ new_state after in
+    let id =
+      if awaitable then AwaitableId.get_unique () else Id.get_unique ()
+    in
+    let awaiters = if awaitable then make_awaiters id else Obj.magic () in
+    { state; id; awaiters }
 
   let get_id loc = loc.id [@@inline]
 
@@ -446,9 +470,41 @@ module Xt = struct
 
   let call { tx } = tx [@@inline]
 
+  let rec take_awaiters ~xt awaiters = function
+    | NIL -> awaiters
+    | CASN (loc, state, l, r) ->
+        let awaiters =
+          if l != NIL then take_awaiters ~xt awaiters l else awaiters
+        in
+        if AwaitableId.is loc.id then
+          let awaiters =
+            if state.before != state.after then
+              match exchange ~xt loc.awaiters [] with
+              | [] -> awaiters
+              | aws -> aws :: awaiters
+            else awaiters
+          in
+          take_awaiters ~xt awaiters r
+        else awaiters
+
+  let resume_awaiters awaiters =
+    awaiters
+    |> List.iter @@ List.iter
+       @@ fun awaiter ->
+       match Atomic.exchange awaiter `Resumed with
+       | `Waiting resume -> resume ()
+       | _ -> ()
+    [@@inline never]
+
+  let resume_awaiters = function
+    | [] -> ()
+    | awaiters -> resume_awaiters awaiters
+    [@@inline]
+
   let attempt (mode : Mode.t) tx =
     let xt = { casn = Atomic.make (mode :> status); cass = NIL } in
     let result = tx ~xt in
+    let awaiters = take_awaiters ~xt [] xt.cass in
     match xt.cass with
     | NIL -> result
     | CASN (loc, state, NIL, NIL) ->
@@ -457,20 +513,83 @@ module Xt = struct
           let before = state.before in
           state.before <- state.after;
           if cas loc before state then result else exit ()
-    | cass -> if determine_for_owner xt.casn cass then result else exit ()
+    | cass ->
+        if determine_for_owner xt.casn cass then (
+          resume_awaiters awaiters;
+          result)
+        else exit ()
 
-  let rec commit backoff mode tx =
-    match attempt mode tx with
-    | result -> result
-    | exception Mode.Interference ->
-        commit (Backoff.once backoff) Mode.lock_free tx
-    | exception Exit -> commit (Backoff.once backoff) mode tx
+  let attempt ?(mode = Mode.lock_free) tx = attempt mode tx.tx [@@inline]
 
-  let commit ?(backoff = Backoff.default) ?(mode = Mode.obstruction_free) tx =
-    commit backoff mode tx.tx
-    [@@inline]
+  type scheduler = ((unit -> unit) -> unit) -> unit
+
+  let rec add_awaiter ~xt awaiter = function
+    | NIL -> ()
+    | CASN (loc, _, l, r) ->
+        if l != NIL then add_awaiter ~xt awaiter l;
+        if AwaitableId.is loc.id then (
+          modify ~xt loc.awaiters (List.cons awaiter);
+          add_awaiter ~xt awaiter r)
+
+  let rec reset casn = function
+    | NIL -> NIL
+    | CASN (loc, state, l, r) as old ->
+        let l' = reset casn l and r' = reset casn r in
+        if is_cmp casn state then
+          if l == l' && r == r' then old else CASN (loc, state, l', r')
+        else
+          let state' = Atomic.get loc.state in
+          let current = eval state' in
+          if current != state.before then exit () else CASN (loc, state', l, r)
 
-  let attempt ?(mode = Mode.lock_free) tx = attempt mode tx.tx [@@inline]
+  let rec commit backoff (mode : Mode.t) scheduler_opt tx =
+    let xt = { casn = Atomic.make (mode :> status); cass = NIL } in
+    match tx ~xt with
+    | result -> (
+        let awaiters = take_awaiters ~xt [] xt.cass in
+        match determine_for_owner xt.casn xt.cass with
+        | true ->
+            resume_awaiters awaiters;
+            result
+        | false -> commit (Backoff.once backoff) mode scheduler_opt tx
+        | exception Mode.Interference ->
+            commit (Backoff.once backoff) Mode.lock_free scheduler_opt tx)
+    | exception Exit -> (
+        match scheduler_opt with
+        | None -> commit (Backoff.once backoff) mode scheduler_opt tx
+        | Some scheduler -> (
+            match reset xt.casn xt.cass with
+            | cass -> (
+                xt.cass <- cass;
+                let self = Atomic.make `Init in
+                add_awaiter ~xt self cass;
+                if xt.cass == cass then
+                  commit (Backoff.once backoff) mode scheduler_opt tx
+                else
+                  match determine_for_owner xt.casn xt.cass with
+                  | true ->
+                      if Atomic.get self == `Init then
+                        scheduler (fun resume ->
+                            if
+                              not
+                                (Atomic.compare_and_set self `Init
+                                   (`Waiting resume))
+                            then resume ());
+                      (* TODO: remove awaiters *)
+                      commit
+                        (Backoff.once (Backoff.reset backoff))
+                        mode scheduler_opt tx
+                  | false -> commit (Backoff.once backoff) mode scheduler_opt tx
+                  | exception Mode.Interference ->
+                      commit (Backoff.once backoff) Mode.lock_free scheduler_opt
+                        tx)
+            | exception Exit ->
+                commit (Backoff.once backoff) mode scheduler_opt tx))
+
+  let commit ?(backoff = Backoff.default) ?(mode = Mode.obstruction_free)
+      ?scheduler { tx } =
+    commit backoff mode scheduler tx
+    [@@inline]
 
   let of_tx tx ~xt =
     let (_, cass), x = tx (xt.casn, xt.cass) in
diff --git a/src/kcas.mli b/src/kcas.mli
@@ -13,9 +13,13 @@ module Loc : sig
   type 'a t
   (** Type of shared memory locations. *)
 
-  val make : 'a -> 'a t
+  val make : ?awaitable:bool -> 'a -> 'a t
   (** [make initial] creates a new shared memory location with the [initial]
-      value. *)
+      value.
+
+      By default a location is considered to be unawaitable to avoid overheads.
+      {!Xt.commit} cannot block to await for changes to unawaitable locations.
+      To make a location awaitable, simply pass [~awaitable:true]. *)
 
   val get_id : 'a t -> int
   (** [get_id r] returns the unique id of the shared memory location [r]. *)
@@ -411,7 +415,13 @@ module Xt : sig
       the transaction or returns the result of the transaction.  The default for
       [attempt] is {!Mode.lock_free}. *)
 
-  val commit : ?backoff:Backoff.t -> ?mode:Mode.t -> 'a tx -> 'a
+  type scheduler = ((unit -> unit) -> unit) -> unit
+  (** Type of a scheduler that can be called to suspend execution of the current
+      domain or fiber until the resume function given by the scheduler is
+      called. *)
+
+  val commit :
+    ?backoff:Backoff.t -> ?mode:Mode.t -> ?scheduler:scheduler -> 'a tx -> 'a
   (** [commit tx] repeats [attempt tx] until it does not raise [Exit] or
       {!Mode.Interference} and then either returns or raises whatever attempt
       returned or raised.
@@ -422,7 +432,8 @@ module Xt : sig
 
       Note that, aside from using exponential backoff to reduce contention, the
       transaction mechanism has no way to intelligently wait until shared memory
-      locations are modified by other domains. *)
+      locations are modified by other domains unless a suitable {!scheduler} is
+      given. *)
 
   (** {1 Conversions} *)
 
