io-sim: added flushTBQueue to MonadSTM

IntersectMBO · Sep 4, 2020 · c1ee15a · c1ee15a
1 parent bd3b148
commit c1ee15a
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 0 deletions.
diff --git 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
@@ -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
@@ -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