@@ -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 )
@@ -46,19 +51,22 @@ import GHC.Weak (Weak(Weak))
4651-- We use a TCHan rather than a Chan so that we have a non-blocking read
4752data CQueue a = CQueue (StrictMVar (StrictList a)) -- Arrived
4853 (TChan a) -- Incoming
54+ (TVar Int) -- Queue size
4955
5056newCQueue :: IO (CQueue a)
51−newCQueue = CQueue <$> newMVar Nil <*> atomically newTChan
57+newCQueue = CQueue <$> newMVar Nil <*> atomically newTChan <*> newTVarIO 0
5258
5359-- | Enqueue an element
5460--
5561-- Enqueue is strict.
5662enqueue :: CQueue a -> a -> IO ()
57−enqueue (CQueue _arrived incoming) !a = atomically $ writeTChan incoming a
63+enqueue c !a = atomically (enqueueSTM c a)
5864
5965-- | Variant of enqueue for use in the STM monad.
6066enqueueSTM :: CQueue a -> a -> STM ()
61−enqueueSTM (CQueue _arrived incoming) !a = writeTChan incoming a
67+enqueueSTM (CQueue _arrived incoming size) !a = do
68+ writeTChan incoming a
69+ modifyTVar' size succ
6270
6371data BlockSpec =
6472 NonBlocking
@@ -101,7 +109,7 @@ dequeue :: forall m a.
101109 -> BlockSpec -- ^ Blocking behaviour
102110 -> [MatchOn m a] -- ^ List of matches
103111 -> IO (Maybe a) -- ^ 'Nothing' only on timeout
104−dequeue (CQueue arrived incoming) blockSpec matchons = mask_ $
112+dequeue (CQueue arrived incoming size) blockSpec matchons = mask_ $ decrementJust $
105113 case blockSpec of
106114 Timeout n -> timeout n $ fmap fromJust run
107115 _other ->
@@ -113,6 +121,16 @@ dequeue (CQueue arrived incoming) blockSpec matchons = mask_ $
113121 -- no onException needed
114122 _other -> run
115123 where
124+ -- Decrement counter is smth is returned from the queue,
125+ -- this is safe to use as method is called under a mask
126+ -- and there is no 'unmasked' operation inside
127+ decrementJust f = do
128+ mx <- f
129+ case mx of
130+ Just{} -> atomically $ modifyTVar' size pred
131+ Nothing -> return ()
132+ return mx
133+
116134 chunks = chunkMatches matchons
117135
118136 run = do
@@ -245,5 +263,8 @@ dequeue (CQueue arrived incoming) blockSpec matchons = mask_ $
245263
246264-- | Weak reference to a CQueue
247265mkWeakCQueue :: CQueue a -> IO () -> IO (Weak (CQueue a))
248−mkWeakCQueue m@(CQueue (StrictMVar (MVar m#)) _) f = IO $ \s ->
266+mkWeakCQueue m@(CQueue (StrictMVar (MVar m#)) _ _) f = IO $ \s ->
249267 case mkWeak# m# m f s of (# s1, w #) -> (# s1, Weak w #)
268+
269+queueSize :: CQueue a -> IO Int
270+queueSize (CQueue _ _ size) = readTVarIO size