module Ki.Internal.Scope
  ( Scope,
    scoped,
    awaitAll,
    fork,
    forkWith,
    forkWith_,
    fork_,
    forkTry,
    forkTryWith,
  )
where

import Control.Concurrent (ThreadId, myThreadId, throwTo)
import Control.Concurrent.MVar (MVar, newEmptyMVar, tryPutMVar, tryTakeMVar)
import Control.Exception
  ( Exception (fromException, toException),
    MaskingState (..),
    SomeAsyncException,
    SomeException,
    asyncExceptionFromException,
    asyncExceptionToException,
    throwIO,
    try,
    uninterruptibleMask,
    pattern ErrorCall,
  )
import Control.Monad (guard, when)
import Data.Foldable (for_)
import Data.Functor (void)
import Data.IntMap (IntMap)
import qualified Data.IntMap.Lazy as IntMap.Lazy
import Data.Void (Void, absurd)
import GHC.Conc
  ( STM,
    TVar,
    atomically,
    enableAllocationLimit,
    labelThread,
    newTVarIO,
    readTVar,
    retry,
    setAllocationCounter,
    throwSTM,
    writeTVar,
  )
import GHC.Conc.Sync (readTVarIO)
import GHC.IO (unsafeUnmask)
import IntSupply (IntSupply)
import qualified IntSupply
import Ki.Internal.ByteCount (byteCountToInt64)
import Ki.Internal.IO
  ( IOResult (..),
    UnexceptionalIO (..),
    assertM,
    exceptionIs,
    interruptiblyMasked,
    unexceptionalTry,
    unexceptionalTryEither,
    uninterruptiblyMasked,
  )
import Ki.Internal.NonblockingSTM
import Ki.Internal.Propagating (Tid, peelOffPropagating, propagate)
import Ki.Internal.Thread (Thread, makeThread)
import Ki.Internal.ThreadAffinity (forkWithAffinity)
import Ki.Internal.ThreadOptions (ThreadOptions (..), defaultThreadOptions)

