git.iko.soy/distributed-process
on 39ed082 help sign in
distributed-process/src/Control/Distributed/Process/Internal/CQueue.hs
269 lines 9.2 kB history raw
1{-# LANGUAGE BangPatterns #-}
2{-# LANGUAGE MagicHash, UnboxedTuples, PatternGuards, ScopedTypeVariables, RankNTypes #-}
3-- | Concurrent queue for single reader, single writer
4module Control.Distributed.Process.Internal.CQueue
5 ( CQueue
6 , BlockSpec(..)
7 , MatchOn(..)
8 , newCQueue
9 , enqueue
10 , enqueueSTM
11 , dequeue
12 , mkWeakCQueue
13 , queueSize
14 ) where
16import Prelude hiding (length, reverse)
17import Control.Concurrent.STM
18 ( atomically
19 , STM
20 , TChan
21 , TVar
22 , modifyTVar'
23 , tryReadTChan
24 , newTChan
25 , newTVarIO
26 , writeTChan
27 , readTChan
28 , readTVarIO
29 , orElse
30 , retry
31 )
32import Control.Applicative ((<$>), (<*>))
33import Control.Exception (mask_, onException)
34import System.Timeout (timeout)
35import Control.Distributed.Process.Internal.StrictMVar
36 ( StrictMVar(StrictMVar)
37 , newMVar
38 , takeMVar
39 , putMVar
40 )
41import Control.Distributed.Process.Internal.StrictList
42 ( StrictList(..)
43 , append
44 )
45import Data.Maybe (fromJust)
46import Data.Traversable (traverse)
47import GHC.MVar (MVar(MVar))
48import GHC.IO (IO(IO))
49import GHC.Prim (mkWeak#)
50import GHC.Weak (Weak(Weak))
52-- We use a TCHan rather than a Chan so that we have a non-blocking read
53data CQueue a = CQueue (StrictMVar (StrictList a)) -- Arrived
54 (TChan a) -- Incoming
55 (TVar Int) -- Queue size
57newCQueue :: IO (CQueue a)
58newCQueue = CQueue <$> newMVar Nil <*> atomically newTChan <*> newTVarIO 0
60-- | Enqueue an element
61--
62-- Enqueue is strict.
63enqueue :: CQueue a -> a -> IO ()
64enqueue c !a = atomically (enqueueSTM c a)
66-- | Variant of enqueue for use in the STM monad.
67enqueueSTM :: CQueue a -> a -> STM ()
68enqueueSTM (CQueue _arrived incoming size) !a = do
69 writeTChan incoming a
70 modifyTVar' size succ
72data BlockSpec =
73 NonBlocking
74 | Blocking
75 | Timeout Int
77data MatchOn m a
78 = MatchMsg (m -> Maybe a)
79 | MatchChan (STM a)
81type MatchChunks m a = [Either [m -> Maybe a] [STM a]]
83chunkMatches :: [MatchOn m a] -> MatchChunks m a
84chunkMatches [] = []
85chunkMatches (MatchMsg m : ms) = Left (m : chk) : chunkMatches rest
86 where (chk, rest) = spanMatchMsg ms
87chunkMatches (MatchChan r : ms) = Right (r : chk) : chunkMatches rest
88 where (chk, rest) = spanMatchChan ms
90spanMatchMsg :: [MatchOn m a] -> ([m -> Maybe a], [MatchOn m a])
91spanMatchMsg [] = ([],[])
92spanMatchMsg (m : ms)
93 | MatchMsg msg <- m = (msg:msgs, rest)
94 | otherwise = ([], m:ms)
95 where !(msgs,rest) = spanMatchMsg ms
97spanMatchChan :: [MatchOn m a] -> ([STM a], [MatchOn m a])
98spanMatchChan [] = ([],[])
99spanMatchChan (m : ms)
100 | MatchChan stm <- m = (stm:stms, rest)
101 | otherwise = ([], m:ms)
102 where !(stms,rest) = spanMatchChan ms
104-- | Dequeue an element
105--
106-- The timeout (if any) is applied only to waiting for incoming messages, not
107-- to checking messages that have already arrived
108dequeue :: forall m a.
109 CQueue m -- ^ Queue
110 -> BlockSpec -- ^ Blocking behaviour
111 -> [MatchOn m a] -- ^ List of matches
112 -> IO (Maybe a) -- ^ 'Nothing' only on timeout
113dequeue (CQueue arrived incoming size) blockSpec matchons = mask_ $ decrementJust $
114 case blockSpec of
115 Timeout n -> timeout n $ fmap fromJust run
116 _other ->
117 case chunks of
118 [Right ports] -> -- channels only, this is easy:
119 case blockSpec of
120 NonBlocking -> atomically $ waitChans ports (return Nothing)
121 _ -> atomically $ waitChans ports retry
122 -- no onException needed
123 _other -> run
124 where
125 -- Decrement counter is smth is returned from the queue,
126 -- this is safe to use as method is called under a mask
127 -- and there is no 'unmasked' operation inside
128 decrementJust f =
129 traverse (either return (\x -> decrement >> return x)) =<< f
130 decrement = atomically $ modifyTVar' size pred
132 chunks = chunkMatches matchons
134 run = do
135 arr <- takeMVar arrived
136 let grabNew xs = do
137 r <- atomically $ tryReadTChan incoming
138 case r of
139 Nothing -> return xs
140 Just x -> grabNew (Snoc xs x)
141 arr' <- grabNew arr
142 goCheck chunks arr'
144 waitChans ports on_block =
145 foldr orElse on_block (map (fmap (Just . Left)) ports)
147 --
148 -- First check the MatchChunks against the messages already in the
149 -- mailbox. For channel matches, we do a non-blocking check at
150 -- this point.
151 --
152 goCheck :: MatchChunks m a
153 -> StrictList m -- messages to check, in this order
154 -> IO (Maybe (Either a a))
156 goCheck [] old = goWait old
158 goCheck (Right ports : rest) old = do
159 r <- atomically $ waitChans ports (return Nothing) -- does not block
160 case r of
161 Just _ -> returnOld old r
162 Nothing -> goCheck rest old
164 goCheck (Left matches : rest) old = do
165 -- checkArrived might in principle take arbitrary time, so
166 -- we ought to call restore and use an exception handler. However,
167 -- the check is usually fast (just a comparison), and the overhead
168 -- of passing around restore and setting up exception handlers is
169 -- high. So just don't use expensive matchIfs!
170 case checkArrived matches old of
171 (old', Just r) -> returnOld old' (Just (Right r))
172 (old', Nothing) -> goCheck rest old'
173 -- use the result list, which is now left-biased
175 --
176 -- Construct an STM transaction that looks at the relevant channels
177 -- in the correct order.
178 --
179 mkSTM :: MatchChunks m a -> STM (Either m a)
180 mkSTM [] = retry
181 mkSTM (Left _ : rest)
182 = fmap Left (readTChan incoming) `orElse` mkSTM rest
183 mkSTM (Right ports : rest)
184 = foldr orElse (mkSTM rest) (map (fmap Right) ports)
186 waitIncoming :: IO (Maybe (Either m a))
187 waitIncoming = case blockSpec of
188 NonBlocking -> atomically $ fmap Just stm `orElse` return Nothing
189 _ -> atomically $ fmap Just stm
190 where
191 stm = mkSTM chunks
193 --
194 -- The initial pass didn't find a message, so now we go into blocking
195 -- mode.
196 --
197 -- Contents of 'arrived' from now on is (old ++ new), and
198 -- messages that arrive are snocced onto new.
199 --
200 goWait :: StrictList m -> IO (Maybe (Either a a))
201 goWait old = do
202 r <- waitIncoming `onException` putMVar arrived old
203 case r of
204 -- Nothing => non-blocking and no message
205 Nothing -> returnOld old Nothing
206 Just e -> case e of
207 --
208 -- Left => message arrived in the process mailbox. We now have to
209 -- run through the MatchChunks checking each one, because we might
210 -- have a situation where the first chunk fails to match and the
211 -- second chunk is a channel match and there *is* a message in the
212 -- channel. In that case the channel wins.
213 --
214 Left m -> goCheck1 chunks m old
215 --
216 -- Right => message arrived on a channel first
217 --
218 Right a -> returnOld old (Just (Left a))
220 --
221 -- A message arrived in the process inbox; check the MatchChunks for
222 -- a valid match.
223 --
224 goCheck1 :: MatchChunks m a
225 -> m -- single message to check
226 -> StrictList m -- old messages we have already checked
227 -> IO (Maybe (Either a a))
229 goCheck1 [] m old = goWait (Snoc old m)
231 goCheck1 (Right ports : rest) m old = do
232 r <- atomically $ waitChans ports (return Nothing) -- does not block
233 case r of
234 Nothing -> goCheck1 rest m old
235 Just _ -> returnOld (Snoc old m) r
237 goCheck1 (Left matches : rest) m old = do
238 case checkMatches matches m of
239 Nothing -> goCheck1 rest m old
240 Just p -> returnOld old (Just (Right p))
242 -- a common pattern for putting back the arrived queue at the end
243 returnOld :: StrictList m -> Maybe (Either a a) -> IO (Maybe (Either a a))
244 returnOld old r = do putMVar arrived old; return r
246 -- as a side-effect, this left-biases the list
247 checkArrived :: [m -> Maybe a] -> StrictList m -> (StrictList m, Maybe a)
248 checkArrived matches list = go list Nil
249 where
250 go Nil Nil = (Nil, Nothing)
251 go Nil r = go r Nil
252 go (Append xs ys) tl = go xs (append ys tl)
253 go (Snoc xs x) tl = go xs (Cons x tl)
254 go (Cons x xs) tl
255 | Just y <- checkMatches matches x = (append xs tl, Just y)
256 | otherwise = let !(rest,r) = go xs tl in (Cons x rest, r)
258 checkMatches :: [m -> Maybe a] -> m -> Maybe a
259 checkMatches [] _ = Nothing
260 checkMatches (m:ms) a = case m a of Nothing -> checkMatches ms a
261 Just b -> Just b
263-- | Weak reference to a CQueue
264mkWeakCQueue :: CQueue a -> IO () -> IO (Weak (CQueue a))
265mkWeakCQueue m@(CQueue (StrictMVar (MVar m#)) _ _) f = IO $ \s ->
266 case mkWeak# m# m f s of (# s1, w #) -> (# s1, Weak w #)
268queueSize :: CQueue a -> IO Int
269queueSize (CQueue _ _ size) = readTVarIO size