@@ -10,6 +10,7 @@ module Control.Distributed.Process.Internal.CQueue
1010 , enqueueSTM
1111 , dequeue
1212 , mkWeakCQueue
13+ , queueSize
1314 ) where
1415
1516import Prelude hiding (length, reverse)
@@ -17,10 +18,14 @@ import Control.Concurrent.STM
1718 ( atomically
1819 , STM
1920 , TChan
21+ , TVar
22+ , modifyTVar'
2023 , tryReadTChan
2124 , newTChan
25+ , newTVarIO
2226 , writeTChan
2327 , readTChan
28+ , readTVarIO
2429 , orElse
2530 , retry
2631 )
@@ -38,6 +43,7 @@ import Control.Distributed.Process.Internal.StrictList
3843 , append
3944 )
4045import Data.Maybe (fromJust)
46+import Data.Foldable (traverse_)
4147import GHC.MVar (MVar(MVar))
4248import GHC.IO (IO(IO))
4349import GHC.Prim (mkWeak#)
@@ -46,19 +52,22 @@ import GHC.Weak (Weak(Weak))
4652-- We use a TCHan rather than a Chan so that we have a non-blocking read
4753data CQueue a = CQueue (StrictMVar (StrictList a)) -- Arrived
4854 (TChan a) -- Incoming
55+ (TVar Int) -- Queue size
4956
5057newCQueue :: IO (CQueue a)
51−newCQueue = CQueue <$> newMVar Nil <*> atomically newTChan
58+newCQueue = CQueue <$> newMVar Nil <*> atomically newTChan <*> newTVarIO 0
5259
5360-- | Enqueue an element
5461--
5562-- Enqueue is strict.
5663enqueue :: CQueue a -> a -> IO ()
57−enqueue (CQueue _arrived incoming) !a = atomically $ writeTChan incoming a
64+enqueue c !a = atomically (enqueueSTM c a)
5865
5966-- | Variant of enqueue for use in the STM monad.
6067enqueueSTM :: CQueue a -> a -> STM ()
61−enqueueSTM (CQueue _arrived incoming) !a = writeTChan incoming a
68+enqueueSTM (CQueue _arrived incoming size) !a = do
69+ writeTChan incoming a
70+ modifyTVar' size succ
6271
6372data BlockSpec =
6473 NonBlocking
@@ -101,7 +110,7 @@ dequeue :: forall m a.
101110 -> BlockSpec -- ^ Blocking behaviour
102111 -> [MatchOn m a] -- ^ List of matches
103112 -> IO (Maybe a) -- ^ 'Nothing' only on timeout
104−dequeue (CQueue arrived incoming) blockSpec matchons = mask_ $
113+dequeue (CQueue arrived incoming size) blockSpec matchons = mask_ $ checkDecrement $
105114 case blockSpec of
106115 Timeout n -> timeout n $ fmap fromJust run
107116 _other ->
@@ -113,6 +122,11 @@ dequeue (CQueue arrived incoming) blockSpec matchons = mask_ $
113122 -- no onException needed
114123 _other -> run
115124 where
125+ checkDecrement f = do
126+ mx <- f
127+ mask_ $ traverse_ (const $ atomically $ modifyTVar' size pred) mx
128+ return mx
129+
116130 chunks = chunkMatches matchons
117131
118132 run = do
@@ -245,5 +259,8 @@ dequeue (CQueue arrived incoming) blockSpec matchons = mask_ $
245259
246260-- | Weak reference to a CQueue
247261mkWeakCQueue :: CQueue a -> IO () -> IO (Weak (CQueue a))
248−mkWeakCQueue m@(CQueue (StrictMVar (MVar m#)) _) f = IO $ \s ->
262+mkWeakCQueue m@(CQueue (StrictMVar (MVar m#)) _ _) f = IO $ \s ->
249263 case mkWeak# m# m f s of (# s1, w #) -> (# s1, Weak w #)
264+
265+queueSize :: CQueue a -> IO Int
266+queueSize (CQueue _ _ size) = readTVarIO size