Make takeMVar exception safe

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
 

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

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
 --