git.iko.soy/streamly
on 5054813 help sign in
streamly/src/Streamly/Internal/Data/Stream/SVar.hs
205 lines 7.6 kB history raw
1{-# LANGUAGE CPP #-}
2{-# LANGUAGE FlexibleContexts #-}
4-- |
5-- Module : Streamly.Internal.Data.Stream.SVar
6-- Copyright : (c) 2017 Harendra Kumar
7--
8-- License : BSD3
9-- Maintainer : streamly@composewell.com
10-- Stability : experimental
11-- Portability : GHC
12--
13--
14module Streamly.Internal.Data.Stream.SVar
15 ( fromSVar
16 , fromStreamVar
17 , fromProducer
18 , fromConsumer
19 , toSVar
20 , pushToFold
21 )
22where
24import Control.Exception (fromException)
25import Control.Monad (when, void)
26import Control.Monad.Catch (throwM)
27import Control.Monad.IO.Class (MonadIO(liftIO))
28import Data.IORef (newIORef, readIORef, mkWeakIORef, writeIORef)
29import Data.Maybe (isNothing)
30import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)
31import System.Mem (performMajorGC)
33import Streamly.Internal.Data.SVar
34import Streamly.Internal.Data.Stream.StreamK hiding (reverse)
36-- | Pull a stream from an SVar.
37{-# NOINLINE fromStreamVar #-}
38fromStreamVar :: MonadAsync m => SVar Stream m a -> Stream m a
39fromStreamVar sv = mkStream $ \st yld sng stp -> do
40 list <- readOutputQ sv
41 -- Reversing the output is important to guarantee that we process the
42 -- outputs in the same order as they were generated by the constituent
43 -- streams.
44 foldStream st yld sng stp $ processEvents $ reverse list
46 where
48 allDone stp = do
49 when (svarInspectMode sv) $ do
50 t <- liftIO $ getTime Monotonic
51 liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
52 liftIO $ printSVar sv "SVar Done"
53 stp
55 {-# INLINE processEvents #-}
56 processEvents [] = mkStream $ \st yld sng stp -> do
57 done <- postProcess sv
58 if done
59 then allDone stp
60 else foldStream st yld sng stp $ fromStreamVar sv
62 processEvents (ev : es) = mkStream $ \st yld sng stp -> do
63 let rest = processEvents es
64 case ev of
65 ChildYield a -> yld a rest
66 ChildStop tid e -> do
67 accountThread sv tid
68 case e of
69 Nothing -> do
70 stop <- shouldStop tid
71 if stop
72 then liftIO (cleanupSVar sv) >> allDone stp
73 else foldStream st yld sng stp rest
74 Just ex ->
75 case fromException ex of
76 Just ThreadAbort ->
77 foldStream st yld sng stp rest
78 Nothing -> liftIO (cleanupSVar sv) >> throwM ex
79 shouldStop tid =
80 case svarStopStyle sv of
81 StopNone -> return False
82 StopAny -> return True
83 StopBy -> do
84 sid <- liftIO $ readIORef (svarStopBy sv)
85 return $ if tid == sid then True else False
87{-# INLINE fromSVar #-}
88fromSVar :: (MonadAsync m, IsStream t) => SVar Stream m a -> t m a
89fromSVar sv =
90 mkStream $ \st yld sng stp -> do
91 ref <- liftIO $ newIORef ()
92 _ <- liftIO $ mkWeakIORef ref hook
93 -- We pass a copy of sv to fromStreamVar, so that we know that it has
94 -- no other references, when that copy gets garbage collected "ref"
95 -- will get garbage collected and our hook will be called.
96 foldStreamShared st yld sng stp $
97 fromStream $ fromStreamVar sv{svarRef = Just ref}
98 where
100 hook = do
101 when (svarInspectMode sv) $ do
102 r <- liftIO $ readIORef (svarStopTime (svarStats sv))
103 when (isNothing r) $
104 printSVar sv "SVar Garbage Collected"
105 cleanupSVar sv
106 -- If there are any SVars referenced by this SVar a GC will prompt
107 -- them to be cleaned up quickly.
108 when (svarInspectMode sv) performMajorGC
110-- | Write a stream to an 'SVar' in a non-blocking manner. The stream can then
111-- be read back from the SVar using 'fromSVar'.
112toSVar :: (IsStream t, MonadAsync m) => SVar Stream m a -> t m a -> m ()
113toSVar sv m = toStreamVar sv (toStream m)
115-------------------------------------------------------------------------------
116-- Process events received by a fold consumer from a stream producer
117-------------------------------------------------------------------------------
119-- | Pull a stream from an SVar.
120{-# NOINLINE fromProducer #-}
121fromProducer :: MonadAsync m => SVar Stream m a -> Stream m a
122fromProducer sv = mkStream $ \st yld sng stp -> do
123 list <- readOutputQ sv
124 -- Reversing the output is important to guarantee that we process the
125 -- outputs in the same order as they were generated by the constituent
126 -- streams.
127 foldStream st yld sng stp $ processEvents $ reverse list
129 where
131 allDone stp = do
132 when (svarInspectMode sv) $ do
133 t <- liftIO $ getTime Monotonic
134 liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
135 liftIO $ printSVar sv "SVar Done"
136 sendStopToProducer sv
137 stp
139 {-# INLINE processEvents #-}
140 processEvents [] = mkStream $ \st yld sng stp -> do
141 foldStream st yld sng stp $ fromProducer sv
143 processEvents (ev : es) = mkStream $ \_ yld _ stp -> do
144 let rest = processEvents es
145 case ev of
146 ChildYield a -> yld a rest
147 ChildStop tid e -> do
148 accountThread sv tid
149 case e of
150 Nothing -> allDone stp
151 Just _ -> error "Bug: fromProducer: received exception"
153-------------------------------------------------------------------------------
154-- Process events received by the producer thread from the consumer side
155-------------------------------------------------------------------------------
157-- XXX currently only one event is sent by a fold consumer to the stream
158-- producer. But we can potentially have multiple events e.g. the fold step can
159-- generate exception more than once and the producer can ignore those
160-- exceptions or handle them and still keep driving the fold.
161--
162{-# NOINLINE fromConsumer #-}
163fromConsumer :: MonadAsync m => SVar Stream m a -> m Bool
164fromConsumer sv = do
165 (list, _) <- liftIO $ readOutputQBasic (outputQueueFromConsumer sv)
166 -- Reversing the output is important to guarantee that we process the
167 -- outputs in the same order as they were generated by the constituent
168 -- streams.
169 processEvents $ reverse list
171 where
173 {-# INLINE processEvents #-}
174 processEvents [] = return False
175 processEvents (ev : _) = do
176 case ev of
177 ChildStop _ e -> do
178 case e of
179 Nothing -> return True
180 Just ex -> throwM ex
181 ChildYield _ -> error "Bug: fromConsumer: invalid ChildYield event"
183-- push values to a fold worker via an SVar. Returns whether the fold is done.
184{-# INLINE pushToFold #-}
185pushToFold :: MonadAsync m => SVar Stream m a -> a -> m Bool
186pushToFold sv a = do
187 -- Check for exceptions before decrement so that we do not
188 -- block forever if the child already exited with an exception.
189 --
190 -- We avoid a race between the consumer fold sending an event and we
191 -- blocking on decrementBufferLimit by waking up the producer thread in
192 -- sendToProducer before any event is sent by the fold to the producer
193 -- stream.
194 let qref = outputQueueFromConsumer sv
195 done <- do
196 (_, n) <- liftIO $ readIORef qref
197 if (n > 0)
198 then fromConsumer sv
199 else return False
200 if done
201 then return True
202 else liftIO $ do
203 decrementBufferLimit sv
204 void $ send sv (ChildYield a)
205 return False