diff --git a/Control/Concurrent/STM/TBQueue.hs b/Control/Concurrent/STM/TBQueue.hs index e38a2a2..c2948ea 100644 --- a/Control/Concurrent/STM/TBQueue.hs +++ b/Control/Concurrent/STM/TBQueue.hs @@ -35,6 +35,7 @@ module Control.Concurrent.STM.TBQueue ( newTBQueueIO, readTBQueue, tryReadTBQueue, + snapshotTBQueue, flushTBQueue, peekTBQueue, tryPeekTBQueue, @@ -146,6 +147,14 @@ readTBQueue (TBQueue rsize read _wsize write _size) = do tryReadTBQueue :: TBQueue a -> STM (Maybe a) tryReadTBQueue c = fmap Just (readTBQueue c) `orElse` return Nothing +-- | Efficiently read the entire contents of a 'TBQueue' into a list without changing queue contents. +-- This function never retries. +snapshotTBQueue :: TBQueue a -> STM [a] +snapshotTBQueue (TBQueue _ read _ write _) = do + xs <- readTVar read + ys <- readTVar write + return (xs ++ reverse ys) + -- | Efficiently read the entire contents of a 'TBQueue' into a list. This -- function never retries. -- diff --git a/Control/Concurrent/STM/TQueue.hs b/Control/Concurrent/STM/TQueue.hs index 720cfa7..679a706 100644 --- a/Control/Concurrent/STM/TQueue.hs +++ b/Control/Concurrent/STM/TQueue.hs @@ -39,6 +39,7 @@ module Control.Concurrent.STM.TQueue ( newTQueueIO, readTQueue, tryReadTQueue, + snapshotTQueue, flushTQueue, peekTQueue, tryPeekTQueue, @@ -108,6 +109,14 @@ readTQueue (TQueue read write) = do tryReadTQueue :: TQueue a -> STM (Maybe a) tryReadTQueue c = fmap Just (readTQueue c) `orElse` return Nothing +-- | Efficiently read the entire contents of a 'TQueue' into a list without changing queue contents. +-- This function never retries. +snapshotTQueue :: TQueue a -> STM [a] +snapshotTQueue (TQueue read write) = do + xs <- readTVar read + ys <- readTVar write + return (xs ++ reverse ys) + -- | Efficiently read the entire contents of a 'TQueue' into a list. This -- function never retries. -- diff --git a/testsuite/src/Issue17.hs b/testsuite/src/Issue17.hs index 06b72f0..f04f9b7 100644 --- a/testsuite/src/Issue17.hs +++ b/testsuite/src/Issue17.hs @@ -63,6 +63,9 @@ assertEmptyTBQueue queue = do atomically (tryPeekTBQueue queue) >>= assertEqual "Expected empty: tryPeekTBQueue should return Nothing" Nothing + atomically (snapshotTBQueue queue) >>= + assertEqual "Expected empty: snapshotTBQueue should return []" [] + atomically (flushTBQueue queue) >>= assertEqual "Expected empty: flushTBQueue should return []" [] diff --git a/testsuite/src/Issue9.hs b/testsuite/src/Issue9.hs index 88c0036..ef92113 100644 --- a/testsuite/src/Issue9.hs +++ b/testsuite/src/Issue9.hs @@ -24,6 +24,9 @@ main = do -- Read 1 1 <- atomically (readTBQueue queue) + -- Snapshot [2..5] + [2,3,4,5] <- atomically (snapshotTBQueue queue) + -- Flush [2..5] [2,3,4,5] <- atomically (flushTBQueue queue)