-- | A scope.
--
-- ==== __👉 Details__
--
-- * A scope delimits the lifetime of all threads created within it.
--
-- * A scope is only valid during the callback provided to 'Ki.scoped'.
--
-- * The thread that creates a scope is considered the parent of all threads created within it.
--
-- * All threads created within a scope can be awaited together (see 'Ki.awaitAll').
--
-- * All threads created within a scope are terminated when the scope closes.
data Scope = Scope
  { -- The MVar that a child tries to put to, in the case that it tries to propagate an exception to its parent, but
    -- gets delivered an exception from its parent concurrently (which interrupts the throw). The parent must raise
    -- exceptions in its children with asynchronous exceptions uninterruptibly masked for correctness, yet we don't want
    -- a parent in the process of tearing down to miss/ignore this exception that we're trying to propagate?
    --
    -- Why a single-celled MVar? What if two siblings are fighting to inform their parent of their death? Well, only
    -- one exception can be propagated by the parent anyway, so we wouldn't need or want both.
    Scope -> MVar SomeException
childExceptionVar :: {-# UNPACK #-} !(MVar SomeException),
    -- The set of child threads that are currently running, each keyed by a monotonically increasing int.
    Scope -> TVar (IntMap ThreadId)
childrenVar :: {-# UNPACK #-} !(TVar (IntMap ThreadId)),
    -- The supply that holds the (int) key to use for the next child thread.
    Scope -> IntSupply
nextChildIdSupply :: {-# UNPACK #-} !IntSupply,
    -- The id of the thread that created the scope, which is considered the parent of all threads created within it.
    Scope -> ThreadId
parentThreadId :: {-# UNPACK #-} !ThreadId,
    Scope -> TVar ScopeStatus
statusVar :: {-# UNPACK #-} !(TVar ScopeStatus)
  }

-- The scope status: either open (allowing new threads to be created), closing (disallowing new threads to be
-- created, and in the process of killing running children), or closed (at the very end of `scoped`)
type ScopeStatus = Int

-- The number of child threads that are guaranteed to be about to start, in the sense that only the GHC scheduler
-- can continue to delay; there's no opportunity for an async exception to strike and prevent one of these threads
-- from starting.
pattern Open :: Int
pattern $mOpen :: forall {r}. ScopeStatus -> ((# #) -> r) -> ((# #) -> r) -> r
Open <- ((>= 0) -> True)

-- The scope is closing.
pattern Closing :: Int
pattern $mClosing :: forall {r}. ScopeStatus -> ((# #) -> r) -> ((# #) -> r) -> r
$bClosing :: ScopeStatus
Closing = -1

-- The scope is closed.
pattern Closed :: Int
pattern $mClosed :: forall {r}. ScopeStatus -> ((# #) -> r) -> ((# #) -> r) -> r
$bClosed :: ScopeStatus
Closed = -2

{-# COMPLETE Open, Closing, Closed #-}

-- Internal async exception thrown by a parent thread to its children when the scope is closing.
--
-- In various places we trust without verifying that any 'ScopeClosing' exception, which is not exported by this module,
-- was indeed thrown to a thread by its parent. It is possible to write a program that violates this (just catch the
-- async exception and throw it to some other thread)... but who would do that?
data ScopeClosing
  = ScopeClosing

instance Show ScopeClosing where
  show :: ScopeClosing -> String
show ScopeClosing
_ = String
"<<internal ki exception: scope closing>>"

instance Exception ScopeClosing where
  toException :: ScopeClosing -> SomeException
toException = ScopeClosing -> SomeException
forall e. Exception e => e -> SomeException
asyncExceptionToException
  fromException :: SomeException -> Maybe ScopeClosing
fromException = SomeException -> Maybe ScopeClosing
forall e. Exception e => SomeException -> Maybe e
asyncExceptionFromException

-- | Open a scope, perform an IO action with it, then close the scope.
--
-- ==== __👉 Details__
--
-- * The thread that creates a scope is considered the parent of all threads created within it.
--
-- * A scope is only valid during the callback provided to 'Ki.scoped'.
--
-- * When a scope closes (/i.e./ just before 'Ki.scoped' returns):
--
--     * The parent thread raises an exception in all of its living children.
--     * The parent thread blocks until those threads terminate.
scoped :: (Scope -> IO a) -> IO a
scoped :: forall a. (Scope -> IO a) -> IO a
scoped Scope -> IO a
action = do
  scope :: Scope
scope@Scope {MVar SomeException
$sel:childExceptionVar:Scope :: Scope -> MVar SomeException
childExceptionVar :: MVar SomeException
childExceptionVar, TVar (IntMap ThreadId)
$sel:childrenVar:Scope :: Scope -> TVar (IntMap ThreadId)
childrenVar :: TVar (IntMap ThreadId)
childrenVar, TVar ScopeStatus
$sel:statusVar:Scope :: Scope -> TVar ScopeStatus
statusVar :: TVar ScopeStatus
statusVar} <- IO Scope
allocateScope

  ((forall a. IO a -> IO a) -> IO a) -> IO a
forall b. ((forall a. IO a -> IO a) -> IO b) -> IO b
uninterruptibleMask \forall a. IO a -> IO a
restore -> do
    Either SomeException a
result <- IO a -> IO (Either SomeException a)
forall e a. Exception e => IO a -> IO (Either e a)
try (IO a -> IO a
forall a. IO a -> IO a
restore (Scope -> IO a
action Scope
scope))

    !IntMap ThreadId
runningChildren <- do
      STM (IntMap ThreadId) -> IO (IntMap ThreadId)
forall a. STM a -> IO a
atomically do
        -- Block until we haven't committed to starting any threads. Without this, we may create a thread concurrently
        -- with closing its scope, and not grab its thread id to throw an exception to.
        ScopeStatus
starting <- TVar ScopeStatus -> STM ScopeStatus
forall a. TVar a -> STM a
readTVar TVar ScopeStatus
statusVar
        Bool -> STM ()
forall (m :: * -> *). Applicative m => Bool -> m ()
assertM (ScopeStatus
starting ScopeStatus -> ScopeStatus -> Bool
forall a. Ord a => a -> a -> Bool
>= ScopeStatus
0)
        Bool -> STM ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (ScopeStatus
starting ScopeStatus -> ScopeStatus -> Bool
forall a. Eq a => a -> a -> Bool
== ScopeStatus
0)
        -- Indicate that this scope is closing, so attempts to create a new thread within it will throw ScopeClosing
        -- (as if the calling thread was a parent of this scope, which it should be, and we threw it a ScopeClosing
        -- ourselves).
        TVar ScopeStatus -> ScopeStatus -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar ScopeStatus
statusVar ScopeStatus
Closing
        -- Return the list of currently-running children to kill. Some of them may have *just* started (e.g. if we
        -- initially retried in `guard (n == 0)` above). That's fine - kill them all!
        TVar (IntMap ThreadId) -> STM (IntMap ThreadId)
forall a. TVar a -> STM a
readTVar TVar (IntMap ThreadId)
childrenVar

    -- Deliver a ScopeClosing exception to every running child.
    --
    -- This happens to throw in the order the children were created, but that isn't an important/useful enough feature
    -- to be worth documenting, so users shouldn't rely on it. It's definitely not the case that child 1 will completely
    -- terminate before child 2 is delivered an exception: each child may delay arbitrarily while cleaning up.
    IntMap ThreadId -> (ThreadId -> IO ()) -> IO ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ IntMap ThreadId
runningChildren \ThreadId
child -> ThreadId -> ScopeClosing -> IO ()
forall e. Exception e => ThreadId -> e -> IO ()
throwTo ThreadId
child ScopeClosing
ScopeClosing

    STM () -> IO ()
forall a. STM a -> IO a
atomically do
      -- Block until all children have terminated; this relies on children respecting the async exception, which they
      -- must, for correctness. Otherwise, a thread could indeed outlive the scope in which it's created, which is
      -- definitely not structured concurrency!
      IntMap ThreadId
children <- TVar (IntMap ThreadId) -> STM (IntMap ThreadId)
forall a. TVar a -> STM a
readTVar TVar (IntMap ThreadId)
childrenVar
      Bool -> STM ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (IntMap ThreadId -> Bool
forall a. IntMap a -> Bool
IntMap.Lazy.null IntMap ThreadId
children)

      -- Record the scope as closed (from closing), so subsequent attempts to use it will throw a runtime exception
      TVar ScopeStatus -> ScopeStatus -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar ScopeStatus
statusVar ScopeStatus
Closed

    -- By now there are three sources of exception:
    --
    --   1) A sync or async exception thrown during the callback, captured in `result`. If applicable, we want to peel
    --      the `Propagating` off of this, which was only used to indicate it came from one of our children.
    --
    --   2) A sync or async exception left for us in `childExceptionVar` by a child that tried to propagate it to us
    --      directly, but failed (because we killed it concurrently).
    --
    --   3) An async exception waiting in our exception queue, because we still have async exceptions uninterruptibly
    --      masked.
    --
    -- We cannot throw more than one, so throw them in that priority order.
    case Either SomeException a
result of
      Left SomeException
exception -> SomeException -> IO a
forall e a. Exception e => e -> IO a
throwIO (SomeException -> SomeException
peelOffPropagating SomeException
exception)
      Right a
value ->
        MVar SomeException -> IO (Maybe SomeException)
forall a. MVar a -> IO (Maybe a)
tryTakeMVar MVar SomeException
childExceptionVar IO (Maybe SomeException) -> (Maybe SomeException -> IO a) -> IO a
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
          Maybe SomeException
Nothing -> a -> IO a
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
value
          Just SomeException
exception -> SomeException -> IO a
forall e a. Exception e => e -> IO a
throwIO SomeException
exception

-- Allocate a new scope.
allocateScope :: IO Scope
allocateScope :: IO Scope
allocateScope = do
  MVar SomeException
childExceptionVar <- IO (MVar SomeException)
forall a. IO (MVar a)
newEmptyMVar
  TVar (IntMap ThreadId)
childrenVar <- IntMap ThreadId -> IO (TVar (IntMap ThreadId))
forall a. a -> IO (TVar a)
newTVarIO IntMap ThreadId
forall a. IntMap a
IntMap.Lazy.empty
  IntSupply
nextChildIdSupply <- IO IntSupply
IntSupply.new
  ThreadId
parentThreadId <- IO ThreadId
myThreadId
  TVar ScopeStatus
statusVar <- ScopeStatus -> IO (TVar ScopeStatus)
forall a. a -> IO (TVar a)
newTVarIO ScopeStatus
0
  Scope -> IO Scope
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Scope {MVar SomeException
$sel:childExceptionVar:Scope :: MVar SomeException
childExceptionVar :: MVar SomeException
childExceptionVar, TVar (IntMap ThreadId)
$sel:childrenVar:Scope :: TVar (IntMap ThreadId)
childrenVar :: TVar (IntMap ThreadId)
childrenVar, IntSupply
$sel:nextChildIdSupply:Scope :: IntSupply
nextChildIdSupply :: IntSupply
nextChildIdSupply, ThreadId
$sel:parentThreadId:Scope :: ThreadId
parentThreadId :: ThreadId
parentThreadId, TVar ScopeStatus
$sel:statusVar:Scope :: TVar ScopeStatus
statusVar :: TVar ScopeStatus
statusVar}

-- Spawn a thread in a scope, providing it its child id and a function that sets the masking state to the requested
-- masking state. The given action is called with async exceptions interruptibly masked.
spawn :: Scope -> ThreadOptions -> (Tid -> (forall x. IO x -> IO x) -> UnexceptionalIO ()) -> IO ChildIds
spawn :: Scope
-> ThreadOptions
-> (ScopeStatus -> (forall a. IO a -> IO a) -> UnexceptionalIO ())
-> IO ChildIds
spawn scope :: Scope
scope@Scope {TVar (IntMap ThreadId)
$sel:childrenVar:Scope :: Scope -> TVar (IntMap ThreadId)
childrenVar :: TVar (IntMap ThreadId)
childrenVar, TVar ScopeStatus
$sel:statusVar:Scope :: Scope -> TVar ScopeStatus
statusVar :: TVar ScopeStatus
statusVar} ThreadOptions
options ScopeStatus -> (forall a. IO a -> IO a) -> UnexceptionalIO ()
action = do
  -- Interruptible mask is enough so long as none of the STM operations below block.
  --
  -- Unconditionally set masking state to MaskedInterruptible, even though we might already be at MaskedInterruptible
  -- or MaskedUninterruptible, to avoid a branch on parentMaskingState.
  IO ChildIds -> IO ChildIds
forall a. IO a -> IO a
interruptiblyMasked do
    -- Record the thread as being about to start. Not allowed to retry.
    NonblockingSTM () -> IO ()
forall a. NonblockingSTM a -> IO a
nonblockingAtomically do
      ScopeStatus
status <- TVar ScopeStatus -> NonblockingSTM ScopeStatus
forall a. TVar a -> NonblockingSTM a
nonblockingReadTVar TVar ScopeStatus
statusVar
      Bool -> NonblockingSTM ()
forall (m :: * -> *). Applicative m => Bool -> m ()
assertM (ScopeStatus
status ScopeStatus -> ScopeStatus -> Bool
forall a. Ord a => a -> a -> Bool
>= -ScopeStatus
2)
      case ScopeStatus
status of
        ScopeStatus
Open -> TVar ScopeStatus -> ScopeStatus -> NonblockingSTM ()
forall a. TVar a -> a -> NonblockingSTM ()
nonblockingWriteTVar' TVar ScopeStatus
statusVar (ScopeStatus
status ScopeStatus -> ScopeStatus -> ScopeStatus
forall a. Num a => a -> a -> a
+ ScopeStatus
1)
        ScopeStatus
Closing -> ScopeClosing -> NonblockingSTM ()
forall e x. Exception e => e -> NonblockingSTM x
nonblockingThrowSTM ScopeClosing
ScopeClosing
        ScopeStatus
Closed -> ErrorCall -> NonblockingSTM ()
forall e x. Exception e => e -> NonblockingSTM x
nonblockingThrowSTM (String -> ErrorCall
ErrorCall String
"ki: scope closed")

    ChildIds
childIds <- Scope
-> ThreadOptions
-> (ScopeStatus -> (forall a. IO a -> IO a) -> UnexceptionalIO ())
-> IO ChildIds
spawnChild Scope
scope ThreadOptions
options ScopeStatus -> (forall a. IO a -> IO a) -> UnexceptionalIO ()
action

    -- Record the child as having started. Not allowed to retry.
    NonblockingSTM () -> IO ()
forall a. NonblockingSTM a -> IO a
nonblockingAtomically do
      ScopeStatus
starting <- TVar ScopeStatus -> NonblockingSTM ScopeStatus
forall a. TVar a -> NonblockingSTM a
nonblockingReadTVar TVar ScopeStatus
statusVar
      Bool -> NonblockingSTM ()
forall (m :: * -> *). Applicative m => Bool -> m ()
assertM (ScopeStatus
starting ScopeStatus -> ScopeStatus -> Bool
forall a. Ord a => a -> a -> Bool
>= ScopeStatus
1)
      TVar ScopeStatus -> ScopeStatus -> NonblockingSTM ()
forall a. TVar a -> a -> NonblockingSTM ()
nonblockingWriteTVar' TVar ScopeStatus
statusVar (ScopeStatus
starting ScopeStatus -> ScopeStatus -> ScopeStatus
forall a. Num a => a -> a -> a
- ScopeStatus
1)
      TVar (IntMap ThreadId) -> ChildIds -> NonblockingSTM ()
recordChild TVar (IntMap ThreadId)
childrenVar ChildIds
childIds

    ChildIds -> IO ChildIds
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ChildIds
childIds

data ChildIds
  = ChildIds
      {-# UNPACK #-} !Tid
      {-# UNPACK #-} !ThreadId

spawnChild :: Scope -> ThreadOptions -> (Tid -> (forall x. IO x -> IO x) -> UnexceptionalIO ()) -> IO ChildIds
spawnChild :: Scope
-> ThreadOptions
-> (ScopeStatus -> (forall a. IO a -> IO a) -> UnexceptionalIO ())
-> IO ChildIds
spawnChild Scope
scope ThreadOptions
options ScopeStatus -> (forall a. IO a -> IO a) -> UnexceptionalIO ()
action = do
  ScopeStatus
childId <- IntSupply -> IO ScopeStatus
IntSupply.next IntSupply
nextChildIdSupply
  ThreadId
childThreadId <-
    ThreadAffinity -> IO () -> IO ThreadId
forkWithAffinity ThreadAffinity
affinity do
      Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Bool -> Bool
not (String -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null String
label)) do
        ThreadId
childThreadId <- IO ThreadId
myThreadId
        ThreadId -> String -> IO ()
labelThread ThreadId
childThreadId String
label

      Maybe ByteCount -> (ByteCount -> IO ()) -> IO ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ Maybe ByteCount
allocationLimit \ByteCount
bytes -> do
        Int64 -> IO ()
setAllocationCounter (ByteCount -> Int64
byteCountToInt64 ByteCount
bytes)
        IO ()
enableAllocationLimit

      let -- Action that sets the masking state from the current (MaskedInterruptible) to the requested one.
          atRequestedMaskingState :: IO a -> IO a
          atRequestedMaskingState :: forall a. IO a -> IO a
atRequestedMaskingState =
            case MaskingState
requestedChildMaskingState of
              MaskingState
Unmasked -> IO a -> IO a
forall a. IO a -> IO a
unsafeUnmask
              MaskingState
MaskedInterruptible -> IO a -> IO a
forall a. a -> a
id
              MaskingState
MaskedUninterruptible -> IO a -> IO a
forall a. IO a -> IO a
uninterruptiblyMasked

      UnexceptionalIO () -> IO ()
forall a. UnexceptionalIO a -> IO a
runUnexceptionalIO (ScopeStatus -> (forall a. IO a -> IO a) -> UnexceptionalIO ()
action ScopeStatus
childId IO x -> IO x
forall a. IO a -> IO a
atRequestedMaskingState)

      NonblockingSTM () -> IO ()
forall a. NonblockingSTM a -> IO a
nonblockingAtomically (TVar (IntMap ThreadId) -> ScopeStatus -> NonblockingSTM ()
unrecordChild TVar (IntMap ThreadId)
childrenVar ScopeStatus
childId)
  ChildIds -> IO ChildIds
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ScopeStatus -> ThreadId -> ChildIds
ChildIds ScopeStatus
childId ThreadId
childThreadId)
  where
    Scope {TVar (IntMap ThreadId)
$sel:childrenVar:Scope :: Scope -> TVar (IntMap ThreadId)
childrenVar :: TVar (IntMap ThreadId)
childrenVar, IntSupply
$sel:nextChildIdSupply:Scope :: Scope -> IntSupply
nextChildIdSupply :: IntSupply
nextChildIdSupply} = Scope
scope
    ThreadOptions {ThreadAffinity
affinity :: ThreadAffinity
$sel:affinity:ThreadOptions :: ThreadOptions -> ThreadAffinity
affinity, Maybe ByteCount
allocationLimit :: Maybe ByteCount
$sel:allocationLimit:ThreadOptions :: ThreadOptions -> Maybe ByteCount
allocationLimit, String
label :: String
$sel:label:ThreadOptions :: ThreadOptions -> String
label, $sel:maskingState:ThreadOptions :: ThreadOptions -> MaskingState
maskingState = MaskingState
requestedChildMaskingState} = ThreadOptions
options
{-# INLINE spawnChild #-}

-- Record our child by either:
--
--   * Flipping `Nothing` to `Just childThreadId` (common case: we record child before it unrecords itself)
--   * Flipping `Just _` to `Nothing` (uncommon case: we observe that a child already unrecorded itself)
recordChild :: TVar (IntMap ThreadId) -> ChildIds -> NonblockingSTM ()
recordChild :: TVar (IntMap ThreadId) -> ChildIds -> NonblockingSTM ()
recordChild TVar (IntMap ThreadId)
childrenVar (ChildIds ScopeStatus
childId ThreadId
childThreadId) = do
  IntMap ThreadId
children <- TVar (IntMap ThreadId) -> NonblockingSTM (IntMap ThreadId)
forall a. TVar a -> NonblockingSTM a
nonblockingReadTVar TVar (IntMap ThreadId)
childrenVar
  TVar (IntMap ThreadId) -> IntMap ThreadId -> NonblockingSTM ()
forall a. TVar a -> a -> NonblockingSTM ()
nonblockingWriteTVar' TVar (IntMap ThreadId)
childrenVar ((Maybe ThreadId -> Maybe ThreadId)
-> ScopeStatus -> IntMap ThreadId -> IntMap ThreadId
forall a.
(Maybe a -> Maybe a) -> ScopeStatus -> IntMap a -> IntMap a
IntMap.Lazy.alter (Maybe ThreadId
-> (ThreadId -> Maybe ThreadId) -> Maybe ThreadId -> Maybe ThreadId
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (ThreadId -> Maybe ThreadId
forall a. a -> Maybe a
Just ThreadId
childThreadId) (Maybe ThreadId -> ThreadId -> Maybe ThreadId
forall a b. a -> b -> a
const Maybe ThreadId
forall a. Maybe a
Nothing)) ScopeStatus
childId IntMap ThreadId
children)

-- Unrecord a child (ourselves) by either:
--
--   * Flipping `Just childThreadId` to `Nothing` (common case: parent recorded us first)
--   * Flipping `Nothing` to `Just undefined` (uncommon case: we terminate and unrecord before parent can record us).
unrecordChild :: TVar (IntMap ThreadId) -> Tid -> NonblockingSTM ()
unrecordChild :: TVar (IntMap ThreadId) -> ScopeStatus -> NonblockingSTM ()
unrecordChild TVar (IntMap ThreadId)
childrenVar ScopeStatus
childId = do
  IntMap ThreadId
children <- TVar (IntMap ThreadId) -> NonblockingSTM (IntMap ThreadId)
forall a. TVar a -> NonblockingSTM a
nonblockingReadTVar TVar (IntMap ThreadId)
childrenVar
  TVar (IntMap ThreadId) -> IntMap ThreadId -> NonblockingSTM ()
forall a. TVar a -> a -> NonblockingSTM ()
nonblockingWriteTVar' TVar (IntMap ThreadId)
childrenVar ((Maybe ThreadId -> Maybe ThreadId)
-> ScopeStatus -> IntMap ThreadId -> IntMap ThreadId
forall a.
(Maybe a -> Maybe a) -> ScopeStatus -> IntMap a -> IntMap a
IntMap.Lazy.alter (Maybe ThreadId
-> (ThreadId -> Maybe ThreadId) -> Maybe ThreadId -> Maybe ThreadId
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (ThreadId -> Maybe ThreadId
forall a. a -> Maybe a
Just ThreadId
forall a. HasCallStack => a
undefined) (Maybe ThreadId -> ThreadId -> Maybe ThreadId
forall a b. a -> b -> a
const Maybe ThreadId
forall a. Maybe a
Nothing)) ScopeStatus
childId IntMap ThreadId
children)

-- | Wait until all threads created within a scope terminate.
awaitAll :: Scope -> STM ()
awaitAll :: Scope -> STM ()
awaitAll Scope {TVar (IntMap ThreadId)
$sel:childrenVar:Scope :: Scope -> TVar (IntMap ThreadId)
childrenVar :: TVar (IntMap ThreadId)
childrenVar, TVar ScopeStatus
$sel:statusVar:Scope :: Scope -> TVar ScopeStatus
statusVar :: TVar ScopeStatus
statusVar} = do
  IntMap ThreadId
children <- TVar (IntMap ThreadId) -> STM (IntMap ThreadId)
forall a. TVar a -> STM a
readTVar TVar (IntMap ThreadId)
childrenVar
  Bool -> STM ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (IntMap ThreadId -> Bool
forall a. IntMap a -> Bool
IntMap.Lazy.null IntMap ThreadId
children)
  ScopeStatus
status <- TVar ScopeStatus -> STM ScopeStatus
forall a. TVar a -> STM a
readTVar TVar ScopeStatus
statusVar
  Bool -> STM ()
forall (m :: * -> *). Applicative m => Bool -> m ()
assertM (ScopeStatus
status ScopeStatus -> ScopeStatus -> Bool
forall a. Ord a => a -> a -> Bool
>= -ScopeStatus
2)
  case ScopeStatus
status of
    ScopeStatus
Open -> Bool -> STM ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (ScopeStatus
status ScopeStatus -> ScopeStatus -> Bool
forall a. Eq a => a -> a -> Bool
== ScopeStatus
0)
    ScopeStatus
Closing -> STM ()
forall a. STM a
retry -- block until closed
    ScopeStatus
Closed -> () -> STM ()
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

-- | Create a child thread to execute an action within a scope.
--
-- /Note/: The child thread does not mask asynchronous exceptions, regardless of the parent thread's masking state. To
-- create a child thread with a different initial masking state, use 'Ki.forkWith'.
fork :: Scope -> IO a -> IO (Thread a)
fork :: forall a. Scope -> IO a -> IO (Thread a)
fork Scope
scope =
  Scope -> ThreadOptions -> IO a -> IO (Thread a)
forall a. Scope -> ThreadOptions -> IO a -> IO (Thread a)
forkWith Scope
scope ThreadOptions
defaultThreadOptions

-- | Variant of 'Ki.fork' for threads that never return.
fork_ :: Scope -> IO Void -> IO ()
fork_ :: Scope -> IO Void -> IO ()
fork_ Scope
scope =
  Scope -> ThreadOptions -> IO Void -> IO ()
forkWith_ Scope
scope ThreadOptions
defaultThreadOptions

-- | Variant of 'Ki.fork' that takes an additional options argument.
forkWith :: Scope -> ThreadOptions -> IO a -> IO (Thread a)
forkWith :: forall a. Scope -> ThreadOptions -> IO a -> IO (Thread a)
forkWith Scope
scope ThreadOptions
opts IO a
action = do
  TVar (Result a)
resultVar <- Result a -> IO (TVar (Result a))
forall a. a -> IO (TVar a)
newTVarIO Result a
forall a. Result a
NoResultYet
  let done :: Result a -> UnexceptionalIO ()
done Result a
result = IO () -> UnexceptionalIO ()
forall a. IO a -> UnexceptionalIO a
UnexceptionalIO (STM () -> IO ()
forall a. STM a -> IO a
atomically (TVar (Result a) -> Result a -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar (Result a)
resultVar Result a
result))
  ChildIds ScopeStatus
_ ThreadId
childThreadId <-
    Scope
-> ThreadOptions
-> (ScopeStatus -> (forall a. IO a -> IO a) -> UnexceptionalIO ())
-> IO ChildIds
spawn Scope
scope ThreadOptions
opts \ScopeStatus
childId forall a. IO a -> IO a
masking -> do
      IO a -> UnexceptionalIO (IOResult a)
forall a. IO a -> UnexceptionalIO (IOResult a)
unexceptionalTry (IO a -> IO a
forall a. IO a -> IO a
masking IO a
action) UnexceptionalIO (IOResult a)
-> (IOResult a -> UnexceptionalIO ()) -> UnexceptionalIO ()
forall a b.
UnexceptionalIO a -> (a -> UnexceptionalIO b) -> UnexceptionalIO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
        Failure SomeException
exception -> do
          Bool -> UnexceptionalIO () -> UnexceptionalIO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Bool -> Bool
not (forall e. Exception e => SomeException -> Bool
exceptionIs @ScopeClosing SomeException
exception)) do
            Scope -> ScopeStatus -> SomeException -> UnexceptionalIO ()
propagateException Scope
scope ScopeStatus
childId SomeException
exception
          -- even put async exceptions that we propagated. this isn't totally ideal because a caller awaiting this
          -- thread would not be able to distinguish between async exceptions delivered to this thread, or itself
          Result a -> UnexceptionalIO ()
done (SomeException -> Result a
forall a. SomeException -> Result a
BadResult SomeException
exception)
        Success a
value -> Result a -> UnexceptionalIO ()
done (a -> Result a
forall a. a -> Result a
GoodResult a
value)
  let doAwait :: STM a
doAwait =
        TVar (Result a) -> STM (Result a)
forall a. TVar a -> STM a
readTVar TVar (Result a)
resultVar STM (Result a) -> (Result a -> STM a) -> STM a
forall a b. STM a -> (a -> STM b) -> STM b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
          Result a
NoResultYet -> STM a
forall a. STM a
retry
          BadResult SomeException
exception -> SomeException -> STM a
forall e a. Exception e => e -> STM a
throwSTM SomeException
exception
          GoodResult a
value -> a -> STM a
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
value
  Thread a -> IO (Thread a)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ThreadId -> STM a -> Thread a
forall a. ThreadId -> STM a -> Thread a
makeThread ThreadId
childThreadId STM a
doAwait)

-- | Variant of 'Ki.forkWith' for threads that never return.
forkWith_ :: Scope -> ThreadOptions -> IO Void -> IO ()
forkWith_ :: Scope -> ThreadOptions -> IO Void -> IO ()
forkWith_ Scope
scope ThreadOptions
opts IO Void
action = do
  ChildIds
_childThreadId <-
    Scope
-> ThreadOptions
-> (ScopeStatus -> (forall a. IO a -> IO a) -> UnexceptionalIO ())
-> IO ChildIds
spawn Scope
scope ThreadOptions
opts \ScopeStatus
childId forall a. IO a -> IO a
masking ->
      (SomeException -> UnexceptionalIO ())
-> (Void -> UnexceptionalIO ()) -> IO Void -> UnexceptionalIO ()
forall a b.
(SomeException -> UnexceptionalIO b)
-> (a -> UnexceptionalIO b) -> IO a -> UnexceptionalIO b
unexceptionalTryEither
        ( \SomeException
exception ->
            Bool -> UnexceptionalIO () -> UnexceptionalIO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Bool -> Bool
not (forall e. Exception e => SomeException -> Bool
exceptionIs @ScopeClosing SomeException
exception)) do
              Scope -> ScopeStatus -> SomeException -> UnexceptionalIO ()
propagateException Scope
scope ScopeStatus
childId SomeException
exception
        )
        Void -> UnexceptionalIO ()
forall a. Void -> a
absurd
        (IO Void -> IO Void
forall a. IO a -> IO a
masking IO Void
action)
  () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

-- | Like 'Ki.fork', but the child thread does not propagate exceptions that are both:
--
-- * Synchronous (/i.e./ not an instance of 'SomeAsyncException').
-- * An instance of @e@.
forkTry :: forall e a. (Exception e) => Scope -> IO a -> IO (Thread (Either e a))
forkTry :: forall e a.
Exception e =>
Scope -> IO a -> IO (Thread (Either e a))
forkTry Scope
scope =
  Scope -> ThreadOptions -> IO a -> IO (Thread (Either e a))
forall e a.
Exception e =>
Scope -> ThreadOptions -> IO a -> IO (Thread (Either e a))
forkTryWith Scope
scope ThreadOptions
defaultThreadOptions

data Result a
  = NoResultYet
  | BadResult !SomeException -- sync or async
  | GoodResult a

-- | Variant of 'Ki.forkTry' that takes an additional options argument.
forkTryWith :: forall e a. (Exception e) => Scope -> ThreadOptions -> IO a -> IO (Thread (Either e a))
forkTryWith :: forall e a.
Exception e =>
Scope -> ThreadOptions -> IO a -> IO (Thread (Either e a))
forkTryWith Scope
scope ThreadOptions
opts IO a
action = do
  TVar (Result a)
resultVar <- Result a -> IO (TVar (Result a))
forall a. a -> IO (TVar a)
newTVarIO Result a
forall a. Result a
NoResultYet
  let done :: Result a -> UnexceptionalIO ()
done Result a
result = IO () -> UnexceptionalIO ()
forall a. IO a -> UnexceptionalIO a
UnexceptionalIO (STM () -> IO ()
forall a. STM a -> IO a
atomically (TVar (Result a) -> Result a -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar (Result a)
resultVar Result a
result))
  ChildIds ScopeStatus
_ ThreadId
childThreadId <-
    Scope
-> ThreadOptions
-> (ScopeStatus -> (forall a. IO a -> IO a) -> UnexceptionalIO ())
-> IO ChildIds
spawn Scope
scope ThreadOptions
opts \ScopeStatus
childId forall a. IO a -> IO a
masking -> do
      IOResult a
result <- IO a -> UnexceptionalIO (IOResult a)
forall a. IO a -> UnexceptionalIO (IOResult a)
unexceptionalTry (IO a -> IO a
forall a. IO a -> IO a
masking IO a
action)
      case IOResult a
result of
        Failure SomeException
exception -> do
          -- then-branch explanation: if the user calls `forkTry @MyAsyncException` for some reason, we want to ignore
          -- this request and propagate the async exception. `forkTry` can only be used to catch synchronous exceptions.
          let shouldPropagate :: Bool
shouldPropagate =
                if forall e. Exception e => SomeException -> Bool
exceptionIs @e SomeException
exception
                  then forall e. Exception e => SomeException -> Bool
exceptionIs @SomeAsyncException SomeException
exception
                  else Bool -> Bool
not (forall e. Exception e => SomeException -> Bool
exceptionIs @ScopeClosing SomeException
exception)
          Bool -> UnexceptionalIO () -> UnexceptionalIO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
shouldPropagate (Scope -> ScopeStatus -> SomeException -> UnexceptionalIO ()
propagateException Scope
scope ScopeStatus
childId SomeException
exception)
          Result a -> UnexceptionalIO ()
done (SomeException -> Result a
forall a. SomeException -> Result a
BadResult SomeException
exception)
        Success a
value -> Result a -> UnexceptionalIO ()
done (a -> Result a
forall a. a -> Result a
GoodResult a
value)
  let doAwait :: STM (Either e a)
doAwait =
        TVar (Result a) -> STM (Result a)
forall a. TVar a -> STM a
readTVar TVar (Result a)
resultVar STM (Result a)
-> (Result a -> STM (Either e a)) -> STM (Either e a)
forall a b. STM a -> (a -> STM b) -> STM b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
          Result a
NoResultYet -> STM (Either e a)
forall a. STM a
retry
          BadResult SomeException
exception ->
            case forall e. Exception e => SomeException -> Maybe e
fromException @e SomeException
exception of
              Maybe e
Nothing -> SomeException -> STM (Either e a)
forall e a. Exception e => e -> STM a
throwSTM SomeException
exception
              Just e
expectedException -> Either e a -> STM (Either e a)
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (e -> Either e a
forall a b. a -> Either a b
Left e
expectedException)
          GoodResult a
value -> Either e a -> STM (Either e a)
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> Either e a
forall a b. b -> Either a b
Right a
value)
  Thread (Either e a) -> IO (Thread (Either e a))
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ThreadId -> STM (Either e a) -> Thread (Either e a)
forall a. ThreadId -> STM a -> Thread a
makeThread ThreadId
childThreadId STM (Either e a)
doAwait)

-- We have a non-`ScopeClosing` exception to propagate to our parent.
--
-- If our scope has already begun closing (`statusVar` is Closing), then either...
--
--   (A) We already received a `ScopeClosing`, but then ended up trying to propagate an exception anyway, because we
--   threw a synchronous exception (or were hit by a different asynchronous exception) during our teardown procedure.
--
--   or
--
--   (B) We will receive a `ScopeClosing` imminently, because our parent has *just* finished setting `statusVar` to
--   Closing, and will proceed to throw ScopeClosing to all of its children.
--
-- If (A), our parent has asynchronous exceptions masked, so we must inform it of our exception via `childExceptionVar`
-- rather than throwTo. If (B), either mechanism would work. And because we don't if we're in case (A) or (B), we just
-- `childExceptionVar`.
--
-- And if our scope has not already begun closing (`statusVar` is not Closing), then we ought to throw our exception to
-- it. But that might fail due to either...
--
--   (C) Our parent concurrently closing the scope and sending us a `ScopeClosing`; because it has asynchronous
--   exceptions uninterruptibly masked and we only have asynchronous exception *synchronously* masked, its `throwTo`
--   will return `()`, and ours will throw that `ScopeClosing` asynchronous exception. In this case, since we now know
--   our parent is tearing down and has asynchronous exceptions masked, we again inform it via `childExceptionVar`.
--
--   (D) Some *other* non-`ScopeClosing` asynchronous exception is raised here. This is truly odd: maybe it's a heap
--   overflow exception from the GHC runtime? Maybe some other thread has smuggled our `ThreadId` out and has manually
--   thrown us an exception for some reason? Either way, because we already have an exception that we are trying to
--   propagate, we just scoot these freaky exceptions under the rug.
--
-- Precondition: interruptibly masked
propagateException :: Scope -> Tid -> SomeException -> UnexceptionalIO ()
propagateException :: Scope -> ScopeStatus -> SomeException -> UnexceptionalIO ()
propagateException Scope {MVar SomeException
$sel:childExceptionVar:Scope :: Scope -> MVar SomeException
childExceptionVar :: MVar SomeException
childExceptionVar, ThreadId
$sel:parentThreadId:Scope :: Scope -> ThreadId
parentThreadId :: ThreadId
parentThreadId, TVar ScopeStatus
$sel:statusVar:Scope :: Scope -> TVar ScopeStatus
statusVar :: TVar ScopeStatus
statusVar} ScopeStatus
childId SomeException
exception =
  IO ScopeStatus -> UnexceptionalIO ScopeStatus
forall a. IO a -> UnexceptionalIO a
UnexceptionalIO (TVar ScopeStatus -> IO ScopeStatus
forall a. TVar a -> IO a
readTVarIO TVar ScopeStatus
statusVar) UnexceptionalIO ScopeStatus
-> (ScopeStatus -> UnexceptionalIO ()) -> UnexceptionalIO ()
forall a b.
UnexceptionalIO a -> (a -> UnexceptionalIO b) -> UnexceptionalIO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    ScopeStatus
Closing -> UnexceptionalIO ()
tryPutChildExceptionVar -- (A) or (B), we don't care which
    ScopeStatus
status -> do
      Bool -> UnexceptionalIO ()
forall (m :: * -> *). Applicative m => Bool -> m ()
assertM (ScopeStatus
status ScopeStatus -> ScopeStatus -> Bool
forall a. Ord a => a -> a -> Bool
>= ScopeStatus
0) -- we know status is Open (0+) here; can't be Closed (-2)
      UnexceptionalIO ()
loop
  where
    loop :: UnexceptionalIO ()
    loop :: UnexceptionalIO ()
loop =
      IO () -> UnexceptionalIO (IOResult ())
forall a. IO a -> UnexceptionalIO (IOResult a)
unexceptionalTry (SomeException -> ScopeStatus -> ThreadId -> IO ()
propagate SomeException
exception ScopeStatus
childId ThreadId
parentThreadId) UnexceptionalIO (IOResult ())
-> (IOResult () -> UnexceptionalIO ()) -> UnexceptionalIO ()
forall a b.
UnexceptionalIO a -> (a -> UnexceptionalIO b) -> UnexceptionalIO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
        Failure SomeException
secondException
          | forall e. Exception e => SomeException -> Bool
exceptionIs @ScopeClosing SomeException
secondException -> UnexceptionalIO ()
tryPutChildExceptionVar -- (C)
          | Bool
otherwise -> UnexceptionalIO ()
loop -- (D)
        Success ()
_ -> () -> UnexceptionalIO ()
forall a. a -> UnexceptionalIO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

    tryPutChildExceptionVar :: UnexceptionalIO ()
    tryPutChildExceptionVar :: UnexceptionalIO ()
tryPutChildExceptionVar =
      IO () -> UnexceptionalIO ()
forall a. IO a -> UnexceptionalIO a
UnexceptionalIO (IO Bool -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (MVar SomeException -> SomeException -> IO Bool
forall a. MVar a -> a -> IO Bool
tryPutMVar MVar SomeException
childExceptionVar SomeException
exception))