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