@@ -0,0 +1,330 @@
1+-- |
2+-- Module : Streamly.Internal.Data.Fold.Concurrent.Channel.Type
3+-- Copyright : (c) 2017, 2022 Composewell Technologies
4+-- License : BSD-3-Clause
5+-- Maintainer : streamly@composewell.com
6+-- Stability : experimental
7+-- Portability : GHC
8+
9+module Streamly.Internal.Data.Fold.Concurrent.Channel.Type
10+ ( Channel (..)
11+ , newChannel
12+ , Config
13+ , sendToWorker
14+ , checkFoldStatus
15+ , dumpSVar
16+ )
17+where
18+
19+#include "inline.hs"
20+
21+import Control.Concurrent (ThreadId, myThreadId, tryPutMVar)
22+import Control.Concurrent.MVar (MVar, newEmptyMVar, takeMVar)
23+import Control.Exception (SomeException(..))
24+import Control.Monad (void)
25+import Control.Monad.Catch (throwM)
26+import Control.Monad.IO.Class (MonadIO(..))
27+import Control.Monad.Trans.Control (MonadBaseControl)
28+import Data.IORef (IORef, newIORef, readIORef)
29+import Data.List (intersperse)
30+import Streamly.Internal.Control.Concurrent (MonadAsync, askRunInIO)
31+import Streamly.Internal.Control.ForkLifted (doFork)
32+import Streamly.Internal.Data.Fold (Fold(..))
33+import Streamly.Internal.Data.Stream.Channel.Dispatcher (dumpSVarStats)
34+import Streamly.Internal.Data.Stream.Channel.Worker (sendWithDoorBell)
35+
36+import qualified Streamly.Internal.Data.Fold as Fold
37+import qualified Streamly.Internal.Data.Stream.StreamD.Type as D
38+
39+import Streamly.Internal.Data.Stream.Channel.Types
40+
41+data Channel m a b = Channel
42+ {
43+ -- FORWARD FLOW: Flow of data from the driver to the consuming fold
44+
45+ -- XXX This is inputQueue instead.
46+
47+ -- Shared output queue (events, length)
48+ --
49+ -- [LOCKING] Frequent locked access. This is updated by the driver on each
50+ -- yield and once in a while read by the consumer fold thread.
51+ --
52+ -- XXX Use a different type than ChildEvent. We can do with a simpler type
53+ -- in folds.
54+ outputQueue :: IORef ([ChildEvent a], Int)
55+ -- This is capped to maxBufferLimit if set to more than that. Otherwise
56+ -- potentially each worker may yield one value to the buffer in the worst
57+ -- case exceeding the requested buffer size.
58+ , maxBufferLimit :: Limit
59+
60+ -- [LOCKING] Infrequent MVar. Used when the outputQ transitions from empty
61+ -- to non-empty.
62+ , outputDoorBell :: MVar () -- signal the consumer about output
63+ , readOutputQ :: m [ChildEvent a]
64+
65+ -- receive async events from the fold consumer to the driver.
66+ , outputQueueFromConsumer :: IORef ([ChildEvent b], Int)
67+ , outputDoorBellFromConsumer :: MVar ()
68+ , bufferSpaceDoorBell :: MVar ()
69+
70+ -- cleanup: to track garbage collection of SVar --
71+ , svarRef :: Maybe (IORef ())
72+
73+ -- Stats --
74+ , svarStats :: SVarStats
75+
76+ -- Diagnostics --
77+ , svarInspectMode :: Bool
78+ , svarCreator :: ThreadId
79+ }
80+
81+{-# NOINLINE dumpSVar #-}
82+dumpSVar :: Channel m a b -> IO String
83+dumpSVar sv = do
84+ xs <- sequence $ intersperse (return "\n")
85+ [ return (dumpCreator (svarCreator sv))
86+ , return "---------CURRENT STATE-----------"
87+ , dumpOutputQ (outputQueue sv)
88+ -- XXX print the types of events in the outputQueue, first 5
89+ , dumpDoorBell (outputDoorBell sv)
90+ , return "---------STATS-----------\n"
91+ , dumpSVarStats (svarInspectMode sv) Nothing (svarStats sv)
92+ ]
93+ return $ concat xs
94+
95+-------------------------------------------------------------------------------
96+-- Support for running folds concurrently
97+-------------------------------------------------------------------------------
98+
99+-- $concurrentFolds
100+--
101+-- To run folds concurrently, we need to decouple the fold execution from the
102+-- stream production. We use the SVar to do that, we have a single worker
103+-- pushing the stream elements to the SVar and on the consumer side a fold
104+-- driver pulls the values and folds them.
105+--
106+-- @
107+--
108+-- Fold worker <------Channel<------Fold driver
109+-- | exceptions |
110+-- --------------->
111+--
112+-- @
113+--
114+-- We need a channel for pushing exceptions from the fold worker to the fold
115+-- driver. The stream may be pushed to multiple folds at the same time. For
116+-- that we need one Channel per fold:
117+--
118+-- @
119+--
120+-- Fold worker <------Channel--
121+-- | |
122+-- Fold worker <------Channel------Driver
123+-- | |
124+-- Fold worker <------Channel--
125+--
126+-- @
127+--
128+-- Note: If the stream pusher terminates due to an exception, we do not
129+-- actively terminate the fold. It gets cleaned up by the GC.
130+
131+-------------------------------------------------------------------------------
132+-- Process events received by a fold worker from a fold driver
133+-------------------------------------------------------------------------------
134+
135+sendToDriver :: Channel m a b -> ChildEvent b -> IO Int
136+sendToDriver sv msg = do
137+ -- In case the producer stream is blocked on pushing to the fold buffer
138+ -- then wake it up so that it can check for the stop event or exception
139+ -- being sent to it otherwise we will be deadlocked.
140+ -- void $ tryPutMVar (pushBufferMVar sv) ()
141+ sendWithDoorBell (outputQueueFromConsumer sv)
142+ (outputDoorBellFromConsumer sv) msg
143+
144+sendYieldToDriver :: MonadIO m => Channel m a b -> b -> m ()
145+sendYieldToDriver sv res = liftIO $ do
146+ void $ sendToDriver sv (ChildYield res)
147+
148+{-# NOINLINE sendExceptionToDriver #-}
149+sendExceptionToDriver :: Channel m a b -> SomeException -> IO ()
150+sendExceptionToDriver sv e = do
151+ tid <- myThreadId
152+ void $ sendToDriver sv (ChildStop tid (Just e))
153+
154+data FromSVarState m a b =
155+ FromSVarRead (Channel m a b)
156+ | FromSVarLoop (Channel m a b) [ChildEvent a]
157+
158+{-# INLINE_NORMAL fromProducerD #-}
159+fromProducerD :: MonadIO m => Channel m a b -> D.Stream m a
160+fromProducerD svar = D.Stream step (FromSVarRead svar)
161+
162+ where
163+
164+ {-# INLINE_LATE step #-}
165+ step _ (FromSVarRead sv) = do
166+ list <- readOutputQ sv
167+ -- Reversing the output is important to guarantee that we process the
168+ -- outputs in the same order as they were generated by the constituent
169+ -- streams.
170+ return $ D.Skip $ FromSVarLoop sv (Prelude.reverse list)
171+
172+ step _ (FromSVarLoop sv []) = return $ D.Skip $ FromSVarRead sv
173+ step _ (FromSVarLoop sv (ev : es)) = do
174+ case ev of
175+ ChildYield a -> return $ D.Yield a (FromSVarLoop sv es)
176+ ChildStopChannel -> return D.Stop
177+ _ -> undefined
178+
179+{-# INLINE readOutputQChan #-}
180+readOutputQChan :: Channel m a b -> IO ([ChildEvent a], Int)
181+readOutputQChan chan = do
182+ let ss = if svarInspectMode chan then Just (svarStats chan) else Nothing
183+ r@(_, n) <- readOutputQRaw (outputQueue chan) ss
184+ if n <= 0
185+ then do
186+ liftIO
187+ $ withDiagMVar
188+ (svarInspectMode chan)
189+ (dumpSVar chan)
190+ "readOutputQChan: nothing to do"
191+ $ takeMVar (outputDoorBell chan)
192+ readOutputQRaw (outputQueue chan) ss
193+ else return r
194+
195+{-# INLINE readOutputQDB #-}
196+readOutputQDB :: Channel m a b -> IO ([ChildEvent a], Int)
197+readOutputQDB chan = do
198+ r <- readOutputQChan chan
199+ -- XXX We can do this only if needed, if someone sleeps because of buffer
200+ -- then they can set a flag and we ring the doorbell only if the flag is
201+ -- set. Like we do in sendWorkerWait for streams.
202+ _ <- tryPutMVar (bufferSpaceDoorBell chan) ()
203+ return r
204+
205+mkNewChannel :: forall m a b. MonadIO m => Config -> IO (Channel m a b)
206+mkNewChannel cfg = do
207+ outQ <- newIORef ([], 0)
208+ outQMv <- newEmptyMVar
209+ outQRev <- newIORef ([], 0)
210+ outQMvRev <- newEmptyMVar
211+ bufferMv <- newEmptyMVar
212+
213+ stats <- newSVarStats
214+ tid <- myThreadId
215+
216+ let getSVar :: Channel m a b -> Channel m a b
217+ getSVar sv = Channel
218+ { outputQueue = outQ
219+ , outputDoorBell = outQMv
220+ , outputQueueFromConsumer = outQRev
221+ , outputDoorBellFromConsumer = outQMvRev
222+ , bufferSpaceDoorBell = bufferMv
223+ , maxBufferLimit = getMaxBuffer cfg
224+ , readOutputQ = liftIO $ fmap fst (readOutputQDB sv)
225+ , svarRef = Nothing
226+ , svarInspectMode = getInspectMode cfg
227+ , svarCreator = tid
228+ , svarStats = stats
229+ }
230+
231+ let sv = getSVar sv in return sv
232+
233+{-# INLINABLE newChannel #-}
234+{-# SPECIALIZE newChannel ::
235+ (Config -> Config) -> Fold IO a b -> IO (Channel IO a b) #-}
236+newChannel :: (MonadIO m, MonadBaseControl IO m) =>
237+ (Config -> Config) -> Fold m a b -> m (Channel m a b)
238+newChannel modifier f = do
239+ sv <- liftIO $ mkNewChannel (modifier defaultConfig)
240+ mrun <- askRunInIO
241+ void $ doFork (work sv) mrun (sendExceptionToDriver sv)
242+ return sv
243+
244+ where
245+
246+ {-# NOINLINE work #-}
247+ work chan =
248+ let f1 = Fold.rmapM (void . sendYieldToDriver chan) f
249+ in D.fold f1 $ fromProducerD chan
250+
251+-------------------------------------------------------------------------------
252+-- Process events received by the driver thread from the fold worker side
253+-------------------------------------------------------------------------------
254+
255+-- XXX currently only one event is sent by a fold consumer to the stream
256+-- producer. But we can potentially have multiple events e.g. the fold step can
257+-- generate exception more than once and the producer can ignore those
258+-- exceptions or handle them and still keep driving the fold.
259+
260+-- XXX In case of scan this could be a stream.
261+
262+-- | Poll for events sent by the fold worker to the fold driver. The fold
263+-- consumer can send a "Stop" event or an exception. When a "Stop" is received
264+-- this function returns 'True'. If an exception is recieved then it throws the
265+-- exception.
266+--
267+{-# NOINLINE checkFoldStatus #-}
268+checkFoldStatus :: MonadAsync m => Channel m a b -> m (Maybe b)
269+checkFoldStatus sv = do
270+ (list, _) <- liftIO $ readOutputQBasic (outputQueueFromConsumer sv)
271+ -- Reversing the output is important to guarantee that we process the
272+ -- outputs in the same order as they were generated by the constituent
273+ -- streams.
274+ processEvents $ reverse list
275+
276+ where
277+
278+ {-# INLINE processEvents #-}
279+ processEvents [] = return Nothing
280+ processEvents (ev : _) = do
281+ case ev of
282+ ChildStop _ e -> maybe undefined throwM e
283+ ChildStopChannel -> undefined
284+ ChildYield b -> return (Just b)
285+
286+{-# INLINE isBufferAvailable #-}
287+isBufferAvailable :: MonadIO m => Channel m a b -> m Bool
288+isBufferAvailable sv = do
289+ let limit = maxBufferLimit sv
290+ case limit of
291+ Unlimited -> return True
292+ Limited lim -> do
293+ (_, n) <- liftIO $ readIORef (outputQueue sv)
294+ return $ fromIntegral lim > n
295+
296+-- | Push values from a driver to a fold worker via a Channel. Before pushing a
297+-- value to the Channel it polls for events received from the fold worker. If a
298+-- stop event is received then it returns 'True' otherwise false. Propagates
299+-- exceptions received from the fold wroker.
300+--
301+{-# INLINE sendToWorker #-}
302+sendToWorker :: MonadAsync m => Channel m a b -> a -> m (Maybe b)
303+sendToWorker chan a = go
304+
305+ where
306+
307+ -- Recursive function, should we use SPEC?
308+ go = do
309+ let qref = outputQueueFromConsumer chan
310+ status <- do
311+ (_, n) <- liftIO $ readIORef qref
312+ if n > 0
313+ then checkFoldStatus chan
314+ else return Nothing
315+ case status of
316+ Just _ -> return status
317+ Nothing -> do
318+ r <- isBufferAvailable chan
319+ if r
320+ then do
321+ liftIO
322+ $ void
323+ $ sendWithDoorBell
324+ (outputQueue chan)
325+ (outputDoorBell chan)
326+ (ChildYield a)
327+ return Nothing
328+ else do
329+ () <- liftIO $ takeMVar (bufferSpaceDoorBell chan)
330+ go