git.iko.soy/distributed-process
on 3e44157 help sign in
distributed-process/src/Control/Distributed/Process/Internal/CQueue.hs
249 lines 8.5 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 ) where
15import Prelude hiding (length, reverse)
16import Control.Concurrent.STM
17 ( atomically
18 , STM
19 , TChan
20 , tryReadTChan
21 , newTChan
22 , writeTChan
23 , readTChan
24 , orElse
25 , retry
26 )
27import Control.Applicative ((<$>), (<*>))
28import Control.Exception (mask_, onException)
29import System.Timeout (timeout)
30import Control.Distributed.Process.Internal.StrictMVar
31 ( StrictMVar(StrictMVar)
32 , newMVar
33 , takeMVar
34 , putMVar
35 )
36import Control.Distributed.Process.Internal.StrictList
37 ( StrictList(..)
38 , append
39 )
40import Data.Maybe (fromJust)
41import GHC.MVar (MVar(MVar))
42import GHC.IO (IO(IO))
43import GHC.Prim (mkWeak#)
44import GHC.Weak (Weak(Weak))
46-- We use a TCHan rather than a Chan so that we have a non-blocking read
47data CQueue a = CQueue (StrictMVar (StrictList a)) -- Arrived
48 (TChan a) -- Incoming
50newCQueue :: IO (CQueue a)
51newCQueue = CQueue <$> newMVar Nil <*> atomically newTChan
53-- | Enqueue an element
54--
55-- Enqueue is strict.
56enqueue :: CQueue a -> a -> IO ()
57enqueue (CQueue _arrived incoming) !a = atomically $ writeTChan incoming a
59-- | Variant of enqueue for use in the STM monad.
60enqueueSTM :: CQueue a -> a -> STM ()
61enqueueSTM (CQueue _arrived incoming) !a = writeTChan incoming a
63data BlockSpec =
64 NonBlocking
65 | Blocking
66 | Timeout Int
68data MatchOn m a
69 = MatchMsg (m -> Maybe a)
70 | MatchChan (STM a)
72type MatchChunks m a = [Either [m -> Maybe a] [STM a]]
74chunkMatches :: [MatchOn m a] -> MatchChunks m a
75chunkMatches [] = []
76chunkMatches (MatchMsg m : ms) = Left (m : chk) : chunkMatches rest
77 where (chk, rest) = spanMatchMsg ms
78chunkMatches (MatchChan r : ms) = Right (r : chk) : chunkMatches rest
79 where (chk, rest) = spanMatchChan ms
81spanMatchMsg :: [MatchOn m a] -> ([m -> Maybe a], [MatchOn m a])
82spanMatchMsg [] = ([],[])
83spanMatchMsg (m : ms)
84 | MatchMsg msg <- m = (msg:msgs, rest)
85 | otherwise = ([], m:ms)
86 where !(msgs,rest) = spanMatchMsg ms
88spanMatchChan :: [MatchOn m a] -> ([STM a], [MatchOn m a])
89spanMatchChan [] = ([],[])
90spanMatchChan (m : ms)
91 | MatchChan stm <- m = (stm:stms, rest)
92 | otherwise = ([], m:ms)
93 where !(stms,rest) = spanMatchChan ms
95-- | Dequeue an element
96--
97-- The timeout (if any) is applied only to waiting for incoming messages, not
98-- to checking messages that have already arrived
99dequeue :: forall m a.
100 CQueue m -- ^ Queue
101 -> BlockSpec -- ^ Blocking behaviour
102 -> [MatchOn m a] -- ^ List of matches
103 -> IO (Maybe a) -- ^ 'Nothing' only on timeout
104dequeue (CQueue arrived incoming) blockSpec matchons = mask_ $
105 case blockSpec of
106 Timeout n -> timeout n $ fmap fromJust run
107 _other ->
108 case chunks of
109 [Right ports] -> -- channels only, this is easy:
110 case blockSpec of
111 NonBlocking -> atomically $ waitChans ports (return Nothing)
112 _ -> atomically $ waitChans ports retry
113 -- no onException needed
114 _other -> run
115 where
116 chunks = chunkMatches matchons
118 run = do
119 arr <- takeMVar arrived
120 let grabNew xs = do
121 r <- atomically $ tryReadTChan incoming
122 case r of
123 Nothing -> return xs
124 Just x -> grabNew (Snoc xs x)
125 arr' <- grabNew arr
126 goCheck chunks arr'
128 waitChans ports on_block =
129 foldr orElse on_block (map (fmap Just) ports)
131 --
132 -- First check the MatchChunks against the messages already in the
133 -- mailbox. For channel matches, we do a non-blocking check at
134 -- this point.
135 --
136 goCheck :: MatchChunks m a
137 -> StrictList m -- messages to check, in this order
138 -> IO (Maybe a)
140 goCheck [] old = goWait old
142 goCheck (Right ports : rest) old = do
143 r <- atomically $ waitChans ports (return Nothing) -- does not block
144 case r of
145 Just _ -> returnOld old r
146 Nothing -> goCheck rest old
148 goCheck (Left matches : rest) old = do
149 -- checkArrived might in principle take arbitrary time, so
150 -- we ought to call restore and use an exception handler. However,
151 -- the check is usually fast (just a comparison), and the overhead
152 -- of passing around restore and setting up exception handlers is
153 -- high. So just don't use expensive matchIfs!
154 case checkArrived matches old of
155 (old', Just r) -> returnOld old' (Just r)
156 (old', Nothing) -> goCheck rest old'
157 -- use the result list, which is now left-biased
159 --
160 -- Construct an STM transaction that looks at the relevant channels
161 -- in the correct order.
162 --
163 mkSTM :: MatchChunks m a -> STM (Either m a)
164 mkSTM [] = retry
165 mkSTM (Left _ : rest)
166 = fmap Left (readTChan incoming) `orElse` mkSTM rest
167 mkSTM (Right ports : rest)
168 = foldr orElse (mkSTM rest) (map (fmap Right) ports)
170 waitIncoming :: IO (Maybe (Either m a))
171 waitIncoming = case blockSpec of
172 NonBlocking -> atomically $ fmap Just stm `orElse` return Nothing
173 _ -> atomically $ fmap Just stm
174 where
175 stm = mkSTM chunks
177 --
178 -- The initial pass didn't find a message, so now we go into blocking
179 -- mode.
180 --
181 -- Contents of 'arrived' from now on is (old ++ new), and
182 -- messages that arrive are snocced onto new.
183 --
184 goWait old = do
185 r <- waitIncoming `onException` putMVar arrived old
186 case r of
187 -- Nothing => non-blocking and no message
188 Nothing -> returnOld old Nothing
189 Just e -> case e of
190 --
191 -- Left => message arrived in the process mailbox. We now have to
192 -- run through the MatchChunks checking each one, because we might
193 -- have a situation where the first chunk fails to match and the
194 -- second chunk is a channel match and there *is* a message in the
195 -- channel. In that case the channel wins.
196 --
197 Left m -> goCheck1 chunks m old
198 --
199 -- Right => message arrived on a channel first
200 --
201 Right a -> returnOld old (Just a)
203 --
204 -- A message arrived in the process inbox; check the MatchChunks for
205 -- a valid match.
206 --
207 goCheck1 :: MatchChunks m a
208 -> m -- single message to check
209 -> StrictList m -- old messages we have already checked
210 -> IO (Maybe a)
212 goCheck1 [] m old = goWait (Snoc old m)
214 goCheck1 (Right ports : rest) m old = do
215 r <- atomically $ waitChans ports (return Nothing) -- does not block
216 case r of
217 Nothing -> goCheck1 rest m old
218 Just _ -> returnOld (Snoc old m) r
220 goCheck1 (Left matches : rest) m old = do
221 case checkMatches matches m of
222 Nothing -> goCheck1 rest m old
223 Just p -> returnOld old (Just p)
225 -- a common pattern for putting back the arrived queue at the end
226 returnOld :: StrictList m -> Maybe a -> IO (Maybe a)
227 returnOld old r = do putMVar arrived old; return r
229 -- as a side-effect, this left-biases the list
230 checkArrived :: [m -> Maybe a] -> StrictList m -> (StrictList m, Maybe a)
231 checkArrived matches list = go list Nil
232 where
233 go Nil Nil = (Nil, Nothing)
234 go Nil r = go r Nil
235 go (Append xs ys) tl = go xs (append ys tl)
236 go (Snoc xs x) tl = go xs (Cons x tl)
237 go (Cons x xs) tl
238 | Just y <- checkMatches matches x = (append xs tl, Just y)
239 | otherwise = let !(rest,r) = go xs tl in (Cons x rest, r)
241 checkMatches :: [m -> Maybe a] -> m -> Maybe a
242 checkMatches [] _ = Nothing
243 checkMatches (m:ms) a = case m a of Nothing -> checkMatches ms a
244 Just b -> Just b
246-- | Weak reference to a CQueue
247mkWeakCQueue :: CQueue a -> IO () -> IO (Weak (CQueue a))
248mkWeakCQueue m@(CQueue (StrictMVar (MVar m#)) _) f = IO $ \s ->
249 case mkWeak# m# m f s of (# s1, w #) -> (# s1, Weak w #)