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