git.iko.soy/distributed-process
on master help sign in
Merge pull request #176 from qnikst/DP-02-IO
4d3472aTim Watson
src/Control/Distributed/Process/Internal/CQueue.hs+26−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 )
@@ -46,19 +51,22 @@ import GHC.Weak (Weak(Weak))
46-- We use a TCHan rather than a Chan so that we have a non-blocking read
51-- We use a TCHan rather than a Chan so that we have a non-blocking read
47data CQueue a = CQueue (StrictMVar (StrictList a)) -- Arrived
52data CQueue a = CQueue (StrictMVar (StrictList a)) -- Arrived
48 (TChan a) -- Incoming
53 (TChan a) -- Incoming
54+ (TVar Int) -- Queue size
49
55
50newCQueue :: IO (CQueue a)
56newCQueue :: IO (CQueue a)
51−newCQueue = CQueue <$> newMVar Nil <*> atomically newTChan
57+newCQueue = CQueue <$> newMVar Nil <*> atomically newTChan <*> newTVarIO 0
52
58
53-- | Enqueue an element
59-- | Enqueue an element
54--
60--
55-- Enqueue is strict.
61-- Enqueue is strict.
56enqueue :: CQueue a -> a -> IO ()
62enqueue :: CQueue a -> a -> IO ()
57−enqueue (CQueue _arrived incoming) !a = atomically $ writeTChan incoming a
63+enqueue c !a = atomically (enqueueSTM c a)
58
64
59-- | Variant of enqueue for use in the STM monad.
65-- | Variant of enqueue for use in the STM monad.
60enqueueSTM :: CQueue a -> a -> STM ()
66enqueueSTM :: 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
62
70
63data BlockSpec =
71data BlockSpec =
64 NonBlocking
72 NonBlocking
@@ -101,7 +109,7 @@ dequeue :: forall m a.
101 -> BlockSpec -- ^ Blocking behaviour
109 -> BlockSpec -- ^ Blocking behaviour
102 -> [MatchOn m a] -- ^ List of matches
110 -> [MatchOn m a] -- ^ List of matches
103 -> IO (Maybe a) -- ^ 'Nothing' only on timeout
111 -> 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 $
105 case blockSpec of
113 case blockSpec of
106 Timeout n -> timeout n $ fmap fromJust run
114 Timeout n -> timeout n $ fmap fromJust run
107 _other ->
115 _other ->
@@ -113,6 +121,16 @@ dequeue (CQueue arrived incoming) blockSpec matchons = mask_ $
113 -- no onException needed
121 -- no onException needed
114 _other -> run
122 _other -> run
115 where
123 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+
116 chunks = chunkMatches matchons
134 chunks = chunkMatches matchons
117
135
118 run = do
136 run = do
@@ -245,5 +263,8 @@ dequeue (CQueue arrived incoming) blockSpec matchons = mask_ $
245
263
246-- | Weak reference to a CQueue
264-- | Weak reference to a CQueue
247mkWeakCQueue :: CQueue a -> IO () -> IO (Weak (CQueue a))
265mkWeakCQueue :: 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 ->
249 case mkWeak# m# m f s of (# s1, w #) -> (# s1, Weak w #)
267 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