diff --git a/Control/Concurrent/STM/TBQueue.hs b/Control/Concurrent/STM/TBQueue.hs index 8722511..7b82388 100644 --- a/Control/Concurrent/STM/TBQueue.hs +++ b/Control/Concurrent/STM/TBQueue.hs @@ -161,10 +161,20 @@ flushTBQueue (TBQueue rsize read wsize write size) = do -- | Get the next value from the @TBQueue@ without removing it, -- retrying if the channel is empty. peekTBQueue :: TBQueue a -> STM a -peekTBQueue c = do - x <- readTBQueue c - unGetTBQueue c x - return x +peekTBQueue (TBQueue _ read _ write _) = do + xs <- readTVar read + case xs of + (x:_) -> return x + [] -> do + ys <- readTVar write + case ys of + [] -> retry + _ -> do + let (z:zs) = reverse ys -- NB. lazy: we want the transaction to be + -- short, otherwise it will conflict + writeTVar write [] + writeTVar read (z:zs) + return z -- | A version of 'peekTBQueue' which does not retry. Instead it -- returns @Nothing@ if no value is available. diff --git a/Control/Concurrent/STM/TQueue.hs b/Control/Concurrent/STM/TQueue.hs index 483db15..33052cf 100644 --- a/Control/Concurrent/STM/TQueue.hs +++ b/Control/Concurrent/STM/TQueue.hs @@ -122,10 +122,20 @@ flushTQueue (TQueue read write) = do -- | Get the next value from the @TQueue@ without removing it, -- retrying if the channel is empty. peekTQueue :: TQueue a -> STM a -peekTQueue c = do - x <- readTQueue c - unGetTQueue c x - return x +peekTQueue (TQueue read write) = do + xs <- readTVar read + case xs of + (x:_) -> return x + [] -> do + ys <- readTVar write + case ys of + [] -> retry + _ -> do + let (z:zs) = reverse ys -- NB. lazy: we want the transaction to be + -- short, otherwise it will conflict + writeTVar write [] + writeTVar read (z:zs) + return z -- | A version of 'peekTQueue' which does not retry. Instead it -- returns @Nothing@ if no value is available.