TMChan.hs

{-# LANGUAGE CPP #-}
{-# LANGUAGE NoMonomorphismRestriction #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE KindSignatures #-}

-- | * Introduction
--
--   Contains a simple source and sink for linking together conduits in
--   in different threads. Usage is so easy, it's best explained with an
--   example:
--
--   We first create a channel for communication...
--
--   > do chan <- atomically $ newTBMChan 16
--
--   Then we fork a new thread loading a wackton of pictures into memory. The
--   data (pictures, in this case) will be streamed down the channel to whatever
--   is on the other side.
--
--   >    _ <- forkIO . runResourceT $ loadTextures lotsOfPictures $$ sinkTBMChan chan
--
--   Finally, we connect something to the other end of the channel. In this
--   case, we connect a sink which uploads the textures one by one to the
--   graphics card.
--
--   >    runResourceT $ sourceTBMChan chan $$ Conduit.mapM_ (liftIO . uploadToGraphicsCard)
--
--   By running the two tasks in parallel, we no longer have to wait for one
--   texture to upload to the graphics card before reading the next one from
--   disk. This avoids the common switching of bottlenecks (such as between the
--   disk and graphics memory) that most loading processes seem to love.
--
--   Control.Concurrent.STM.TMChan and Control.Concurrent.STM.TBMChan are
--   re-exported for convenience.
--
--   * Caveats
--
--   It is recommended to use TBMChan as much as possible, and generally avoid
--   TMChan usage. TMChans are unbounded, and if used, the conduit pipeline
--   will no longer use a bounded amount of space. They will essentially leak
--   memory if the writer is faster than the reader.
--
--   Therefore, use bounded channels as much as possible, preferably with a
--   high bound so it will be hit infrequently.
module Data.Conduit.TMChan ( -- * Bounded Channel Connectors
                             module Control.Concurrent.STM.TBMChan
                           , sourceTBMChan
                           , sinkTBMChan
                           -- * Unbounded Channel Connectors
                           , module Control.Concurrent.STM.TMChan
                           , sourceTMChan
                           , sinkTMChan
                           -- * Parallel Combinators
                           , (>=<)
                           , mergeSources
                           , (<=>)
                           , mergeConduits
                           ) where

#if __GLASGOW_HASKELL__ < 710
import           Control.Applicative
#endif
import Control.Monad
import Control.Monad.IO.Class ( liftIO, MonadIO )
import Control.Monad.Trans.Class
import Control.Monad.Trans.Resource
import Control.Concurrent (killThread)
import Control.Concurrent.STM
import Control.Concurrent.STM.TBMChan
import Control.Concurrent.STM.TMChan

import qualified Control.Exception.Lifted as Lifted

import Data.Conduit
import qualified Data.Conduit.List as CL

chanSource
    :: MonadIO m
    => chan                     -- ^ The channel.
    -> (chan -> STM (Maybe a))  -- ^ The 'read' function.
    -> (chan -> IO  ())         -- ^ The 'close' function.
    -> Source m a
chanSource ch reader closer =
    loop
  where
    loop = do
        a <- liftSTM $ reader ch
        case a of
            Just x  -> yieldOr x close >> loop
            Nothing -> return ()
    close = liftIO $ closer ch
{-# INLINE chanSource #-}

chanSink
    :: MonadIO m
    => chan                     -- ^ The channel.
    -> (chan -> a -> STM ())    -- ^ The 'write' function.
    -> (chan -> STM ())         -- ^ The 'close' function.
    -> Sink a m ()
chanSink ch writer closer = do
    CL.mapM_ $ liftIO . atomically . writer ch
    liftSTM $ closer ch
{-# INLINE chanSink #-}

-- | A simple wrapper around a TBMChan. As data is pushed into the channel, the
--   source will read it and pass it down the conduit pipeline. When the
--   channel is closed, the source will close also.
--
--   If the channel fills up, the pipeline will stall until values are read.
sourceTBMChan :: MonadIO m => TBMChan a -> Source m a
sourceTBMChan ch = chanSource ch readTBMChan (atomically . closeTBMChan)
{-# INLINE sourceTBMChan #-}

-- | A simple wrapper around a TMChan. As data is pushed into the channel, the
--   source will read it and pass it down the conduit pipeline. When the
--   channel is closed, the source will close also.
sourceTMChan :: MonadIO m => TMChan a -> Source m a
sourceTMChan ch = chanSource ch readTMChan (atomically . closeTMChan)
{-# INLINE sourceTMChan #-}

-- | A simple wrapper around a TBMChan. As data is pushed into the sink, it
--   will magically begin to appear in the channel. If the channel is full,
--   the sink will block until space frees up.
sinkTBMChan :: MonadIO m
            => TBMChan a
            -> Bool -- ^ Should the channel be closed when the sink is closed?
            -> Sink a m ()
sinkTBMChan ch close = chanSink ch writeTBMChan (when close . closeTBMChan)
{-# INLINE sinkTBMChan #-}

-- | A simple wrapper around a TMChan. As data is pushed into this sink, it
--   will magically begin to appear in the channel.
sinkTMChan :: MonadIO m
           => TMChan a
           -> Bool -- ^ Should the channel be closed when the sink is closed?
           -> Sink a m ()
sinkTMChan ch close = chanSink ch writeTMChan (when close . closeTMChan)
{-# INLINE sinkTMChan #-}

infixl 5 >=<

-- | Modifies a TVar, returning its new value.
modifyTVar'' :: TVar a -> (a -> a) -> STM a
modifyTVar'' tv f = do x <- f <$> readTVar tv
                       writeTVar tv x
                       return x

liftSTM :: forall (m :: * -> *) a. MonadIO m => STM a -> m a
liftSTM = liftIO . atomically

-- | Combines two sources with an unbounded channel, creating a new source
--   which pulls data from a mix of the two sources: whichever produces first.
--
--   The order of the new source's data is undefined, but it will be some
--   combination of the two given sources.
(>=<) :: (MonadResource mi, MonadIO mo, MonadBaseControl IO mi)
      => Source mi a
      -> Source mi a
      -> mi (Source mo a)
sa >=< sb = mergeSources [ sa, sb ] 16
{-# INLINE (>=<) #-}

decRefcount :: TVar Int -> TBMChan a ->  STM ()
decRefcount tv chan = do n <- modifyTVar'' tv (subtract 1)
                         when (n == 0) $
                            closeTBMChan chan

-- | Merges a list of sources, putting them all into a bounded channel, and
--   returns a source which can be pulled from to pull from all the given
--   sources in a first-come-first-serve basis.
--
--   The order of the new source's data is undefined, but it will be some
--   combination of the given sources. The monad of the resultant source
--   (@mo@) is independent of the monads of the input sources (@mi@).
--
--   @since 3.0
--   All spawned threads will be removed when source is closed or upon an
--   exit from 'ResourceT' region. This means that result can only be used
--   within a 'runResourceT' scope.
--
--   @before 3.0
--   Spawned threads are not guaranteed to be closed. This may happen if
--   Source was closed before all it's input were closed.
mergeSources :: (MonadResource mi, MonadIO mo, MonadBaseControl IO mi)
             => [Source mi a] -- ^ The sources to merge.
             -> Int -- ^ The bound of the intermediate channel.
             -> mi (Source mo a)
mergeSources sx bound = do
     c <- liftSTM $ newTBMChan bound
     refcount <- liftSTM . newTVar $ length sx
     regs <- forM (map (transPipe lift) sx) $ \s -> Lifted.mask_ $ do
       register . killThread <=< runResourceT $ resourceForkIO $ s $$ chanSink c writeTBMChan $ decRefcount refcount
     return $ chanSource c readTBMChan
                           (\chan -> do liftSTM $ closeTBMChan chan
                                        mapM_ release regs)

-- | Combines two conduits with unbounded channels, creating a new conduit
--   which pulls data from a mix of the two: whichever produces first.
--
--   The order of the new conduit's output is undefined, but it will be some
--   combination of the two given conduits.
(<=>) :: (MonadIO mi, MonadThrow mi, MonadIO mo, MonadBaseControl IO mi)
      => Show i
      => Conduit i (ResourceT mi) i
      -> Conduit i (ResourceT mi) i
      -> ResourceT mi (Conduit i mo i)
sa <=> sb = mergeConduits [ sa, sb ] 16
{-# INLINE (<=>) #-}

-- | Provide an input across several conduits, putting them all into a bounded
--   channel. Returns a conduit which can be pulled from to pull from all the
--   given conduits in a first-come-first-serve basis.
--
--   The order of the new conduits's outputs is undefined, but it will be some
--   combination of the given conduits. The monad of the resultant conduit
--   (@mo@) is independent of the monads of the input conduits (@mi@).
--
-- @since 3.0
--   Closes all worker processes when resulting conduit is closed or when execution
--   leaves ResourceT context. This means that conduit is only valid inside
--   'runResouceT' scope.
--
-- @before 3.0
--   Spawned threads are not guaranteed to be closed, This may happen if threads
--   Conduit was closed before all threads have finished execution.
mergeConduits :: (MonadIO mi, MonadIO mi, MonadThrow mi, MonadIO mo, MonadBaseControl IO mi)
              => [Conduit i (ResourceT mi) o] -- ^ The conduits to merge.
              -> Int -- ^ The bound for the channels.
              -> ResourceT mi (Conduit i mo o)
mergeConduits conduits bound = do
        let len = length conduits
        refcount <- liftSTM $ newTVar len
        iChannels <- replicateM len $ liftSTM $ newTBMChan bound
        oChannel <- liftSTM $ newTBMChan bound
        regs <- forM (zip iChannels conduits)
            $ \(iChannel, conduit) -> Lifted.mask_ $
              register . killThread <=< resourceForkIO
                $  sourceTBMChan iChannel
                $$ conduit
                =$ chanSink oChannel writeTBMChan (decRefcount refcount)
        return
            $  toConsumer (chanSink iChannels writeTBMChans (mapM_ closeTBMChan))
            >> toProducer (chanSource oChannel readTBMChan (\c -> do atomically $ closeTBMChan c
                                                                     mapM_ release regs))
  where
    writeTBMChans channels a = forM_ channels $ \c -> writeTBMChan c a