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 (..),
    SomeException,
    assert,
    asyncExceptionFromException,
    asyncExceptionToException,
    throwIO,
    try,
    uninterruptibleMask,
    pattern ErrorCall,
  )
import Control.Monad (when, guard)
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 Ki.Internal.ByteCount
import Ki.Internal.Counter
import Ki.Internal.IO
  ( IOResult (..),
    UnexceptionalIO (..),
    interruptiblyMasked,
    isAsyncException,
    unexceptionalTry,
    unexceptionalTryEither,
    uninterruptiblyMasked,
  )
import Ki.Internal.Thread
import Data.Maybe (isJust)

-- | 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 counter that holds the (int) key to use for the next child thread.
    Scope -> Counter
nextChildIdCounter :: {-# UNPACK #-} !Counter,
    -- 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 living 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.
data ScopeClosing
  = ScopeClosing

instance Show ScopeClosing where
  show :: ScopeClosing -> String
show ScopeClosing
_ = String
"ScopeClosing"

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

-- Trust without verifying that any 'ScopeClosed' 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?
isScopeClosingException :: SomeException -> Bool
isScopeClosingException :: SomeException -> Bool
isScopeClosingException SomeException
exception =
  Maybe ScopeClosing -> Bool
forall a. Maybe a -> Bool
isJust (forall e. Exception e => SomeException -> Maybe e
fromException @ScopeClosing SomeException
exception)

pattern IsScopeClosingException :: SomeException
pattern $mIsScopeClosingException :: forall {r}. SomeException -> ((# #) -> r) -> ((# #) -> r) -> r
IsScopeClosingException <- (isScopeClosingException -> True)

-- | 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
livingChildren <- do
      IntMap ThreadId
livingChildren0 <-
        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
n <- TVar ScopeStatus -> STM ScopeStatus
forall a. TVar a -> STM a
readTVar TVar ScopeStatus
statusVar
          Bool -> STM () -> STM ()
forall a. HasCallStack => Bool -> a -> a
assert (ScopeStatus
n ScopeStatus -> ScopeStatus -> Bool
forall a. Ord a => a -> a -> Bool
>= ScopeStatus
0) (Bool -> STM ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (ScopeStatus
n 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

      -- If one of our children propagated an exception to us, then we know it's about to terminate, so we don't bother
      -- throwing an exception to it.
      IntMap ThreadId -> IO (IntMap ThreadId)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure case Either SomeException a
result of
        Left (SomeException -> Maybe ThreadFailed
forall e. Exception e => SomeException -> Maybe e
fromException -> Just ThreadFailed {ScopeStatus
childId :: ScopeStatus
$sel:childId:ThreadFailed :: ThreadFailed -> ScopeStatus
childId}) -> ScopeStatus -> IntMap ThreadId -> IntMap ThreadId
forall a. ScopeStatus -> IntMap a -> IntMap a
IntMap.Lazy.delete ScopeStatus
childId IntMap ThreadId
livingChildren0
        Either SomeException a
_ -> IntMap ThreadId
livingChildren0

    -- Deliver a ScopeClosing exception to every living child.
    --
    -- This happens to throw in the order the children were created... but I think we decided this feature isn't very
    -- useful in practice, so maybe we should simplify the internals and just keep a set of children?
    [ThreadId] -> (ThreadId -> IO ()) -> IO ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ (IntMap ThreadId -> [ThreadId]
forall a. IntMap a -> [a]
IntMap.Lazy.elems IntMap ThreadId
livingChildren) \ThreadId
livingChild -> ThreadId -> ScopeClosing -> IO ()
forall e. Exception e => ThreadId -> e -> IO ()
throwTo ThreadId
livingChild 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!
      TVar (IntMap ThreadId) -> STM ()
forall a. TVar (IntMap a) -> STM ()
blockUntilEmpty TVar (IntMap ThreadId)
childrenVar
      -- 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 unwrap
    --      the `ThreadFailed` 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
unwrapThreadFailed 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
  Counter
nextChildIdCounter <- IO Counter
newCounter
  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, Counter
$sel:nextChildIdCounter:Scope :: Counter
nextChildIdCounter :: Counter
nextChildIdCounter, 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 ThreadId
spawn :: Scope
-> ThreadOptions
-> (ScopeStatus -> (forall a. IO a -> IO a) -> UnexceptionalIO ())
-> IO ThreadId
spawn
  Scope {TVar (IntMap ThreadId)
$sel:childrenVar:Scope :: Scope -> TVar (IntMap ThreadId)
childrenVar :: TVar (IntMap ThreadId)
childrenVar, Counter
$sel:nextChildIdCounter:Scope :: Scope -> Counter
nextChildIdCounter :: Counter
nextChildIdCounter, TVar ScopeStatus
$sel:statusVar:Scope :: Scope -> TVar ScopeStatus
statusVar :: TVar ScopeStatus
statusVar}
  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}
  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 ThreadId -> IO ThreadId
forall a. IO a -> IO a
interruptiblyMasked do
      -- Record the thread as being about to start. Not allowed to retry.
      STM () -> IO ()
forall a. STM a -> IO a
atomically do
        ScopeStatus
n <- TVar ScopeStatus -> STM ScopeStatus
forall a. TVar a -> STM a
readTVar TVar ScopeStatus
statusVar
        Bool -> STM () -> STM ()
forall a. HasCallStack => Bool -> a -> a
assert (ScopeStatus
n ScopeStatus -> ScopeStatus -> Bool
forall a. Ord a => a -> a -> Bool
>= -ScopeStatus
2) do
          case ScopeStatus
n of
            ScopeStatus
Open -> TVar ScopeStatus -> ScopeStatus -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar ScopeStatus
statusVar (ScopeStatus -> STM ()) -> ScopeStatus -> STM ()
forall a b. (a -> b) -> a -> b
$! ScopeStatus
n ScopeStatus -> ScopeStatus -> ScopeStatus
forall a. Num a => a -> a -> a
+ ScopeStatus
1
            ScopeStatus
Closing -> ScopeClosing -> STM ()
forall e a. Exception e => e -> STM a
throwSTM ScopeClosing
ScopeClosing
            ScopeStatus
Closed -> ErrorCall -> STM ()
forall e a. Exception e => e -> STM a
throwSTM (String -> ErrorCall
ErrorCall String
"ki: scope closed")

      ScopeStatus
childId <- Counter -> IO ScopeStatus
incrCounter Counter
nextChildIdCounter

      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

          case Maybe ByteCount
allocationLimit of
            Maybe ByteCount
Nothing -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
            Just 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)

          STM () -> IO ()
forall a. STM a -> IO a
atomically (TVar (IntMap ThreadId) -> ScopeStatus -> STM ()
unrecordChild TVar (IntMap ThreadId)
childrenVar ScopeStatus
childId)

      -- Record the child as having started. Not allowed to retry.
      STM () -> IO ()
forall a. STM a -> IO a
atomically do
        ScopeStatus
n <- TVar ScopeStatus -> STM ScopeStatus
forall a. TVar a -> STM a
readTVar TVar ScopeStatus
statusVar
        TVar ScopeStatus -> ScopeStatus -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar ScopeStatus
statusVar (ScopeStatus -> STM ()) -> ScopeStatus -> STM ()
forall a b. (a -> b) -> a -> b
$! ScopeStatus
n ScopeStatus -> ScopeStatus -> ScopeStatus
forall a. Num a => a -> a -> a
- ScopeStatus
1
        TVar (IntMap ThreadId) -> ScopeStatus -> ThreadId -> STM ()
recordChild TVar (IntMap ThreadId)
childrenVar ScopeStatus
childId ThreadId
childThreadId

      ThreadId -> IO ThreadId
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ThreadId
childThreadId

-- 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)
--
-- Never retries.
recordChild :: TVar (IntMap ThreadId) -> Tid -> ThreadId -> STM ()
recordChild :: TVar (IntMap ThreadId) -> ScopeStatus -> ThreadId -> STM ()
recordChild TVar (IntMap ThreadId)
childrenVar ScopeStatus
childId ThreadId
childThreadId = do
  IntMap ThreadId
children <- TVar (IntMap ThreadId) -> STM (IntMap ThreadId)
forall a. TVar a -> STM a
readTVar TVar (IntMap ThreadId)
childrenVar
  TVar (IntMap ThreadId) -> IntMap ThreadId -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar (IntMap ThreadId)
childrenVar (IntMap ThreadId -> STM ()) -> IntMap ThreadId -> STM ()
forall a b. (a -> b) -> a -> b
$! (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).
--
-- Never retries.
unrecordChild :: TVar (IntMap ThreadId) -> Tid -> STM ()
unrecordChild :: TVar (IntMap ThreadId) -> ScopeStatus -> STM ()
unrecordChild TVar (IntMap ThreadId)
childrenVar ScopeStatus
childId = do
  IntMap ThreadId
children <- TVar (IntMap ThreadId) -> STM (IntMap ThreadId)
forall a. TVar a -> STM a
readTVar TVar (IntMap ThreadId)
childrenVar
  TVar (IntMap ThreadId) -> IntMap ThreadId -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar (IntMap ThreadId)
childrenVar (IntMap ThreadId -> STM ()) -> IntMap ThreadId -> STM ()
forall a b. (a -> b) -> a -> b
$! (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
  TVar (IntMap ThreadId) -> STM ()
forall a. TVar (IntMap a) -> STM ()
blockUntilEmpty TVar (IntMap ThreadId)
childrenVar
  ScopeStatus
n <- TVar ScopeStatus -> STM ScopeStatus
forall a. TVar a -> STM a
readTVar TVar ScopeStatus
statusVar
  case ScopeStatus
n of
    ScopeStatus
Open -> Bool -> STM ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (ScopeStatus
n 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 ()

-- Block until an IntMap becomes empty.
blockUntilEmpty :: TVar (IntMap a) -> STM ()
blockUntilEmpty :: forall a. TVar (IntMap a) -> STM ()
blockUntilEmpty TVar (IntMap a)
var = do
  IntMap a
x <- TVar (IntMap a) -> STM (IntMap a)
forall a. TVar a -> STM a
readTVar TVar (IntMap a)
var
  Bool -> STM ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (IntMap a -> Bool
forall a. IntMap a -> Bool
IntMap.Lazy.null IntMap a
x)

-- | 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))
  ThreadId
ident <-
    Scope
-> ThreadOptions
-> (ScopeStatus -> (forall a. IO a -> IO a) -> UnexceptionalIO ())
-> IO ThreadId
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
          Bool -> UnexceptionalIO () -> UnexceptionalIO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when
            (Bool -> Bool
not (SomeException -> Bool
isScopeClosingException SomeException
exception))
            (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
ident 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
  ThreadId
_childThreadId <-
    Scope
-> ThreadOptions
-> (ScopeStatus -> (forall a. IO a -> IO a) -> UnexceptionalIO ())
-> IO ThreadId
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 (SomeException -> Bool
isScopeClosingException SomeException
exception)) (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))
  ThreadId
childThreadId <-
    Scope
-> ThreadOptions
-> (ScopeStatus -> (forall a. IO a -> IO a) -> UnexceptionalIO ())
-> IO ThreadId
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
          let shouldPropagate :: Bool
shouldPropagate =
                if SomeException -> Bool
isScopeClosingException SomeException
exception
                  then Bool
False
                  else case forall e. Exception e => SomeException -> Maybe e
fromException @e SomeException
exception of
                    Maybe e
Nothing -> Bool
True
                    -- if the user calls `forkTry @MyAsyncException`, we still want to propagate the async exception
                    Just e
_ -> SomeException -> Bool
isAsyncException 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) / (B)
    ScopeStatus
_ -> UnexceptionalIO ()
loop -- we know status is Open here
  where
    loop :: UnexceptionalIO ()
    loop :: UnexceptionalIO ()
loop =
      IO () -> UnexceptionalIO (IOResult ())
forall a. IO a -> UnexceptionalIO (IOResult a)
unexceptionalTry (ThreadId -> ThreadFailed -> IO ()
forall e. Exception e => ThreadId -> e -> IO ()
throwTo ThreadId
parentThreadId ThreadFailed {ScopeStatus
$sel:childId:ThreadFailed :: ScopeStatus
childId :: ScopeStatus
childId, SomeException
exception :: SomeException
$sel:exception:ThreadFailed :: SomeException
exception}) 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
IsScopeClosingException -> UnexceptionalIO ()
tryPutChildExceptionVar -- (C)
        Failure SomeException
_ -> 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))