@@ -1,8 +1,9 @@
11-- | Concurrent queue for single reader, single writer
2−{-# LANGUAGE MagicHash, UnboxedTuples #-}
2+{-# LANGUAGE MagicHash, UnboxedTuples, PatternGuards #-}
33module Control.Distributed.Process.Internal.CQueue
44 ( CQueue
55 , BlockSpec(..)
6+ , MatchOn(..)
67 , newCQueue
78 , enqueue
89 , dequeue
@@ -12,14 +13,17 @@ module Control.Distributed.Process.Internal.CQueue
1213import Prelude hiding (length, reverse)
1314import Control.Concurrent.STM
1415 ( atomically
16+ , STM
1517 , TChan
18+ , tryReadTChan
1619 , newTChan
1720 , writeTChan
1821 , readTChan
19− , tryReadTChan
22+ , orElse
23+ , retry
2024 )
2125import Control.Applicative ((<$>), (<*>))
22−import Control.Exception (mask, onException)
26+import Control.Exception (mask_, onException)
2327import System.Timeout (timeout)
2428import Control.Distributed.Process.Internal.StrictMVar
2529 ( StrictMVar(StrictMVar)
@@ -29,9 +33,9 @@ import Control.Distributed.Process.Internal.StrictMVar
2933 )
3034import Control.Distributed.Process.Internal.StrictList
3135 ( StrictList(..)
32− , reverse
33− , reverse'
3436 )
37+import qualified Control.Distributed.Process.Internal.StrictList as S
38+import Data.Maybe (fromJust)
3539import GHC.MVar (MVar(MVar))
3640import GHC.IO (IO(IO))
3741import GHC.Prim (mkWeak#)
@@ -55,71 +59,176 @@ data BlockSpec =
5559 | Blocking
5660 | Timeout Int
5761
62+data MatchOn m a
63+ = MatchMsg (m -> Maybe a)
64+ | MatchChan (STM a)
65+
66+type MatchChunks m a = [Either [m -> Maybe a] [STM a]]
67+
68+chunkMatches :: [MatchOn m a] -> MatchChunks m a
69+chunkMatches [] = []
70+chunkMatches (MatchMsg m : ms) = Left (m : chk) : chunkMatches rest
71+ where (chk, rest) = spanMatchMsg ms
72+chunkMatches (MatchChan r : ms) = Right (r : chk) : chunkMatches rest
73+ where (chk, rest) = spanMatchChan ms
74+
75+spanMatchMsg :: [MatchOn m a] -> ([m -> Maybe a], [MatchOn m a])
76+spanMatchMsg [] = ([],[])
77+spanMatchMsg (m : ms)
78+ | MatchMsg msg <- m = (msg:msgs, rest)
79+ | otherwise = ([], m:ms)
80+ where !(msgs,rest) = spanMatchMsg ms
81+
82+spanMatchChan :: [MatchOn m a] -> ([STM a], [MatchOn m a])
83+spanMatchChan [] = ([],[])
84+spanMatchChan (m : ms)
85+ | MatchChan stm <- m = (stm:stms, rest)
86+ | otherwise = ([], m:ms)
87+ where !(stms,rest) = spanMatchChan ms
88+
5889-- | Dequeue an element
5990--
6091-- The timeout (if any) is applied only to waiting for incoming messages, not
6192-- to checking messages that have already arrived
62−dequeue :: forall a b.
63− CQueue a -- ^ Queue
93+dequeue :: forall m a.
94+ CQueue m -- ^ Queue
6495 -> BlockSpec -- ^ Blocking behaviour
65− -> [a -> Maybe b] -- ^ List of matches
66− -> IO (Maybe b) -- ^ 'Nothing' only on timeout
67−dequeue (CQueue arrived incoming) blockSpec matches = go
96+ -> [MatchOn m a] -- ^ List of matches
97+ -> IO (Maybe a) -- ^ 'Nothing' only on timeout
98+dequeue (CQueue arrived incoming) blockSpec matchons =
99+ case blockSpec of
100+ Timeout n -> timeout n $ fmap fromJust go
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 -> go
68109 where
69− go :: IO (Maybe b)
70− go = mask $ \restore -> do
71− arr <- takeMVar arrived
72− -- We first check the arrived messages. If we get interrupted during this
73− -- search, we just put the MVar back (we haven't read from the Chan yet)
74− (arr', mb) <- onException (restore (checkArrived Nil arr))
75− (putMVar arrived arr)
76− case (mb, blockSpec) of
77− (Just b, _) -> do
78− putMVar arrived arr'
79− return (Just b)
80− (Nothing, NonBlocking) ->
81− checkNonBlocking arr'
82− (Nothing, Blocking) ->
83− Just <$> checkBlocking arr'
84− (Nothing, Timeout n) ->
85− timeout n $ checkBlocking arr'
86−
87− -- We reverse the accumulator on return only if we find a match
88− checkArrived :: StrictList a -> StrictList a -> IO (StrictList a, Maybe b)
89− checkArrived acc Nil = return (acc, Nothing)
90− checkArrived acc (Cons x xs) =
91− case check x of
92− Just y -> return (reverse' acc xs, Just y)
93− Nothing -> checkArrived (Cons x acc) xs
94−
95− -- If we call checkBlocking there may or may not be a timeout
96− checkBlocking :: StrictList a -> IO b
97− checkBlocking acc = do
98− -- readTChan is a blocking call, and hence is interruptable. If it is
99− -- interrupted, we put the value of the accumulator in 'arrived'
100− -- (as opposed to the original value), so that no messages get lost
101− -- (hence the low-level structure using mask rather than modifyMVar)
102− x <- onException (atomically $ readTChan incoming)
103− (putMVar arrived $ reverse acc)
104− case check x of
105− Nothing -> checkBlocking (Cons x acc)
106− Just y -> putMVar arrived (reverse acc) >> return y
107−
108− -- checkNonBlocking is only called if there is no timeout
109− checkNonBlocking :: StrictList a -> IO (Maybe b)
110− checkNonBlocking acc = do
111− -- tryReadTChan is *not* interruptible
112− mx <- atomically $ tryReadTChan incoming
113− case mx of
114− Nothing -> putMVar arrived (reverse acc) >> return Nothing
115− Just x -> case check x of
116− Nothing -> checkNonBlocking (Cons x acc)
117− Just y -> putMVar arrived (reverse acc) >> return (Just y)
118−
119− check :: a -> Maybe b
120− check = checkMatches matches
121−
122− checkMatches :: [a -> Maybe b] -> a -> Maybe b
110+ chunks = chunkMatches matchons
111+
112+ go = 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'
121+
122+ waitChans ports on_block =
123+ foldr orElse on_block (map (fmap Just) ports)
124+
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)
133+
134+ goCheck [] old = goWait old
135+
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
141+
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
152+
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)
163+
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
170+
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)
196+
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)
205+
206+ goCheck1 [] m old = goWait (Snoc old m)
207+
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
213+
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)
218+
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
222+
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 = S.foldr f (Nil, Nothing) list
226+ where
227+ f x (rest,r)
228+ | Just y <- checkMatches matches x = (rest, Just y)
229+ | otherwise = (Cons x rest, r)
230+
231+ checkMatches :: [m -> Maybe a] -> m -> Maybe a
123232 checkMatches [] _ = Nothing
124233 checkMatches (m:ms) a = case m a of Nothing -> checkMatches ms a
125234 Just b -> Just b