{-# LANGUAGE UnboxedTuples #-}

-- |
-- Module      : Streamly.Internal.Control.ForkIO
-- Copyright   : (c) 2017 Composewell Technologies
-- License     : BSD-3-Clause
-- Maintainer  : streamly@composewell.com
-- Stability   : experimental
-- Portability : GHC

module Streamly.Internal.Control.ForkIO
    ( rawForkIO
    , forkIOManaged
    , forkManagedWith
    )
where

import Control.Concurrent (ThreadId, forkIO, killThread)
import Control.Monad.IO.Class (MonadIO(..))
import GHC.Conc (ThreadId(..))
import GHC.Exts
import GHC.IO (IO(..))
import System.Mem.Weak (addFinalizer)

-- | Stolen from the async package. The perf improvement is modest, 2% on a
-- thread heavy benchmark (parallel composition using noop computations).
-- A version of forkIO that does not include the outer exception
-- handler: saves a bit of time when we will be installing our own
-- exception handler.
{-# INLINE rawForkIO #-}
rawForkIO :: IO () -> IO ThreadId
rawForkIO :: IO () -> IO ThreadId
rawForkIO (IO State# RealWorld -> (# State# RealWorld, () #)
action) = (State# RealWorld -> (# State# RealWorld, ThreadId #))
-> IO ThreadId
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO ((State# RealWorld -> (# State# RealWorld, ThreadId #))
 -> IO ThreadId)
-> (State# RealWorld -> (# State# RealWorld, ThreadId #))
-> IO ThreadId
forall a b. (a -> b) -> a -> b
$ \ State# RealWorld
s ->
   case (State# RealWorld -> (# State# RealWorld, () #))
-> State# RealWorld -> (# State# RealWorld, ThreadId# #)
forall a.
a -> State# RealWorld -> (# State# RealWorld, ThreadId# #)
fork# State# RealWorld -> (# State# RealWorld, () #)
action State# RealWorld
s of (# State# RealWorld
s1, ThreadId#
tid #) -> (# State# RealWorld
s1, ThreadId# -> ThreadId
ThreadId ThreadId#
tid #)

-- | Fork a thread that is automatically killed as soon as the reference to the
-- returned threadId is garbage collected.
--
{-# INLINABLE forkManagedWith #-}
forkManagedWith :: MonadIO m => (m () -> m ThreadId) -> m () -> m ThreadId
forkManagedWith :: forall (m :: * -> *).
MonadIO m =>
(m () -> m ThreadId) -> m () -> m ThreadId
forkManagedWith m () -> m ThreadId
fork m ()
action = do
    ThreadId
tid <- m () -> m ThreadId
fork m ()
action
    IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ ThreadId -> IO () -> IO ()
forall key. key -> IO () -> IO ()
addFinalizer ThreadId
tid (ThreadId -> IO ()
killThread ThreadId
tid)
    ThreadId -> m ThreadId
forall (m :: * -> *) a. Monad m => a -> m a
return ThreadId
tid

-- | Fork a thread that is automatically killed as soon as the reference to the
-- returned threadId is garbage collected.
--
{-# INLINABLE forkIOManaged #-}
forkIOManaged :: IO () -> IO ThreadId
forkIOManaged :: IO () -> IO ThreadId
forkIOManaged = (IO () -> IO ThreadId) -> IO () -> IO ThreadId
forall (m :: * -> *).
MonadIO m =>
(m () -> m ThreadId) -> m () -> m ThreadId
forkManagedWith IO () -> IO ThreadId
forkIO