diff --git a/src/fib/coop_lock.ml b/src/fib/coop_lock.ml new file mode 100644 index 00000000..eb42f792 --- /dev/null +++ b/src/fib/coop_lock.ml @@ -0,0 +1,106 @@ +module A = Atomic_ +module Q = Moonpool_private.Fqueue + +type waiter = unit -> unit + +type 'a state = { + value: 'a; + locked: bool; + waiters: waiter Q.t; +} + +type 'a t = { st: 'a state A.t } [@@unboxed] + +let create x : _ t = + { st = A.make { locked = false; value = x; waiters = Q.empty } } + +let try_lock_ (self : 'a t) : 'a option = + let old = A.get self.st in + if old.locked then + None + else if A.compare_and_set self.st old { old with locked = true } then + Some old.value + else + None + +let rec lock_ (self : 'a t) : 'a = + let old = A.get self.st in + if old.locked then ( + (* suspend and add myself to the [waiters] queue *) + Moonpool.Private.Suspend_.suspend + { + handle = + (fun ~run:_ ~resume k -> + let waiter : waiter = fun () -> resume k @@ Ok () in + if + A.compare_and_set self.st old + { old with waiters = Q.push old.waiters waiter } + then + () + else + (* wakeup now to try again *) + resume k @@ Ok ()); + }; + lock_ self + ) else if A.compare_and_set self.st old { old with locked = true } then + old.value + else + lock_ self + +let unlock_ (self : _ t) (x : 'a) : unit = + let waiter_to_awake = ref None in + while + let old = A.get self.st in + assert old.locked; + (* check if there's a waiter to wake up *) + let waiters = + match Q.pop_exn old.waiters with + | exception Q.Empty -> + waiter_to_awake := None; + old.waiters + | w, ws -> + waiter_to_awake := Some w; + ws + in + let new_st = { locked = false; value = x; waiters } in + not (A.compare_and_set self.st old new_st) + do + () + done; + (* wakeup the next waiter, if any *) + Option.iter (fun f -> f ()) !waiter_to_awake + +let with_lock self f = + let x = lock_ self in + match f x with + | res -> + unlock_ self x; + res + | exception e -> + let bt = Printexc.get_raw_backtrace () in + unlock_ self x; + Printexc.raise_with_backtrace e bt + +let with_lock_update self f = + let x = lock_ self in + match f x with + | new_x, res -> + unlock_ self new_x; + res + | exception e -> + let bt = Printexc.get_raw_backtrace () in + unlock_ self x; + Printexc.raise_with_backtrace e bt + +let with_try_lock self f = + match try_lock_ self with + | None -> f None + | Some x -> + (match f (Some x) with + | res -> + unlock_ self x; + res + | exception e -> + let bt = Printexc.get_raw_backtrace () in + unlock_ self x; + Printexc.raise_with_backtrace e bt) diff --git a/src/fib/coop_lock.mli b/src/fib/coop_lock.mli new file mode 100644 index 00000000..894878e5 --- /dev/null +++ b/src/fib/coop_lock.mli @@ -0,0 +1,12 @@ +(** Cooperative lock *) + +type 'a t + +val create : 'a -> 'a t +(** New lock *) + +val with_lock : 'a t -> ('a -> 'b) -> 'b +(** [with_lock lock f] calls [f] with the lock's content *) + +val with_lock_update : 'a t -> ('a -> 'a * 'b) -> 'b +val with_try_lock : 'a t -> ('a option -> 'b) -> 'b