From 3abe5ebea05eeb8bf56467b630b41909fe4ae0ff Mon Sep 17 00:00:00 2001
From: Javier Sagredo <javier.sagredo@iohk.io>
Date: Thu, 26 Sep 2024 10:23:56 +0200
Subject: [PATCH] Modify `withTMVar` to use `generalBracket`

---
 .../Ouroboros/Consensus/Util.hs               | 23 +++++++++++--------
 1 file changed, 13 insertions(+), 10 deletions(-)

diff --git a/ouroboros-consensus/src/ouroboros-consensus/Ouroboros/Consensus/Util.hs b/ouroboros-consensus/src/ouroboros-consensus/Ouroboros/Consensus/Util.hs
index 0b5c88af0c..cb17c656df 100644
--- a/ouroboros-consensus/src/ouroboros-consensus/Ouroboros/Consensus/Util.hs
+++ b/ouroboros-consensus/src/ouroboros-consensus/Ouroboros/Consensus/Util.hs
@@ -487,19 +487,22 @@ withTMVarAnd ::
      IOLike m
   => StrictTMVar m a
   -> (a -> STM m b) -- ^ Additional STM action to run in the same atomically
-                     -- block as the TMVar is acquired
+                    -- block as the TMVar is acquired
   -> (a -> b -> m (c, a)) -- ^ Action
   -> m c
 withTMVarAnd tv guard f =
-  bracketOnError
+  fst . fst <$> generalBracket
     (atomically $ do
-        i <- takeTMVar tv
-        g <- guard i
-        pure (i, g)
+        istate <- takeTMVar tv
+        guarded <- guard istate
+        pure (istate, guarded)
     )
-    (atomically . putTMVar tv . fst)
-    (\(s, g) -> do
-        (x, s') <- f s g
-        atomically $ putTMVar tv s'
-        return x
+    (\(origState, _) -> \case
+        ExitCaseSuccess (_, newState)
+          -> atomically $ putTMVar tv newState
+        ExitCaseException _
+          -> atomically $ putTMVar tv origState
+        ExitCaseAbort
+          -> atomically $ putTMVar tv origState
     )
+    (uncurry f)