content: -- | Monitored threads
  --
  -- Intended for unqualified import.
  module Network.GRPC.Util.Thread (
      ThreadState(..)
      -- * Creating threads
    , ThreadBody
    , newThreadState
    , forkThread
    , threadBody
      -- * Thread debug ID
    , DebugThreadId -- opaque
    , threadDebugId
      -- * Access thread state
    , CancelResult(..)
    , cancelThread
    , ThreadState_(..)
    , getThreadState_
    , unlessAbnormallyTerminated
    , withThreadInterface
    , waitForNormalThreadTermination
    , waitForNormalOrAbnormalThreadTermination
    ) where
  
  import Control.Concurrent
  import Control.Concurrent.STM
  import Control.Exception
  import Control.Monad
  import Data.Void (Void, absurd)
  import Foreign (newStablePtr, freeStablePtr)
  import GHC.Stack
  import System.IO.Unsafe (unsafePerformIO)
  
  import Network.GRPC.Util.GHC
  
  {-------------------------------------------------------------------------------
    Debug thread IDs
  -------------------------------------------------------------------------------}
  
  -- | Debug thread IDs
  --
  -- Unlike 'ThreadId', these do not correspond to a /running/ thread necessarily,
  -- but just enable us to distinguish one thread from another.
  data DebugThreadId = DebugThreadId {
        debugThreadId        :: Word
      , debugThreadCreatedAt :: CallStack
      }
    deriving stock (Show)
  
  nextDebugThreadId :: MVar Word
  {-# NOINLINE nextDebugThreadId #-}
  nextDebugThreadId = unsafePerformIO $ newMVar 0
  
  newDebugThreadId :: HasCallStack => IO DebugThreadId
  newDebugThreadId =
      modifyMVar nextDebugThreadId $ \x -> do
        let !nextId = succ x
        return (
            nextId
          , DebugThreadId x (popIrrelevant callStack)
          )
    where
      -- Pop off the call to 'newDebugThreadId'
      --
      -- We leave the call to 'newThreadState' on the stack because it is useful
      -- to know where that was called /from/.
      popIrrelevant :: CallStack -> CallStack
      popIrrelevant = popCallStack
  
  {-------------------------------------------------------------------------------
    State
  -------------------------------------------------------------------------------}
  
  -- | State of a thread with public interface of type @a@
  data ThreadState a =
      -- | The thread has not yet started
      --
      -- If the thread is cancelled before it is started, then the exception will
      -- be delivered once started. This is important, because it gives the thread
      -- control over /when/ the exception is delivered (that is, when it chooses
      -- to unmask async exceptions).
      --
      -- The alternative would be not to start the thread at all in this case, but
      -- this takes away the control mentioned above; if the thread /needs/ to do
      -- something before it can be killed, it must be given that chance. It may
      -- /seem/ that this alternative would give the caller (which /created/ the
      -- thread) more control, but actually that control is illusory, since the
      -- timing of async exceptions is anyway unpredictable.
      ThreadNotStarted DebugThreadId
  
      -- | The externally visible thread interface is still being initialized
    | ThreadInitializing DebugThreadId ThreadId
  
      -- | Thread is ready
    | ThreadRunning DebugThreadId ThreadId a
  
      -- | Thread terminated normally
      --
      -- This still carries the thread interface: we may need it to query the
      -- thread's final status, for example.
    | ThreadDone DebugThreadId a
  
      -- | Thread terminated with an exception
    | ThreadException DebugThreadId SomeException
    deriving stock (Show, Functor)
  
  threadDebugId :: ThreadState a -> DebugThreadId
  threadDebugId (ThreadNotStarted   debugId    ) = debugId
  threadDebugId (ThreadInitializing debugId _  ) = debugId
  threadDebugId (ThreadRunning      debugId _ _) = debugId
  threadDebugId (ThreadDone         debugId   _) = debugId
  threadDebugId (ThreadException    debugId   _) = debugId
  
  {-------------------------------------------------------------------------------
    Creating threads
  -------------------------------------------------------------------------------}
  
  type ThreadBody a =
            (forall x. IO x -> IO x) -- ^ Unmask exceptions
         -> (a -> IO ())             -- ^ Mark thread ready
         -> DebugThreadId            -- ^ Unique identifier for this thread
         -> IO ()
  
  newThreadState :: HasCallStack => IO (TVar (ThreadState a))
  newThreadState = do
      debugId <- newDebugThreadId
      newTVarIO $ ThreadNotStarted debugId
  
  forkThread ::
       HasCallStack
    => ThreadLabel -> TVar (ThreadState a) -> ThreadBody a -> IO ()
  forkThread label state body =
      void $ mask_ $ forkIOWithUnmask $ \unmask ->
        threadBody label state $ body unmask
  
  -- | Wrap the thread body
  --
  -- This should be wrapped around the body of the thread, and should be called
  -- with exceptions masked.
  --
  -- This is intended for integration with existing libraries (such as @http2@),
  -- which might do the forking under the hood.
  --
  -- If the 'ThreadState' is anything other than 'ThreadNotStarted' on entry,
  -- this function terminates immediately.
  threadBody :: forall a.
       HasCallStack
    => ThreadLabel
    -> TVar (ThreadState a)
    -> ((a -> IO ()) -> DebugThreadId -> IO ())
    -> IO ()
  threadBody label state body = do
      labelThisThread label
      threadId  <- myThreadId
      initState <- readTVarIO state
  
      -- See discussion of 'ThreadNotStarted'
      -- It's critical that async exceptions are masked at this point.
      case initState of
        ThreadNotStarted debugId -> do
          atomically $ writeTVar state $ ThreadInitializing debugId threadId
        ThreadException _ exception ->
          -- We don't change the thread status here: 'cancelThread' offers the
          -- guarantee that the thread status /will/ be in aborted or done state
          -- on return. This means that /externally/ the thread will be
          -- considered done, even if perhaps the thread must still execute some
          -- actions before it can actually terminate.
          void . forkIO $ throwTo threadId exception
        _otherwise -> do
          unexpected "initState" initState
  
      let markReady :: a -> STM ()
          markReady a = do
              modifyTVar state $ \oldState ->
                case oldState of
                  ThreadInitializing debugId _ ->
                    ThreadRunning debugId threadId a
                  ThreadException _ _ ->
                    oldState -- leave alone (see discussion above)
                  _otherwise ->
                    unexpected "markReady" oldState
  
          markDone :: Either SomeException () -> STM ()
          markDone mDone = do
              modifyTVar state $ \oldState ->
                case (oldState, mDone) of
                  (ThreadRunning debugId _ iface, Right ()) ->
                    ThreadDone debugId iface
                  (ThreadException{}, _) ->
                    oldState -- record /first/ exception
                  (_, Left e) ->
                    ThreadException (threadDebugId oldState) e
                  _otherwise ->
                    unexpected "markDone" oldState
  
      res <- try $ body (atomically . markReady) (threadDebugId initState)
      atomically $ markDone res
    where
      unexpected :: String -> ThreadState a -> x
      unexpected msg st = error $ concat [
            msg
          , ": unexpected "
          , show (const () <$> st)
          ]
  
  {-------------------------------------------------------------------------------
    Stopping
  -------------------------------------------------------------------------------}
  
  -- | Result of cancelling a thread
  data CancelResult a =
      -- | The thread terminated normally before we could cancel it
      AlreadyTerminated a
  
      -- | The thread terminated with an exception before we could cancel it
    | AlreadyAborted SomeException
  
      -- | We killed the thread with the specified exception
    | Cancelled
  
  -- | Kill thread if it is running
  --
  -- * If the thread is in `ThreadNotStarted` state, we merely change the state to
  --   'ThreadException'. The thread may still be started (see discussion of
  --   'ThreadNotStarted'), but /externally/ the thread will be considered to have
  --   terminated.
  --
  -- * If the thread is initializing or running, we update the state to
  --   'ThreadException' and then throw the specified exception to the thread.
  --
  -- * If the thread is /already/ in 'ThreadException' state, or if the thread is
  --   in 'ThreadDone' state, we do nothing.
  --
  -- In all cases, the caller is guaranteed that the thread state has been updated
  -- even if perhaps the thread is still shutting down.
  cancelThread :: forall a.
       TVar (ThreadState a)
    -> SomeException
    -> IO (CancelResult a)
  cancelThread state e = do
      (result, mTid) <- atomically aux
      forM_ mTid $ flip throwTo e
      return result
    where
      aux :: STM (CancelResult a, Maybe ThreadId)
      aux = do
          st <- readTVar state
          case st of
            ThreadNotStarted debugId -> do
              writeTVar state $ ThreadException debugId e
              return (Cancelled, Nothing)
            ThreadInitializing debugId threadId -> do
              writeTVar state $ ThreadException debugId e
              return (Cancelled, Just threadId)
            ThreadRunning debugId threadId _ -> do
              writeTVar state $ ThreadException debugId e
              return (Cancelled, Just threadId)
            ThreadException _debugId e' ->
              return (AlreadyAborted e', Nothing)
            ThreadDone _debugId a ->
              return (AlreadyTerminated a, Nothing)
  
  {-------------------------------------------------------------------------------
    Interacting with the thread
  -------------------------------------------------------------------------------}
  
  -- | Get the thread's interface
  --
  -- The behaviour of this 'getThreadInterface' depends on the thread state; it
  --
  -- * blocks if the thread in case of 'ThreadNotStarted' or 'ThreadInitializing'
  -- * throws 'ThreadInterfaceUnavailable' in case of 'ThreadException'.
  -- * returns the thread interface otherwise
  --
  -- We do /not/ distinguish between 'ThreadDone' and 'ThreadRunning' here, as
  -- doing so is inherently racy (we might return that the client is still
  -- running, and then it terminates before the calling code can do anything with
  -- that information).
  --
  -- NOTE: This turns off deadlock detection for the duration of the transaction.
  -- It should therefore only be used for transactions that can never be blocked
  -- indefinitely.
  --
  -- Usage note: in practice we use this to interact with threads that in turn
  -- interact with @http2@ (running 'sendMessageLoop' or 'recvMessageLoop').
  -- Although calls into @http2@ may in fact block indefinitely, we will /catch/
  -- those exception and treat them as network failures. If a @grapesy@ function
  -- ever throws a "blocked indefinitely" exception, this should be reported as a
  -- bug in @grapesy@.
  withThreadInterface :: forall a b.
       TVar (ThreadState a)
    -> (a -> STM b)
    -> IO b
  withThreadInterface state k =
      withoutDeadlockDetection . atomically $
        k =<< getThreadInterface
    where
      getThreadInterface :: STM a
      getThreadInterface = do
          st <- readTVar state
          case st of
            ThreadNotStarted   _     -> retry
            ThreadInitializing _ _   -> retry
            ThreadRunning      _ _ a -> return a
            ThreadDone         _   a -> return a
            ThreadException    _   e -> throwSTM e
  
  -- | Wait for the thread to terminate normally
  --
  -- If the thread terminated with an exception, this rethrows that exception.
  waitForNormalThreadTermination :: TVar (ThreadState a) -> STM ()
  waitForNormalThreadTermination state =
      waitUntilInitialized state >>= \case
        ThreadNotYetRunning_ v -> absurd v
        ThreadRunning_         -> retry
        ThreadDone_            -> return ()
        ThreadException_ e     -> throwSTM e
  
  -- | Wait for the thread to terminate normally or abnormally
  waitForNormalOrAbnormalThreadTermination ::
       TVar (ThreadState a)
    -> STM (Maybe SomeException)
  waitForNormalOrAbnormalThreadTermination state =
      waitUntilInitialized state >>= \case
        ThreadNotYetRunning_ v -> absurd v
        ThreadRunning_         -> retry
        ThreadDone_            -> return $ Nothing
        ThreadException_ e     -> return $ Just e
  
  -- | Run the specified transaction, unless the thread terminated with an
  -- exception
  unlessAbnormallyTerminated ::
       TVar (ThreadState a)
    -> STM b
    -> STM (Either SomeException b)
  unlessAbnormallyTerminated state f =
      waitUntilInitialized state >>= \case
        ThreadNotYetRunning_ v -> absurd v
        ThreadRunning_         -> Right <$> f
        ThreadDone_            -> Right <$> f
        ThreadException_ e     -> return $ Left e
  
  waitUntilInitialized ::
       TVar (ThreadState a)
    -> STM (ThreadState_ Void)
  waitUntilInitialized state = getThreadState_ state retry
  
  -- | An abstraction of 'ThreadState' without the public interface type.
  data ThreadState_ notRunning =
        ThreadNotYetRunning_ notRunning
      | ThreadRunning_
      | ThreadDone_
      | ThreadException_ SomeException
  
  getThreadState_ ::
       TVar (ThreadState a)
    -> STM notRunning
    -> STM (ThreadState_ notRunning)
  getThreadState_ state onNotRunning = do
      st <- readTVar state
      case st of
        ThreadNotStarted   _     -> ThreadNotYetRunning_ <$> onNotRunning
        ThreadInitializing _ _   -> ThreadNotYetRunning_ <$> onNotRunning
        ThreadRunning      _ _ _ -> return $ ThreadRunning_
        ThreadDone         _   _ -> return $ ThreadDone_
        ThreadException    _   e -> return $ ThreadException_ e
  
  {-------------------------------------------------------------------------------
    Internal auxiliary
  -------------------------------------------------------------------------------}
  
  -- | Locally turn off deadlock detection
  --
  -- See also <https://well-typed.com/blog/2024/01/when-blocked-indefinitely-is-not-indefinite/>.
  withoutDeadlockDetection :: IO a -> IO a
  withoutDeadlockDetection k = do
      threadId <- myThreadId
      bracket (newStablePtr threadId) freeStablePtr $ \_ -> k
  