From 975fdaf9cd3dff81287f11e33ce2290480ac0186 Mon Sep 17 00:00:00 2001
From: Vesa Karvonen <vesa.a.j.k@gmail.com>
Date: Fri, 9 Jun 2023 16:54:04 +0300
Subject: [PATCH] Change/Fix to also snapshot and rollback post commit actions

---
 src/kcas/kcas.ml  | 9 ++++++---
 test/kcas/test.ml | 8 ++++++++
 2 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/src/kcas/kcas.ml b/src/kcas/kcas.ml
index be2c2252..5771508e 100644
--- a/src/kcas/kcas.ml
+++ b/src/kcas/kcas.ml
@@ -727,10 +727,13 @@ module Xt = struct
               and gt = rollback casn gt_mark gt in
               CASN { loc; state; lt; gt; awaiters = [] })
 
-  type 'x snap = cass
+  type 'x snap = cass * Action.t
 
-  let snapshot ~xt = xt.cass
-  let rollback ~xt snap = xt.cass <- rollback xt.casn snap xt.cass
+  let snapshot ~xt = (xt.cass, xt.post_commit)
+
+  let rollback ~xt (snap, post_commit) =
+    xt.cass <- rollback xt.casn snap xt.cass;
+    xt.post_commit <- post_commit
 
   let rec first ~xt tx = function
     | [] -> tx ~xt
diff --git a/test/kcas/test.ml b/test/kcas/test.ml
index c8d47fe4..99ead17f 100644
--- a/test/kcas/test.ml
+++ b/test/kcas/test.ml
@@ -524,21 +524,29 @@ let test_rollback () =
     let n_permanent = Random.int n_locs in
     let n_rollbacks = Random.int n_locs in
 
+    let expected = ref false in
+    let unexpected = ref false in
+
     let tx ~xt =
       in_place_shuffle locs;
       for i = 0 to n_permanent - 1 do
         Xt.incr ~xt locs.(i)
       done;
+      Xt.post_commit ~xt (fun () -> expected := true);
 
       let snap = Xt.snapshot ~xt in
       in_place_shuffle locs;
       for i = 0 to n_rollbacks - 1 do
         Xt.incr ~xt locs.(i)
       done;
+      Xt.post_commit ~xt (fun () -> unexpected := true);
       Xt.rollback ~xt snap
     in
     Xt.commit { tx };
 
+    assert !expected;
+    assert (not !unexpected);
+
     accum := n_permanent + !accum
   done;