Tweaks

lib/picos_std.sync/rwlock.slim.ml

@@ -4,40 +4,32 @@ open Picos_std_awaitable
 
 type t = int Awaitable.t
 
-let poisoned = 1 lsl 0
-let frozen = 1 lsl 1
-let exclusive = 1 lsl 2
-let readers = 1 lsl 3
-let writers = 1 lsl 4
+let frozen = 1 lsl 0
+let poisoned = 1 lsl 1
+let readers = 1 lsl 2
+let writers = 1 lsl 3
 let awaiters = readers lor writers
-let locked = 1 lsl 5
+let locked = 1 lsl 4
+let exclusive = 1 lsl (Sys.int_size - 2)
 
 (* *)
 
-let create ?padded () = Awaitable.make ?padded 0
-
 exception Poisoned
 exception Frozen
 
 let rec lock_awaiting t =
   let before = Awaitable.get t in
-  if before land (poisoned lor frozen) <> 0 then
-    raise (if before land poisoned <> 0 then Poisoned else Frozen)
-  else if before < locked then begin
+  if before < locked then begin
     let after = before lor locked lor exclusive lor writers in
     if not (Awaitable.compare_and_set t before after) then lock_awaiting t
   end
-  else if before land writers = 0 then begin
+  else if before land (poisoned lor frozen) <> 0 then
+    raise (if before land poisoned <> 0 then Poisoned else Frozen)
+  else
     let after = before lor writers in
-    if Awaitable.compare_and_set t before after then begin
-      Awaitable.await t after
-    end;
-    lock_awaiting t
-  end
-  else begin
-    Awaitable.await t before;
+    if before = after || Awaitable.compare_and_set t before after then
+      Awaitable.await t after;
     lock_awaiting t
-  end
 
 let lock t =
   let before = Awaitable.get t in
@@ -48,50 +40,41 @@ let lock t =
     not (Awaitable.compare_and_set t before after)
   then lock_awaiting t
 
-let rec lock_ro_awaiting t backoff =
+let rec lock_ro_awaiting t =
   let before = Awaitable.get t in
-  if before land poisoned <> 0 then raise Poisoned
-  else if before land exclusive = 0 then begin
+  if before < exclusive then begin
     let after = (before + locked) lor readers in
-    if not (Awaitable.compare_and_set t before after) then
-      lock_ro_awaiting t (Backoff.once backoff)
+    if not (Awaitable.compare_and_set t before after) then lock_ro_awaiting t
   end
-  else if before land readers = 0 then begin
+  else if before land poisoned <> 0 then raise Poisoned
+  else
     let after = before lor readers in
-    if Awaitable.compare_and_set t before after then begin
+    if before = after || Awaitable.compare_and_set t before after then
       Awaitable.await t after;
-      lock_ro_awaiting t Backoff.default
-    end
-    else lock_ro_awaiting t (Backoff.once backoff)
-  end
-  else begin
-    Awaitable.await t before;
-    lock_ro_awaiting t Backoff.default
-  end
+    lock_ro_awaiting t
 
 let rec lock_ro_contended t =
   let before = Awaitable.get t in
-  if locked * 2 < before then
+  if locked * 2 <= before then
     let after = (before - locked) lor readers in
-    if Awaitable.compare_and_set t before after then
-      lock_ro_awaiting t Backoff.default
+    if Awaitable.compare_and_set t before after then lock_ro_awaiting t
     else lock_ro_contended t
 
 let lock_ro t =
   let prior = Awaitable.fetch_and_add t locked in
-  if prior land exclusive <> 0 then lock_ro_contended t
+  if exclusive <= prior then lock_ro_contended t
 
-let rec unlock_awaiting t backoff =
+let rec unlock_awaiting t =
   let before = Awaitable.get t in
   if before < locked && before land awaiters <> 0 then
     if Awaitable.compare_and_set t before (before land lnot awaiters) then
       if before land readers <> 0 then Awaitable.broadcast t
       else Awaitable.signal t
-    else unlock_awaiting t (Backoff.once backoff)
+    else unlock_awaiting t
 
 let unlock t =
   let before = Awaitable.get t in
-  if before land exclusive <> 0 then begin
+  if exclusive <= before then begin
     let prior =
       Awaitable.fetch_and_add t
         (-(locked lor exclusive) - (before land awaiters))
@@ -100,14 +83,9 @@ let unlock t =
       if prior land readers <> 0 then Awaitable.broadcast t
       else Awaitable.signal t
   end
-  else begin
+  else
     let prior = Awaitable.fetch_and_add t (-locked) in
-    if prior land awaiters <> 0 then unlock_awaiting t Backoff.default
-  end
-
-let is_locked t = Awaitable.get t land exclusive <> 0
-let is_poisoned t = Awaitable.get t land poisoned <> 0
-let is_frozen t = Awaitable.get t land frozen <> 0
+    if prior land awaiters <> 0 then unlock_awaiting t
 
 let poison t =
   let before = Awaitable.get t in
@@ -119,19 +97,18 @@ let poison t =
 
 let freeze t =
   lock_ro t;
+  let before = ref (Awaitable.get t) in
   while
-    let before = Awaitable.get t in
-    before land frozen = 0
+    !before land frozen = 0
     &&
-    let after = before + (locked lor frozen) in
+    let after = !before + (locked lor frozen) in
     (* This leaves the rwlock as read locked. *)
-    not (Awaitable.compare_and_set t before after)
+    not (Awaitable.compare_and_set t !before after)
   do
-    let _ : Backoff.t = Backoff.once Backoff.default in
-    ()
+    before := Awaitable.get t
   done;
   (* We must wake up any writers waiting to obtain the lock. *)
-  if Awaitable.get t land awaiters <> 0 then Awaitable.broadcast t;
+  if !before land awaiters <> 0 then Awaitable.broadcast t;
   unlock t
 
 let protect t thunk =
@@ -156,6 +133,11 @@ let protect_ro t thunk =
       unlock t;
       Printexc.raise_with_backtrace exn bt
 
+let[@inline] create ?padded () = Awaitable.make ?padded 0
+let[@inline] is_locked t = exclusive <= Awaitable.get t
+let[@inline] is_poisoned t = Awaitable.get t land poisoned <> 0
+let[@inline] is_frozen t = Awaitable.get t land frozen <> 0
+
 module Condition = struct
   include Cond