From b14daa3ae3d245c360373319a8e09978ec34742f Mon Sep 17 00:00:00 2001
From: Javier Sagredo <javier.sagredo@iohk.io>
Date: Tue, 8 Oct 2024 15:54:52 +0200
Subject: [PATCH 1/2] Add regression test

---
 .../Control/Concurrent/Class/MonadMVar.hs     | 25 +++++++++++++++++++
 1 file changed, 25 insertions(+)

diff --git a/io-sim/test/Test/Control/Concurrent/Class/MonadMVar.hs b/io-sim/test/Test/Control/Concurrent/Class/MonadMVar.hs
index aa72a787..a9049b97 100644
--- a/io-sim/test/Test/Control/Concurrent/Class/MonadMVar.hs
+++ b/io-sim/test/Test/Control/Concurrent/Class/MonadMVar.hs
@@ -8,6 +8,7 @@ module Test.Control.Concurrent.Class.MonadMVar where
 import Control.Concurrent.Class.MonadMVar
 import Control.Monad.Class.MonadAsync
 import Control.Monad.Class.MonadFork
+import Control.Monad.Class.MonadTest
 import Control.Monad.Class.MonadTime.SI
 import Control.Monad.Class.MonadTimer.SI
 import Data.Bifoldable (bifoldMap)
@@ -64,6 +65,7 @@ tests =
       [ testCase "empty MVar is empty"    unit_isEmptyMVar_empty_sim
       , testCase "full MVar is not empty" unit_isEmptyMVar_full_sim
       ]
+    , testProperty "takeMVar is exception safe" prop_takeMVar_exception_safe
     ]
 
 
@@ -310,6 +312,29 @@ unit_isEmptyMVar_full_sim =
     assertBool "full mvar must not be empty" $
     runSimOrThrow (prop_isEmptyMVar False)
 
+--
+-- takeMVar is exception safe
+--
+prop_takeMVar_exception_safe :: Property
+prop_takeMVar_exception_safe =
+  exploreSimTrace id (do
+      exploreRaces
+      mv <- newMVar (0 :: Int)
+      t1 <- async $ void $ withMVar mv (\v -> pure (v + 1, ()))
+      t2 <- async $ void $ do
+        _ <- withMVar mv (\v -> pure (v + 1, ()))
+        withMVar mv (\v -> pure (v + 1, ()))
+      t3 <- async $ cancel t1
+      wait t3
+      wait t2
+      wait t1
+  ) (\_ trace ->
+       case traceResult False trace of
+         Left FailureDeadlock{} ->
+           counterexample (ppTrace trace) $ property False
+         _ -> property True
+    )
+
 --
 -- Utils
 --

From 64d32adcf54d924158c8d7c832f9c0d58fc659fb Mon Sep 17 00:00:00 2001
From: Javier Sagredo <javier.sagredo@iohk.io>
Date: Tue, 8 Oct 2024 15:55:12 +0200
Subject: [PATCH 2/2] Make takeMVar exception safe

---
 io-sim/CHANGELOG.md                   |  2 ++
 io-sim/src/Control/Monad/IOSim/STM.hs | 19 ++++++++++++++++++-
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/io-sim/CHANGELOG.md b/io-sim/CHANGELOG.md
index 7260b8f7..a4e9a80a 100644
--- a/io-sim/CHANGELOG.md
+++ b/io-sim/CHANGELOG.md
@@ -10,6 +10,8 @@
 - Implement `MonadLabelledMVar` instance for `(IOSim s)`
 - `TVarId` is now a sum type with one constructor per `TVar` role, e.g. `TVar`,
   `TMVar`, `MVar` and a few others - except for `TChan`.
+- A blocked `takeTVar` is now safe in the presence of exceptions. It will relay
+  the value to other waiting threads.
 
 ## 1.6.0.0
 
diff --git a/io-sim/src/Control/Monad/IOSim/STM.hs b/io-sim/src/Control/Monad/IOSim/STM.hs
index 672926ed..15c155a5 100644
--- a/io-sim/src/Control/Monad/IOSim/STM.hs
+++ b/io-sim/src/Control/Monad/IOSim/STM.hs
@@ -398,7 +398,24 @@ takeMVarDefault (MVar tv) = mask_ $ do
                 -- takevar; we need to remove it from 'takeq', otherwise we
                 -- will have a space leak.
                 let takeq' = Deque.filter (/= takevar) takeq
-                writeTVar tv (MVarEmpty takeq' readq)
+                takevalue <- readTVar takevar
+                case takevalue of
+                  Nothing ->
+                    writeTVar tv (MVarEmpty takeq' readq)
+                  -- we were given a value before we could read it. Relay it to any
+                  -- new reading threads and possible the next take thread.
+                  Just x -> do
+                    -- notify readers
+                    mapM_ (\readvar -> writeTVar readvar (Just x)) readq
+
+                    -- notify first `takeMVar` thread
+                    case Deque.uncons takeq' of
+                      Nothing ->
+                        writeTVar tv (MVarFull x mempty)
+
+                      Just (takevar', takeq'') -> do
+                        writeTVar takevar' (Just x)
+                        writeTVar tv (MVarEmpty takeq'' mempty)
 
               -- This case is unlikely but possible if another thread ran
               -- first and modified the mvar. This situation is fine as far as