Provide a exception-safe withTMVar combinator

io-classes/CHANGELOG.md

@@ -6,6 +6,7 @@
 
 * Added `threadLabel` to `MonadThread`
 * Added `MonadLabelledMVar` class.
+* Added `withTMVar` and `withTMVarAnd` functions.
 
 ### 1.7.0.0
 

io-classes/src/Control/Concurrent/Class/MonadSTM/TMVar.hs

@@ -18,6 +18,8 @@ module Control.Concurrent.Class.MonadSTM.TMVar
   , swapTMVar
   , writeTMVar
   , isEmptyTMVar
+  , withTMVar
+  , withTMVarAnd
     -- * MonadLabelledSTM
   , labelTMVar
   , labelTMVarIO

io-classes/src/Control/Monad/Class/MonadSTM/Internal.hs

@@ -6,6 +6,7 @@
 {-# LANGUAGE FlexibleInstances          #-}
 {-# LANGUAGE GADTs                      #-}
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE LambdaCase                 #-}
 {-# LANGUAGE MultiParamTypeClasses      #-}
 {-# LANGUAGE NamedFieldPuns             #-}
 {-# LANGUAGE PatternSynonyms            #-}
@@ -99,6 +100,9 @@ module Control.Monad.Class.MonadSTM.Internal
   , isEmptyTChanDefault
   , cloneTChanDefault
   , labelTChanDefault
+  -- * WithTMVar
+  , withTMVar
+  , withTMVarAnd
   ) where
 
 import Prelude hiding (read)
@@ -1257,3 +1261,40 @@ instance MonadSTM m => MonadSTM (ReaderT r m) where
 writeTMVar' :: STM.TMVar a -> a -> STM.STM ()
 writeTMVar' t new = STM.tryTakeTMVar t >> STM.putTMVar t new
 #endif
+
+
+-- | Apply @f@ with the content of @tv@ as state, restoring the original value when an
+-- exception occurs
+withTMVar ::
+     (MonadSTM m, MonadThrow.MonadCatch m)
+  => TMVar m a
+  -> (a -> m (c, a))
+  -> m c
+withTMVar tv f = withTMVarAnd tv (const $ pure ()) (\a -> const $ f a)
+
+-- | Apply @f@ with the content of @tv@ as state, restoring the original value
+-- when an exception occurs. Additionally run a @STM@ action when acquiring the
+-- value.
+withTMVarAnd ::
+     (MonadSTM m, MonadThrow.MonadCatch m)
+  => TMVar m a
+  -> (a -> STM m b) -- ^ Additional STM action to run in the same atomically
+                    -- block as the TMVar is acquired
+  -> (a -> b -> m (c, a)) -- ^ Action
+  -> m c
+withTMVarAnd tv guard f =
+  fst . fst <$> MonadThrow.generalBracket
+    (atomically $ do
+        istate <- takeTMVar tv
+        guarded <- guard istate
+        pure (istate, guarded)
+    )
+    (\(origState, _) -> \case
+        MonadThrow.ExitCaseSuccess (_, newState)
+          -> atomically $ putTMVar tv newState
+        MonadThrow.ExitCaseException _
+          -> atomically $ putTMVar tv origState
+        MonadThrow.ExitCaseAbort
+          -> atomically $ putTMVar tv origState
+    )
+    (uncurry f)

io-classes/strict-stm/Control/Concurrent/Class/MonadSTM/Strict/TMVar.hs

@@ -1,6 +1,7 @@
 {-# LANGUAGE BangPatterns       #-}
 {-# LANGUAGE ExplicitNamespaces #-}
 {-# LANGUAGE GADTs              #-}
+{-# LANGUAGE LambdaCase         #-}
 {-# LANGUAGE TypeOperators      #-}
 
 -- | This module corresponds to `Control.Concurrent.STM.TMVar` in "stm" package
@@ -25,18 +26,22 @@ module Control.Concurrent.Class.MonadSTM.Strict.TMVar
   , swapTMVar
   , writeTMVar
   , isEmptyTMVar
+  , withTMVar
+  , withTMVarAnd
     -- * MonadLabelledSTM
   , labelTMVar
   , labelTMVarIO
     -- * MonadTraceSTM
   , traceTMVar
   , traceTMVarIO
+  , traceTMVarShow
+  , traceTMVarShowIO
   ) where
 
 
 import Control.Concurrent.Class.MonadSTM.TMVar qualified as Lazy
 import Control.Monad.Class.MonadSTM hiding (traceTMVar, traceTMVarIO)
-
+import Control.Monad.Class.MonadThrow
 
 type LazyTMVar   m = Lazy.TMVar m
 
@@ -59,12 +64,39 @@ traceTMVar :: MonadTraceSTM m
            -> STM m ()
 traceTMVar p (StrictTMVar var) = Lazy.traceTMVar p var
 
+traceTMVarShow :: (MonadTraceSTM m, Show a)
+               => proxy m
+               -> StrictTMVar m a
+               -> STM m ()
+traceTMVarShow p tmvar =
+  traceTMVar p tmvar (\pv v -> pure $ TraceString $ case (pv, v) of
+          (Nothing, Nothing) -> "Created empty"
+          (Nothing, Just st') -> "Created full: " <> show st'
+          (Just Nothing, Just st') -> "Put: " <> show st'
+          (Just Nothing, Nothing) -> "Remains empty"
+          (Just Just{}, Nothing) -> "Take"
+          (Just (Just st'), Just st'') -> "Modified: " <> show st' <> " -> " <> show st''
+      )
+
 traceTMVarIO :: MonadTraceSTM m
              => StrictTMVar m a
              -> (Maybe (Maybe a) -> (Maybe a) -> InspectMonad m TraceValue)
              -> m ()
 traceTMVarIO (StrictTMVar var) = Lazy.traceTMVarIO var
 
+traceTMVarShowIO :: (Show a, MonadTraceSTM m)
+                 => StrictTMVar m a
+                 -> m ()
+traceTMVarShowIO tmvar =
+  traceTMVarIO tmvar (\pv v -> pure $ TraceString $ case (pv, v) of
+          (Nothing, Nothing) -> "Created empty"
+          (Nothing, Just st') -> "Created full: " <> show st'
+          (Just Nothing, Just st') -> "Put: " <> show st'
+          (Just Nothing, Nothing) -> "Remains empty"
+          (Just Just{}, Nothing) -> "Take"
+          (Just (Just st'), Just st'') -> "Modified: " <> show st' <> " -> " <> show st''
+      )
+
 castStrictTMVar :: LazyTMVar m ~ LazyTMVar n
                 => StrictTMVar m a -> StrictTMVar n a
 castStrictTMVar (StrictTMVar var) = StrictTMVar var
@@ -107,3 +139,24 @@ writeTMVar (StrictTMVar tmvar) !a = Lazy.writeTMVar tmvar a
 
 isEmptyTMVar :: MonadSTM m => StrictTMVar m a -> STM m Bool
 isEmptyTMVar (StrictTMVar tmvar) = Lazy.isEmptyTMVar tmvar
+
+withTMVar :: (MonadSTM m, MonadCatch m)
+          => StrictTMVar m a
+          -> (a -> m (c, a))
+          -> m c
+withTMVar (StrictTMVar tmvar) f =
+  Lazy.withTMVar tmvar (\x -> do
+                           !(!c, !a) <- f x
+                           pure $! (c, a)
+                       )
+
+withTMVarAnd :: (MonadSTM m, MonadCatch m)
+             => StrictTMVar m a
+             -> (a -> STM m b)
+             -> (a -> b -> m (c, a))
+             -> m c
+withTMVarAnd (StrictTMVar tmvar) f g =
+  Lazy.withTMVarAnd tmvar f (\x y -> do
+                                !(!c, !a) <- g x y
+                                pure $! (c, a)
+                            )