git.iko.soy/distributed-process
on master help sign in
DP-02: simple approach for maintaining queue size.
5096410Alexander Vershilov
src/Control/Distributed/Process/Internal/CQueue.hs+22−5
compare with master
unifiedsplit
@@ -10,6 +10,7 @@ module Control.Distributed.Process.Internal.CQueue
10 , enqueueSTM
10 , enqueueSTM
11 , dequeue
11 , dequeue
12 , mkWeakCQueue
12 , mkWeakCQueue
13+ , queueSize
13 ) where
14 ) where
14
15
15import Prelude hiding (length, reverse)
16import Prelude hiding (length, reverse)
@@ -17,10 +18,14 @@ import Control.Concurrent.STM
17 ( atomically
18 ( atomically
18 , STM
19 , STM
19 , TChan
20 , TChan
21+ , TVar
22+ , modifyTVar'
20 , tryReadTChan
23 , tryReadTChan
21 , newTChan
24 , newTChan
25+ , newTVarIO
22 , writeTChan
26 , writeTChan
23 , readTChan
27 , readTChan
28+ , readTVarIO
24 , orElse
29 , orElse
25 , retry
30 , retry
26 )
31 )
@@ -38,6 +43,7 @@ import Control.Distributed.Process.Internal.StrictList
38 , append
43 , append
39 )
44 )
40import Data.Maybe (fromJust)
45import Data.Maybe (fromJust)
46+import Data.Foldable (traverse_)
41import GHC.MVar (MVar(MVar))
47import GHC.MVar (MVar(MVar))
42import GHC.IO (IO(IO))
48import GHC.IO (IO(IO))
43import GHC.Prim (mkWeak#)
49import GHC.Prim (mkWeak#)
@@ -46,19 +52,22 @@ import GHC.Weak (Weak(Weak))
46-- We use a TCHan rather than a Chan so that we have a non-blocking read
52-- We use a TCHan rather than a Chan so that we have a non-blocking read
47data CQueue a = CQueue (StrictMVar (StrictList a)) -- Arrived
53data CQueue a = CQueue (StrictMVar (StrictList a)) -- Arrived
48 (TChan a) -- Incoming
54 (TChan a) -- Incoming
55+ (TVar Int) -- Queue size
49
56
50newCQueue :: IO (CQueue a)
57newCQueue :: IO (CQueue a)
51−newCQueue = CQueue <$> newMVar Nil <*> atomically newTChan
58+newCQueue = CQueue <$> newMVar Nil <*> atomically newTChan <*> newTVarIO 0
52
59
53-- | Enqueue an element
60-- | Enqueue an element
54--
61--
55-- Enqueue is strict.
62-- Enqueue is strict.
56enqueue :: CQueue a -> a -> IO ()
63enqueue :: CQueue a -> a -> IO ()
57−enqueue (CQueue _arrived incoming) !a = atomically $ writeTChan incoming a
64+enqueue c !a = atomically (enqueueSTM c a)
58
65
59-- | Variant of enqueue for use in the STM monad.
66-- | Variant of enqueue for use in the STM monad.
60enqueueSTM :: CQueue a -> a -> STM ()
67enqueueSTM :: 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
62
71
63data BlockSpec =
72data BlockSpec =
64 NonBlocking
73 NonBlocking
@@ -101,7 +110,7 @@ dequeue :: forall m a.
101 -> BlockSpec -- ^ Blocking behaviour
110 -> BlockSpec -- ^ Blocking behaviour
102 -> [MatchOn m a] -- ^ List of matches
111 -> [MatchOn m a] -- ^ List of matches
103 -> IO (Maybe a) -- ^ 'Nothing' only on timeout
112 -> 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 $
105 case blockSpec of
114 case blockSpec of
106 Timeout n -> timeout n $ fmap fromJust run
115 Timeout n -> timeout n $ fmap fromJust run
107 _other ->
116 _other ->
@@ -113,6 +122,11 @@ dequeue (CQueue arrived incoming) blockSpec matchons = mask_ $
113 -- no onException needed
122 -- no onException needed
114 _other -> run
123 _other -> run
115 where
124 where
125+ checkDecrement f = do
126+ mx <- f
127+ mask_ $ traverse_ (const $ atomically $ modifyTVar' size pred) mx
128+ return mx
129+
116 chunks = chunkMatches matchons
130 chunks = chunkMatches matchons
117
131
118 run = do
132 run = do
@@ -245,5 +259,8 @@ dequeue (CQueue arrived incoming) blockSpec matchons = mask_ $
245
259
246-- | Weak reference to a CQueue
260-- | Weak reference to a CQueue
247mkWeakCQueue :: CQueue a -> IO () -> IO (Weak (CQueue a))
261mkWeakCQueue :: 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 ->
249 case mkWeak# m# m f s of (# s1, w #) -> (# s1, Weak w #)
263 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