Fix to allow retry from withing Xt.update and Xt.modify

Author: polytypic

src/kcas/kcas.ml

@@ -535,13 +535,17 @@ module Xt = struct
   let update0 loc f xt lt gt =
     let state = fenceless_get (as_atomic loc) in
     let before = eval state in
-    let after = f before in
-    let state =
-      if before == after && is_obstruction_free xt.casn loc then state
-      else { before; after; casn = xt.casn; awaiters = [] }
-    in
-    xt.cass <- CASN { loc; state; lt; gt; awaiters = [] };
-    before
+    match f before with
+    | after ->
+        let state =
+          if before == after && is_obstruction_free xt.casn loc then state
+          else { before; after; casn = xt.casn; awaiters = [] }
+        in
+        xt.cass <- CASN { loc; state; lt; gt; awaiters = [] };
+        before
+    | exception exn ->
+        xt.cass <- CASN { loc; state; lt; gt; awaiters = [] };
+        raise exn
     [@@inline]
 
   let update loc f xt state' lt gt =

test/kcas/test.ml

@@ -328,7 +328,13 @@ let test_backoff () =
 let test_blocking () =
   let state = Loc.make `Spawned in
   let await state' =
-    Loc.get_as (fun state -> Retry.unless (state == state')) state
+    (* Intentionally test that [Xt.modify] allows retry. *)
+    let tx ~xt =
+      Xt.modify ~xt state @@ fun state ->
+      Retry.unless (state == state');
+      state
+    in
+    Xt.commit { tx }
   in
 
   let a = Loc.make 0 and bs = Array.init 10 @@ fun _ -> Loc.make 0 in