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