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