diff --git a/io-sim-classes/src/Control/Monad/Class/MonadSTM.hs b/io-sim-classes/src/Control/Monad/Class/MonadSTM.hs index 985bedc61bb..2ebfe078b50 100644 --- a/io-sim-classes/src/Control/Monad/Class/MonadSTM.hs +++ b/io-sim-classes/src/Control/Monad/Class/MonadSTM.hs @@ -44,6 +44,7 @@ module Control.Monad.Class.MonadSTM , isEmptyTBQueueDefault , isFullTBQueueDefault , lengthTBQueueDefault + , flushTBQueueDefault -- * MonadThrow aliases , throwSTM @@ -131,6 +132,7 @@ class ( Monad stm newTBQueue :: Natural -> stm (TBQueue_ stm a) readTBQueue :: TBQueue_ stm a -> stm a tryReadTBQueue :: TBQueue_ stm a -> stm (Maybe a) + flushTBQueue :: TBQueue_ stm a -> stm [a] writeTBQueue :: TBQueue_ stm a -> a -> stm () -- | @since 0.2.0.0 lengthTBQueue :: TBQueue_ stm a -> stm Natural @@ -214,6 +216,7 @@ instance MonadSTMTx STM.STM where newTQueue = STM.newTQueue readTQueue = STM.readTQueue tryReadTQueue = STM.tryReadTQueue + flushTBQueue = STM.flushTBQueue writeTQueue = STM.writeTQueue isEmptyTQueue = STM.isEmptyTQueue newTBQueue = STM.newTBQueue @@ -479,6 +482,20 @@ lengthTBQueueDefault (TBQueue rsize _read wsize _write size) = do return $! size - r - w +flushTBQueueDefault :: MonadSTM m => TBQueueDefault m a -> STM m [a] +flushTBQueueDefault (TBQueue rsize read wsize write size) = do + xs <- readTVar read + ys <- readTVar write + if null xs && null ys + then return [] + else do + writeTVar read [] + writeTVar write [] + writeTVar rsize 0 + writeTVar wsize size + return (xs ++ reverse ys) + + -- | 'throwIO' specialised to @stm@ monad. -- throwSTM :: (MonadSTMTx stm, MonadThrow.MonadThrow stm, Exception e) diff --git a/io-sim/src/Control/Monad/IOSim/Internal.hs b/io-sim/src/Control/Monad/IOSim/Internal.hs index 6ff202cb98e..496fecc57d9 100644 --- a/io-sim/src/Control/Monad/IOSim/Internal.hs +++ b/io-sim/src/Control/Monad/IOSim/Internal.hs @@ -340,6 +340,7 @@ instance MonadSTMTx (STM s) where newTBQueue = newTBQueueDefault readTBQueue = readTBQueueDefault tryReadTBQueue = tryReadTBQueueDefault + flushTBQueue = flushTBQueueDefault writeTBQueue = writeTBQueueDefault lengthTBQueue = lengthTBQueueDefault isEmptyTBQueue = isEmptyTBQueueDefault diff --git a/ouroboros-consensus/src/Ouroboros/Consensus/Util/EarlyExit.hs b/ouroboros-consensus/src/Ouroboros/Consensus/Util/EarlyExit.hs index 1ce8b95caac..0c9749e0028 100644 --- a/ouroboros-consensus/src/Ouroboros/Consensus/Util/EarlyExit.hs +++ b/ouroboros-consensus/src/Ouroboros/Consensus/Util/EarlyExit.hs @@ -110,6 +110,7 @@ instance MonadSTMTx stm => MonadSTMTx (WithEarlyExit stm) where newTBQueue = lift . newTBQueue readTBQueue = lift . readTBQueue tryReadTBQueue = lift . tryReadTBQueue + flushTBQueue = lift . flushTBQueue writeTBQueue = lift .: writeTBQueue lengthTBQueue = lift . lengthTBQueue isEmptyTBQueue = lift . isEmptyTBQueue