module TheatreDev.Daemon
  ( Daemon,

    -- * Acquisition
    spawn,

    -- * Control
    kill,
    wait,
  )
where

import TheatreDev.Prelude
import TheatreDev.Wait qualified as Wait

-- |
-- Think of an actor that does not process any messages and simply
-- interrupts between each iteration to check whether it's still alive.
data Daemon = Daemon
  { -- | Kill the daemon.
    Daemon -> STM ()
kill :: STM (),
    -- | Wait for the daemon to die due to error or being killed.
    Daemon -> STM (Maybe SomeException)
wait :: STM (Maybe SomeException)
  }

instance Semigroup Daemon where
  Daemon
left <> :: Daemon -> Daemon -> Daemon
<> Daemon
right =
    Daemon
      { $sel:kill:Daemon :: STM ()
kill = Daemon
left.kill forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Daemon
right.kill,
        $sel:wait:Daemon :: STM (Maybe SomeException)
wait = STM (Maybe SomeException)
-> STM (Maybe SomeException) -> STM (Maybe SomeException)
Wait.both Daemon
left.wait Daemon
right.wait
      }

instance Monoid Daemon where
  mempty :: Daemon
mempty =
    Daemon
      { $sel:kill:Daemon :: STM ()
kill = forall (m :: * -> *) a. Monad m => a -> m a
return (),
        $sel:wait:Daemon :: STM (Maybe SomeException)
wait = forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing
      }
  mconcat :: [Daemon] -> Daemon
mconcat [Daemon]
daemons =
    Daemon
      { $sel:kill:Daemon :: STM ()
kill = forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ (.kill) [Daemon]
daemons,
        $sel:wait:Daemon :: STM (Maybe SomeException)
wait = [STM (Maybe SomeException)] -> STM (Maybe SomeException)
Wait.all (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (.wait) [Daemon]
daemons)
      }

-- | Fork a thread to run the daemon loop on
-- returning immediately with a handle to control it.
spawn ::
  -- | Initial state of the daemon.
  state ->
  -- | Iteration action, updating the daemon's state.
  -- It gets executed in a loop,
  -- with checks of whether the daemon is still alive after each one.
  -- Killing the daemon will not interrupt the currently ongoing iteration,
  -- thus providing gracefulness guarantees.
  --
  -- If an exception is thrown by this action,
  -- the iteration loop will stop,
  -- the 'cleanUp' action will get executed and
  -- in all places where 'wait' is called the exception will be rethrown.
  (state -> IO state) ->
  -- | Clean up after the iteration loop is stopped.
  -- You can use that to release resources or
  -- issue notifications about the daemon dying.
  (state -> IO ()) ->
  IO Daemon
spawn :: forall state.
state -> (state -> IO state) -> (state -> IO ()) -> IO Daemon
spawn state
initialState state -> IO state
iterate state -> IO ()
cleanUp = do
  TVar Bool
iteratingVar <- forall a. a -> IO (TVar a)
newTVarIO Bool
True
  TMVar (Maybe SomeException)
resultVar <- forall a. IO (TMVar a)
newEmptyTMVarIO
  ((forall a. IO a -> IO a) -> IO ()) -> IO ThreadId
forkIOWithUnmask forall a b. (a -> b) -> a -> b
$ \forall a. IO a -> IO a
unmask ->
    let go :: state -> IO ()
go !state
state = do
          Bool
iterating <- forall a. TVar a -> IO a
readTVarIO TVar Bool
iteratingVar
          if Bool
iterating
            then do
              Either SomeException state
iterationAttemptResult <- forall e a. Exception e => IO a -> IO (Either e a)
try (forall a. IO a -> IO a
unmask (state -> IO state
iterate state
state))
              case Either SomeException state
iterationAttemptResult of
                Right state
newState -> state -> IO ()
go state
newState
                Left SomeException
exception -> do
                  forall e a. Exception e => IO a -> IO (Either e a)
try @SomeException (forall a. IO a -> IO a
unmask (state -> IO ()
cleanUp state
state))
                  forall a. STM a -> IO a
atomically (forall a. TMVar a -> a -> STM ()
putTMVar TMVar (Maybe SomeException)
resultVar (forall a. a -> Maybe a
Just SomeException
exception))
            else do
              Either SomeException ()
cleanUpResult <- forall e a. Exception e => IO a -> IO (Either e a)
try @SomeException (forall a. IO a -> IO a
unmask (state -> IO ()
cleanUp state
state))
              case Either SomeException ()
cleanUpResult of
                Right () -> forall a. STM a -> IO a
atomically (forall a. TMVar a -> a -> STM ()
putTMVar TMVar (Maybe SomeException)
resultVar forall a. Maybe a
Nothing)
                Left SomeException
exception -> forall a. STM a -> IO a
atomically (forall a. TMVar a -> a -> STM ()
putTMVar TMVar (Maybe SomeException)
resultVar (forall a. a -> Maybe a
Just SomeException
exception))
     in state -> IO ()
go state
initialState
  return
    Daemon
      { $sel:kill:Daemon :: STM ()
kill = forall a. TVar a -> a -> STM ()
writeTVar TVar Bool
iteratingVar Bool
False,
        $sel:wait:Daemon :: STM (Maybe SomeException)
wait = forall a. TMVar a -> STM a
readTMVar TMVar (Maybe SomeException)
resultVar
      }
  where

-- | Command the daemon to stop iterating,
-- finish the ongoing iteration and execute the clean up action.
--
-- This action executes immediately.
-- If you want to block waiting for the daemon to actually die,
-- after 'kill' you can run 'wait'.
kill :: Daemon -> IO ()
kill :: Daemon -> IO ()
kill Daemon
daemon =
  forall a. STM a -> IO a
atomically Daemon
daemon.kill

-- | Block waiting for the daemon to die and execute its 'cleanUp' action
-- either due to getting killed
-- or due to its 'iterate' action throwing an exception.
-- The exception will get rethrown here.
wait :: Daemon -> IO ()
wait :: Daemon -> IO ()
wait Daemon
daemon =
  forall a. STM a -> IO a
atomically Daemon
daemon.wait forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall b a. b -> (a -> b) -> Maybe a -> b
maybe (forall (f :: * -> *) a. Applicative f => a -> f a
pure ()) forall e a. Exception e => e -> IO a
throwIO