Make takeMVar exception safe

input-output-hk · Oct 21, 2024 · 64d32ad · 64d32ad
1 parent b14daa3
commit 64d32ad
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 1 deletion.
diff --git 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
@@ -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