@@ -25,6 +25,7 @@ module Streamly.Streams.Parallel
2525 , parallel
2626 , parallelEndByFirst
2727 , parallelEndByAny
28+ , tee
2829
2930 -- * Function application
3031 , mkParallel
@@ -35,6 +36,8 @@ module Streamly.Streams.Parallel
3536 )
3637where
3738
39+import Control.Concurrent (myThreadId)
40+import Control.Exception (SomeException(..), throwIO)
3841import Control.Monad (ap)
3942import Control.Monad.Base (MonadBase(..), liftBaseDefault)
4043import Control.Monad.Catch (MonadThrow, throwM)
@@ -51,7 +54,7 @@ import Prelude hiding (map)
5154
5255import qualified Data.Set as Set
5356
54−import Streamly.Streams.SVar (fromSVar)
57+import Streamly.Streams.SVar (fromSVar, fromStreamVar)
5558import Streamly.Streams.Serial (map)
5659import Streamly.SVar
5760import Streamly.Streams.StreamK (IsStream(..), Stream, mkStream, foldStream,
@@ -205,6 +208,86 @@ foldParallel f m = do
205208 pushWorkerPar sv (runOne defState{streamVar = Just sv} $ toStream m)
206209 f $ fromSVar sv
207210
211+{-# INLINE teeToSVar #-}
212+teeToSVar :: (IsStream t, MonadAsync m) => SVar Stream m a -> t m a -> t m a
213+teeToSVar svr m = mkStream $ \st yld sng stp -> do
214+ foldStreamShared st yld sng stp (go svr m)
215+
216+ where
217+
218+ go sv m0 = mkStream $ \st yld sng stp -> do
219+ let stop = do
220+ liftIO $ do
221+ incrementBufferLimit sv
222+ sendStop sv Nothing
223+ -- XXX wait for Stop response event on exception channel,
224+ -- drain the exception channel.
225+ stp
226+ sendit a = liftIO $ do
227+ -- XXX check for exceptions before decrement so that we do not
228+ -- block forever if the child already exited with an exception.
229+ decrementBufferLimit sv
230+ void $ send sv (ChildYield a)
231+ single a = sendit a >> (liftIO $ sendStop sv Nothing) >> sng a
232+ yieldk a r = sendit a >> yld a (go sv r)
233+ in foldStreamShared st yieldk single stop m0
234+
235+-- In case of folds the roles of worker and parent on an SVar are reversed. The
236+-- parent stream pushes values to an SVar instead of pulling from it and a
237+-- worker thread running the fold pulls from the SVar and folds the stream. We
238+-- keep a separate channel for pushing exceptions in the reverse direction i.e.
239+-- from the fold to the parent stream.
240+--
241+-- NOTE: If we use fromSVar here it will kill the main computation (the parent)
242+-- when the SVar goes away so we use fromStreamVar instead.
243+
244+{-# NOINLINE handleChildException #-}
245+handleChildException :: SVar t m a -> SomeException -> IO ()
246+handleChildException _sv e = do
247+ -- tid <- myThreadId
248+ -- void $ sendReverse sv (ChildStop tid (Just e))
249+ throwIO e
250+
251+-- | Redirect a copy of the stream to a supplied fold and run it concurrently
252+-- in an independent thread. The fold may buffer some elements. The buffer size
253+-- is determined by the prevailing 'maxBuffer' setting.
254+--
255+-- @
256+-- Stream m a -> m b
257+-- |
258+-- -----stream m a ---------------stream m a-----
259+--
260+-- @
261+--
262+-- @
263+-- > S.drain $ S.tee (S.mapM_ print) (S.enumerateFromTo 1 2)
264+-- 1
265+-- 2
266+-- @
267+--
268+-- Exceptions from the concurrently running fold are propagated to the current
269+-- computation. Note that, because of buffering in the fold, exceptions may be
270+-- delayed and may not correspond to the current element being processed in the
271+-- parent stream, but we guarantee that the tap finishes and all exceptions
272+-- from it are drained before the parent stream stops.
273+--
274+--
275+-- Compare with 'tap'.
276+--
277+-- @since 0.7.0
278+{-# INLINE tee #-}
279+tee :: (IsStream t, MonadAsync m) => (t m a -> m b) -> t m a -> t m a
280+tee f m = mkStream $ \st yld sng stp -> do
281+ -- Buffer size for the SVar is derived from the current state
282+ sv <- newParallelVar StopNone (adaptState st)
283+ -- XXX exception handling
284+ -- XXX if we terminate due to an exception, do we need to actively
285+ -- terminate the fold?
286+ liftIO myThreadId >>= modifyThread sv
287+ void $ doFork (void $ f $ fromStream $ fromStreamVar sv)
288+ (svarMrun sv) (handleChildException sv)
289+ foldStreamShared st yld sng stp (teeToSVar sv m)
290+
208291------------------------------------------------------------------------------
209292-- Concurrent Application
210293------------------------------------------------------------------------------