git.iko.soy/distributed-process
on 8aed555 help sign in
distributed-process/src/Control/Distributed/Process/Internal/CQueue.hs
276 lines 9.4 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
79data MatchOn m a
80 = MatchMsg (m -> Maybe a)
81 | MatchChan (STM a)
82 deriving (Functor)
84type MatchChunks m a = [Either [m -> Maybe a] [STM a]]
86chunkMatches :: [MatchOn m a] -> MatchChunks m a
87chunkMatches [] = []
88chunkMatches (MatchMsg m : ms) = Left (m : chk) : chunkMatches rest
89 where (chk, rest) = spanMatchMsg ms
90chunkMatches (MatchChan r : ms) = Right (r : chk) : chunkMatches rest
91 where (chk, rest) = spanMatchChan ms
93spanMatchMsg :: [MatchOn m a] -> ([m -> Maybe a], [MatchOn m a])
94spanMatchMsg [] = ([],[])
95spanMatchMsg (m : ms)
96 | MatchMsg msg <- m = (msg:msgs, rest)
97 | otherwise = ([], m:ms)
98 where !(msgs,rest) = spanMatchMsg ms
100spanMatchChan :: [MatchOn m a] -> ([STM a], [MatchOn m a])
101spanMatchChan [] = ([],[])
102spanMatchChan (m : ms)
103 | MatchChan stm <- m = (stm:stms, rest)
104 | otherwise = ([], m:ms)
105 where !(stms,rest) = spanMatchChan ms
107-- | Dequeue an element
108--
109-- The timeout (if any) is applied only to waiting for incoming messages, not
110-- to checking messages that have already arrived
111dequeue :: forall m a.
112 CQueue m -- ^ Queue
113 -> BlockSpec -- ^ Blocking behaviour
114 -> [MatchOn m a] -- ^ List of matches
115 -> IO (Maybe a) -- ^ 'Nothing' only on timeout
116dequeue (CQueue arrived incoming size) blockSpec matchons = mask_ $ decrementJust $
117 case blockSpec of
118 Timeout n -> timeout n $ fmap fromJust run
119 _other ->
120 case chunks of
121 [Right ports] -> -- channels only, this is easy:
122 case blockSpec of
123 NonBlocking -> atomically $ waitChans ports (return Nothing)
124 _ -> atomically $ waitChans ports retry
125 -- no onException needed
126 _other -> run
127 where
128 -- Decrement counter is smth is returned from the queue,
129 -- this is safe to use as method is called under a mask
130 -- and there is no 'unmasked' operation inside
131 decrementJust f =
132 traverse (either return (\x -> decrement >> return x)) =<< f
133 decrement = atomically $ modifyTVar' size pred
135 chunks = chunkMatches matchons
137 run = do
138 arr <- takeMVar arrived
139 let grabNew xs = do
140 r <- atomically $ tryReadTChan incoming
141 case r of
142 Nothing -> return xs
143 Just x -> grabNew (Snoc xs x)
144 arr' <- grabNew arr
145 goCheck chunks arr'
147 waitChans ports on_block =
148 foldr orElse on_block (map (fmap (Just . Left)) ports)
150 --
151 -- First check the MatchChunks against the messages already in the
152 -- mailbox. For channel matches, we do a non-blocking check at
153 -- this point.
154 --
155 goCheck :: MatchChunks m a
156 -> StrictList m -- messages to check, in this order
157 -> IO (Maybe (Either a a))
159 goCheck [] old = goWait old
161 goCheck (Right ports : rest) old = do
162 r <- atomically $ waitChans ports (return Nothing) -- does not block
163 case r of
164 Just _ -> returnOld old r
165 Nothing -> goCheck rest old
167 goCheck (Left matches : rest) old = do
168 -- checkArrived might in principle take arbitrary time, so
169 -- we ought to call restore and use an exception handler. However,
170 -- the check is usually fast (just a comparison), and the overhead
171 -- of passing around restore and setting up exception handlers is
172 -- high. So just don't use expensive matchIfs!
173 case checkArrived matches old of
174 (old', Just r) -> returnOld old' (Just (Right r))
175 (old', Nothing) -> goCheck rest old'
176 -- use the result list, which is now left-biased
178 --
179 -- Construct an STM transaction that looks at the relevant channels
180 -- in the correct order.
181 --
182 mkSTM :: MatchChunks m a -> STM (Either m a)
183 mkSTM [] = retry
184 mkSTM (Left _ : rest)
185 = fmap Left (readTChan incoming) `orElse` mkSTM rest
186 mkSTM (Right ports : rest)
187 = foldr orElse (mkSTM rest) (map (fmap Right) ports)
189 waitIncoming :: IO (Maybe (Either m a))
190 waitIncoming = case blockSpec of
191 NonBlocking -> atomically $ fmap Just stm `orElse` return Nothing
192 _ -> atomically $ fmap Just stm
193 where
194 stm = mkSTM chunks
196 --
197 -- The initial pass didn't find a message, so now we go into blocking
198 -- mode.
199 --
200 -- Contents of 'arrived' from now on is (old ++ new), and
201 -- messages that arrive are snocced onto new.
202 --
203 goWait :: StrictList m -> IO (Maybe (Either a a))
204 goWait old = do
205 r <- waitIncoming `onException` putMVar arrived old
206 case r of
207 -- Nothing => non-blocking and no message
208 Nothing -> returnOld old Nothing
209 Just e -> case e of
210 --
211 -- Left => message arrived in the process mailbox. We now have to
212 -- run through the MatchChunks checking each one, because we might
213 -- have a situation where the first chunk fails to match and the
214 -- second chunk is a channel match and there *is* a message in the
215 -- channel. In that case the channel wins.
216 --
217 Left m -> goCheck1 chunks m old
218 --
219 -- Right => message arrived on a channel first
220 --
221 Right a -> returnOld old (Just (Left a))
223 --
224 -- A message arrived in the process inbox; check the MatchChunks for
225 -- a valid match.
226 --
227 goCheck1 :: MatchChunks m a
228 -> m -- single message to check
229 -> StrictList m -- old messages we have already checked
230 -> IO (Maybe (Either a a))
232 goCheck1 [] m old = goWait (Snoc old m)
234 goCheck1 (Right ports : rest) m old = do
235 r <- atomically $ waitChans ports (return Nothing) -- does not block
236 case r of
237 Nothing -> goCheck1 rest m old
238 Just _ -> returnOld (Snoc old m) r
240 goCheck1 (Left matches : rest) m old = do
241 case checkMatches matches m of
242 Nothing -> goCheck1 rest m old
243 Just p -> returnOld old (Just (Right p))
245 -- a common pattern for putting back the arrived queue at the end
246 returnOld :: StrictList m -> Maybe (Either a a) -> IO (Maybe (Either a a))
247 returnOld old r = do putMVar arrived old; return r
249 -- as a side-effect, this left-biases the list
250 checkArrived :: [m -> Maybe a] -> StrictList m -> (StrictList m, Maybe a)
251 checkArrived matches list = go list Nil
252 where
253 go Nil Nil = (Nil, Nothing)
254 go Nil r = go r Nil
255 go (Append xs ys) tl = go xs (append ys tl)
256 go (Snoc xs x) tl = go xs (Cons x tl)
257 go (Cons x xs) tl
258 | Just y <- checkMatches matches x = (append xs tl, Just y)
259 | otherwise = let !(rest,r) = go xs tl in (Cons x rest, r)
261 checkMatches :: [m -> Maybe a] -> m -> Maybe a
262 checkMatches [] _ = Nothing
263 checkMatches (m:ms) a = case m a of Nothing -> checkMatches ms a
264 Just b -> Just b
266-- | Weak reference to a CQueue
267mkWeakCQueue :: CQueue a -> IO () -> IO (Weak (CQueue a))
268mkWeakCQueue m@(CQueue (StrictMVar (MVar m#)) _ _) f = IO $ \s ->
269#if MIN_VERSION_base(4,9,0)
270 case mkWeak# m# m (unIO f) s of (# s1, w #) -> (# s1, Weak w #)
271#else
272 case mkWeak# m# m f s of (# s1, w #) -> (# s1, Weak w #)
273#endif
275queueSize :: CQueue a -> IO Int
276queueSize (CQueue _ _ size) = readTVarIO size