module Ki.Internal.Thread
  ( Thread,
    makeThread,
    await,
    ThreadAffinity (..),
    ThreadOptions (..),
    defaultThreadOptions,
    ThreadFailed (..),
    unwrapThreadFailed,
  )
where

import Control.Exception
  ( BlockedIndefinitelyOnSTM (..),
    Exception (fromException, toException),
    MaskingState (..),
    asyncExceptionFromException,
    asyncExceptionToException,
  )
import GHC.Conc (STM, catchSTM)
import Ki.Internal.ByteCount
import Ki.Internal.Prelude

-- | A thread.
--
-- ==== __👉 Details__
--
-- * A thread's lifetime is delimited by the scope in which it was created.
--
-- * The thread that creates a scope is considered the parent of all threads created within it.
--
-- * If an exception is raised in a child thread, the child either propagates the exception to its parent (see
--   'Ki.fork'), or returns the exception as a value (see 'Ki.forkTry').
--
-- * All threads created within a scope are terminated when the scope closes.
data Thread a = Thread
  { forall a. Thread a -> ThreadId
threadId :: {-# UNPACK #-} !ThreadId,
    forall a. Thread a -> STM a
await_ :: !(STM a)
  }
  deriving stock ((forall a b. (a -> b) -> Thread a -> Thread b)
-> (forall a b. a -> Thread b -> Thread a) -> Functor Thread
forall a b. a -> Thread b -> Thread a
forall a b. (a -> b) -> Thread a -> Thread b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: forall a b. a -> Thread b -> Thread a
$c<$ :: forall a b. a -> Thread b -> Thread a
fmap :: forall a b. (a -> b) -> Thread a -> Thread b
$cfmap :: forall a b. (a -> b) -> Thread a -> Thread b
Functor)

instance Eq (Thread a) where
  Thread ThreadId
ix STM a
_ == :: Thread a -> Thread a -> Bool
== Thread ThreadId
iy STM a
_ =
    ThreadId
ix ThreadId -> ThreadId -> Bool
forall a. Eq a => a -> a -> Bool
== ThreadId
iy

instance Ord (Thread a) where
  compare :: Thread a -> Thread a -> Ordering
compare (Thread ThreadId
ix STM a
_) (Thread ThreadId
iy STM a
_) =
    ThreadId -> ThreadId -> Ordering
forall a. Ord a => a -> a -> Ordering
compare ThreadId
ix ThreadId
iy

makeThread :: ThreadId -> STM a -> Thread a
makeThread :: forall a. ThreadId -> STM a -> Thread a
makeThread ThreadId
threadId STM a
action =
  Thread
    { ThreadId
threadId :: ThreadId
$sel:threadId:Thread :: ThreadId
threadId,
      -- If *they* are deadlocked, we will *both* will be delivered a wakeup from the RTS. We want to shrug this
      -- exception off, because afterwards they'll have put to the result var. But don't shield indefinitely, once will
      -- cover this use case and prevent any accidental infinite loops.
      $sel:await_:Thread :: STM a
await_ = (BlockedIndefinitelyOnSTM -> STM a)
-> (a -> STM a) -> STM a -> STM a
forall e b a.
Exception e =>
(e -> STM b) -> (a -> STM b) -> STM a -> STM b
tryEitherSTM (\BlockedIndefinitelyOnSTM
BlockedIndefinitelyOnSTM -> STM a
action) a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure STM a
action
    }

-- | What, if anything, a thread is bound to.
data ThreadAffinity
  = -- | Unbound.
    Unbound
  | -- | Bound to a capability.
    Capability Int
  | -- | Bound to an OS thread.
    OsThread
  deriving stock (ThreadAffinity -> ThreadAffinity -> Bool
(ThreadAffinity -> ThreadAffinity -> Bool)
-> (ThreadAffinity -> ThreadAffinity -> Bool) -> Eq ThreadAffinity
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: ThreadAffinity -> ThreadAffinity -> Bool
$c/= :: ThreadAffinity -> ThreadAffinity -> Bool
== :: ThreadAffinity -> ThreadAffinity -> Bool
$c== :: ThreadAffinity -> ThreadAffinity -> Bool
Eq, Int -> ThreadAffinity -> ShowS
[ThreadAffinity] -> ShowS
ThreadAffinity -> String
(Int -> ThreadAffinity -> ShowS)
-> (ThreadAffinity -> String)
-> ([ThreadAffinity] -> ShowS)
-> Show ThreadAffinity
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [ThreadAffinity] -> ShowS
$cshowList :: [ThreadAffinity] -> ShowS
show :: ThreadAffinity -> String
$cshow :: ThreadAffinity -> String
showsPrec :: Int -> ThreadAffinity -> ShowS
$cshowsPrec :: Int -> ThreadAffinity -> ShowS
Show)

-- |
--
-- [@affinity@]:
--
--     The affinity of a thread. A thread can be unbound, bound to a specific capability, or bound to a specific OS
--     thread.
--
--     Default: 'Unbound'
--
-- [@allocationLimit@]:
--
--     The maximum number of bytes a thread may allocate before it is delivered an
--     'Control.Exception.AllocationLimitExceeded' exception. If caught, the thread is allowed to allocate an additional
--     100kb (tunable with @+RTS -xq@) to perform any necessary cleanup actions; if exceeded, the thread is delivered
--     another.
--
--     Default: @Nothing@ (no limit)
--
-- [@label@]:
--
--     The label of a thread, visible in the [event log](https://downloads.haskell.org/ghc/latest/docs/html/users_guide/runtime_control.html#rts-eventlog) (@+RTS -l@).
--
--     Default: @""@ (no label)
--
-- [@maskingState@]:
--
--     The masking state a thread is created in. To unmask, use 'GHC.IO.unsafeUnmask'.
--
--     Default: @Unmasked@
data ThreadOptions = ThreadOptions
  { ThreadOptions -> ThreadAffinity
affinity :: ThreadAffinity,
    ThreadOptions -> Maybe ByteCount
allocationLimit :: Maybe ByteCount,
    ThreadOptions -> String
label :: String,
    ThreadOptions -> MaskingState
maskingState :: MaskingState
  }
  deriving stock (ThreadOptions -> ThreadOptions -> Bool
(ThreadOptions -> ThreadOptions -> Bool)
-> (ThreadOptions -> ThreadOptions -> Bool) -> Eq ThreadOptions
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: ThreadOptions -> ThreadOptions -> Bool
$c/= :: ThreadOptions -> ThreadOptions -> Bool
== :: ThreadOptions -> ThreadOptions -> Bool
$c== :: ThreadOptions -> ThreadOptions -> Bool
Eq, Int -> ThreadOptions -> ShowS
[ThreadOptions] -> ShowS
ThreadOptions -> String
(Int -> ThreadOptions -> ShowS)
-> (ThreadOptions -> String)
-> ([ThreadOptions] -> ShowS)
-> Show ThreadOptions
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [ThreadOptions] -> ShowS
$cshowList :: [ThreadOptions] -> ShowS
show :: ThreadOptions -> String
$cshow :: ThreadOptions -> String
showsPrec :: Int -> ThreadOptions -> ShowS
$cshowsPrec :: Int -> ThreadOptions -> ShowS
Show)

-- | Default thread options.
--
-- @
-- 'Ki.ThreadOptions'
--   { 'Ki.affinity' = 'Ki.Unbound'
--   , 'Ki.allocationLimit' = Nothing
--   , 'Ki.label' = ""
--   , 'Ki.maskingState' = 'Unmasked'
--   }
-- @
defaultThreadOptions :: ThreadOptions
defaultThreadOptions :: ThreadOptions
defaultThreadOptions =
  ThreadOptions
    { $sel:affinity:ThreadOptions :: ThreadAffinity
affinity = ThreadAffinity
Unbound,
      $sel:allocationLimit:ThreadOptions :: Maybe ByteCount
allocationLimit = Maybe ByteCount
forall a. Maybe a
Nothing,
      $sel:label:ThreadOptions :: String
label = String
"",
      $sel:maskingState:ThreadOptions :: MaskingState
maskingState = MaskingState
Unmasked
    }

-- Internal exception type thrown by a child thread to its parent, if it fails unexpectedly.
data ThreadFailed = ThreadFailed
  { ThreadFailed -> Int
childId :: {-# UNPACK #-} !Int,
    ThreadFailed -> SomeException
exception :: !SomeException
  }
  deriving stock (Int -> ThreadFailed -> ShowS
[ThreadFailed] -> ShowS
ThreadFailed -> String
(Int -> ThreadFailed -> ShowS)
-> (ThreadFailed -> String)
-> ([ThreadFailed] -> ShowS)
-> Show ThreadFailed
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [ThreadFailed] -> ShowS
$cshowList :: [ThreadFailed] -> ShowS
show :: ThreadFailed -> String
$cshow :: ThreadFailed -> String
showsPrec :: Int -> ThreadFailed -> ShowS
$cshowsPrec :: Int -> ThreadFailed -> ShowS
Show)

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

unwrapThreadFailed :: SomeException -> SomeException
unwrapThreadFailed :: SomeException -> SomeException
unwrapThreadFailed SomeException
e0 =
  case SomeException -> Maybe ThreadFailed
forall e. Exception e => SomeException -> Maybe e
fromException SomeException
e0 of
    Just (ThreadFailed Int
_ SomeException
e1) -> SomeException
e1
    Maybe ThreadFailed
Nothing -> SomeException
e0

-- | Wait for a thread to terminate.
await :: Thread a -> STM a
await :: forall a. Thread a -> STM a
await =
  Thread a -> STM a
forall a. Thread a -> STM a
await_

-- Like try, but with continuations
tryEitherSTM :: Exception e => (e -> STM b) -> (a -> STM b) -> STM a -> STM b
tryEitherSTM :: forall e b a.
Exception e =>
(e -> STM b) -> (a -> STM b) -> STM a -> STM b
tryEitherSTM e -> STM b
onFailure a -> STM b
onSuccess STM a
action =
  STM (STM b) -> STM b
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join (STM (STM b) -> (e -> STM (STM b)) -> STM (STM b)
forall e a. Exception e => STM a -> (e -> STM a) -> STM a
catchSTM (a -> STM b
onSuccess (a -> STM b) -> STM a -> STM (STM b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> STM a
action) (STM b -> STM (STM b)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (STM b -> STM (STM b)) -> (e -> STM b) -> e -> STM (STM b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. e -> STM b
onFailure))