From cae451a35b881f90cbfcbfc943340c8d6ff3c56b Mon Sep 17 00:00:00 2001
From: Vesa Karvonen <vesa.a.j.k@gmail.com>
Date: Sat, 25 Jan 2025 08:14:56 +0000
Subject: [PATCH] Simplify poisoning logic in `Lock`

---
 lib/picos_std.sync/lock.ml | 36 ++++++++++++++++++------------------
 1 file changed, 18 insertions(+), 18 deletions(-)

diff --git a/lib/picos_std.sync/lock.ml b/lib/picos_std.sync/lock.ml
index d144f8b6..2c01ad05 100644
--- a/lib/picos_std.sync/lock.ml
+++ b/lib/picos_std.sync/lock.ml
@@ -1,41 +1,41 @@
 open Picos_std_awaitable
 
-let poisoned = 1 lsl 0
-let no_writers = 1 lsl 1
-let locked = 1 lsl 2
+let no_awaiters = 1 lsl 0
+let locked = 1 lsl 1
 let locked_permanently = 1 lsl (Sys.int_size - 2)
 
 type t = int Awaitable.t
 
 exception Poisoned
 
-let poison t =
+let rec poison t =
   let before = Awaitable.get t in
   if before < locked then invalid_arg "not locked";
   (* Unfortunately we cannot check ownership at this point. *)
-  if before land poisoned = 0 then
-    let _ : int = Awaitable.fetch_and_add t (locked_permanently + poisoned) in
-    Awaitable.broadcast t
+  if before < locked_permanently / 2 then
+    if Awaitable.compare_and_set t before (before + locked_permanently) then
+      Awaitable.broadcast t
+    else poison t
 
 (* *)
 
 let rec acquire_awaiting t before =
   if before < locked then begin
-    (* We know [before] is [0] or [no_writers]. *)
-    let after = locked land lnot no_writers in
+    (* We know [before] is [0] or [no_awaiters]. *)
+    let after = locked land lnot no_awaiters in
     if not (Awaitable.compare_and_set t before after) then
       acquire_awaiting t (Awaitable.get t)
   end
-  else if before land poisoned <> 0 then raise Poisoned
+  else if locked_permanently / 2 <= before then raise Poisoned
   else
-    let after = before land lnot no_writers in
+    let after = before land lnot no_awaiters in
     if before = after || Awaitable.compare_and_set t before after then
       Awaitable.await t after;
     acquire_awaiting t (Awaitable.get t)
 
 let rec acquire_contended t before =
   if locked * 2 <= before then
-    let after = (before - locked) land lnot no_writers in
+    let after = (before - locked) land lnot no_awaiters in
     if Awaitable.compare_and_set t before after then acquire_awaiting t after
     else acquire_contended t (Awaitable.get t)
 
@@ -47,12 +47,12 @@ let acquire t =
 
 let signal_awaiter t =
   (* Lock state was [0] after the [fetch_and_add] in [unlock]. *)
-  if Awaitable.compare_and_set t 0 no_writers then Awaitable.signal t
+  if Awaitable.compare_and_set t 0 no_awaiters then Awaitable.signal t
 
 let release t =
   let prior = Awaitable.fetch_and_add t (-locked) in
-  if prior < locked lor no_writers then signal_awaiter t
-  else if locked_permanently / 2 < prior then
+  if prior < locked lor no_awaiters then signal_awaiter t
+  else if locked_permanently / 2 <= prior then
     (* The lock was poisoned.  We need to undo. *)
     let _ : int = Awaitable.fetch_and_add t locked in
     ()
@@ -69,13 +69,13 @@ let try_acquire t =
   prior < locked
   ||
   let prior = Awaitable.fetch_and_add t (-locked) in
-  prior land poisoned <> 0 && raise Poisoned
+  locked_permanently / 2 <= prior && raise Poisoned
 
 (* *)
 
-let[@inline] create ?padded () = Awaitable.make ?padded no_writers
+let[@inline] create ?padded () = Awaitable.make ?padded no_awaiters
 let[@inline] is_locked t = locked <= Awaitable.get t
-let[@inline] is_poisoned t = Awaitable.get t land poisoned <> 0
+let[@inline] is_poisoned t = locked_permanently / 2 <= Awaitable.get t
 
 module Condition = struct
   type lock = t