{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE NumericUnderscores #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TypeApplications #-}

-- | This module provides a test watchdog - an utility monitoring test cases and killing them if they don't
-- finish in time. 'Watchdog' thread runs in the background, and after specified timeout, it throws
-- 'WatchdogException' to the target thread. A user is able to 'kickWatchdog', which delays the killing and
-- 'poisonWatchdog' which stops the watchdog.
--
-- To wrap a test case in a watchdog just use
--
-- @
-- runWithWatchdog watchdogConfig $ \\watchdog -> do
--   -- body of your test case
-- @
--
module Hedgehog.Extras.Test.TestWatchdog
  (
  -- * Wrap in watchdog
    runWithWatchdog_
  , runWithWatchdog
  , runWithDefaultWatchdog_
  , runWithDefaultWatchdog

  -- * Watchdog control
  , kickWatchdog
  , poisonWatchdog

  -- * Types
  , Watchdog
  , WatchdogConfig(..)
  , WatchdogException(..)

  -- * Low level API
  -- | There is also a lower-level API available, giving the ability to provide target thread ID, which watchdog
  -- will try to kill.

  , makeWatchdog
  , runWatchdog
  ) where

import           Control.Concurrent (myThreadId, threadDelay, throwTo)
import           Control.Concurrent.STM (atomically)
import           Control.Concurrent.STM.TChan (TChan, newTChanIO, tryReadTChan, writeTChan)
import           Control.Exception (Exception)
import           Control.Monad.IO.Class
import           Data.Time (NominalDiffTime, UTCTime, diffUTCTime, getCurrentTime,
                   nominalDiffTimeToSeconds)
import           GHC.Conc (ThreadId)
import           GHC.Stack

import           Control.Monad.Base (MonadBase (..))
import           Control.Monad.Trans.Control (MonadBaseControl)
import qualified Hedgehog.Extras.Test.Concurrent as H
import           Prelude

-- | Configuration for the watchdog.
newtype WatchdogConfig = WatchdogConfig
  { WatchdogConfig -> Int
watchdogTimeout :: Int -- ^ Timeout in seconds after which watchdog will kill the test case
  }

-- | Default watchdog configuration with 10 minutes timeout.
defaultWatchdogConfig :: WatchdogConfig
defaultWatchdogConfig :: WatchdogConfig
defaultWatchdogConfig = WatchdogConfig
  { watchdogTimeout :: Int
watchdogTimeout = Int
600
  }

-- | A watchdog instance. See the module header for more detailed description.
data Watchdog = Watchdog
  { Watchdog -> WatchdogConfig
watchdogConfig :: !WatchdogConfig
  , Watchdog -> ThreadId
watchedThreadId :: !ThreadId -- ^ monitored thread id
  , Watchdog -> UTCTime
startTime :: !UTCTime -- ^ watchdog creation time
  , Watchdog -> TChan WatchdogCommand
kickChan :: TChan WatchdogCommand -- ^ a queue of watchdog commands
  }

instance Show Watchdog where
  show :: Watchdog -> String
show Watchdog{watchdogConfig :: Watchdog -> WatchdogConfig
watchdogConfig=WatchdogConfig{Int
watchdogTimeout :: WatchdogConfig -> Int
watchdogTimeout :: Int
watchdogTimeout}, UTCTime
startTime :: Watchdog -> UTCTime
startTime :: UTCTime
startTime, ThreadId
watchedThreadId :: Watchdog -> ThreadId
watchedThreadId :: ThreadId
watchedThreadId} = [String] -> String
forall a. Monoid a => [a] -> a
mconcat
    [ String
"Watchdog with timeout ", Int -> String
forall a. Show a => a -> String
show Int
watchdogTimeout
    , String
", started at ", UTCTime -> String
forall a. Show a => a -> String
show UTCTime
startTime
    , String
", watching thread ID ", ThreadId -> String
forall a. Show a => a -> String
show ThreadId
watchedThreadId
    ]

-- | Create manually a new watchdog, providing the target thread ID. After all watchdog timeouts expire,
-- the target thread will get 'WatchdogException' thrown to it asynchronously (using 'throwTo').
makeWatchdog :: MonadBase IO m
             => WatchdogConfig
             -> ThreadId -- ^ thread id which will get killed after all kicks expire
             -> m Watchdog
makeWatchdog :: forall (m :: * -> *).
MonadBase IO m =>
WatchdogConfig -> ThreadId -> m Watchdog
makeWatchdog WatchdogConfig
config ThreadId
watchedThreadId' = IO Watchdog -> m Watchdog
forall α. IO α -> m α
forall (b :: * -> *) (m :: * -> *) α. MonadBase b m => b α -> m α
liftBase (IO Watchdog -> m Watchdog) -> IO Watchdog -> m Watchdog
forall a b. (a -> b) -> a -> b
$ do
  Watchdog
watchdog <- WatchdogConfig
-> ThreadId -> UTCTime -> TChan WatchdogCommand -> Watchdog
Watchdog WatchdogConfig
config ThreadId
watchedThreadId' (UTCTime -> TChan WatchdogCommand -> Watchdog)
-> IO UTCTime -> IO (TChan WatchdogCommand -> Watchdog)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO UTCTime
getCurrentTime IO (TChan WatchdogCommand -> Watchdog)
-> IO (TChan WatchdogCommand) -> IO Watchdog
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> IO (TChan WatchdogCommand)
forall a. IO (TChan a)
newTChanIO
  Watchdog -> IO ()
forall (m :: * -> *). MonadIO m => Watchdog -> m ()
kickWatchdog Watchdog
watchdog
  Watchdog -> IO Watchdog
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Watchdog
watchdog

-- | Run watchdog in a loop in the current thread. Usually this function should be used with 'H.withAsync'
-- to run it in the background.
runWatchdog :: MonadBase IO m
            => Watchdog
            -> m ()
runWatchdog :: forall (m :: * -> *). MonadBase IO m => Watchdog -> m ()
runWatchdog w :: Watchdog
w@Watchdog{ThreadId
watchedThreadId :: Watchdog -> ThreadId
watchedThreadId :: ThreadId
watchedThreadId, UTCTime
startTime :: Watchdog -> UTCTime
startTime :: UTCTime
startTime, TChan WatchdogCommand
kickChan :: Watchdog -> TChan WatchdogCommand
kickChan :: TChan WatchdogCommand
kickChan} = IO () -> m ()
forall α. IO α -> m α
forall (b :: * -> *) (m :: * -> *) α. MonadBase b m => b α -> m α
liftBase (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ do
  STM (Maybe WatchdogCommand) -> IO (Maybe WatchdogCommand)
forall a. STM a -> IO a
atomically (TChan WatchdogCommand -> STM (Maybe WatchdogCommand)
forall a. TChan a -> STM (Maybe a)
tryReadTChan TChan WatchdogCommand
kickChan) IO (Maybe WatchdogCommand)
-> (Maybe WatchdogCommand -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Just WatchdogCommand
PoisonPill ->
      -- deactivate watchdog
      () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    Just (Kick Int
timeout) -> do
      -- got a kick, wait for another period
      Int -> IO ()
threadDelay (Int -> IO ()) -> Int -> IO ()
forall a b. (a -> b) -> a -> b
$ Int
timeout Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
1_000_000
      Watchdog -> IO ()
forall (m :: * -> *). MonadBase IO m => Watchdog -> m ()
runWatchdog Watchdog
w
    Maybe WatchdogCommand
Nothing -> do
      -- we are out of scheduled timeouts, kill the monitored thread
      UTCTime
currentTime <- IO UTCTime
getCurrentTime
      ThreadId -> WatchdogException -> IO ()
forall e. Exception e => ThreadId -> e -> IO ()
throwTo ThreadId
watchedThreadId (WatchdogException -> IO ())
-> (NominalDiffTime -> WatchdogException)
-> NominalDiffTime
-> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NominalDiffTime -> WatchdogException
WatchdogException (NominalDiffTime -> IO ()) -> NominalDiffTime -> IO ()
forall a b. (a -> b) -> a -> b
$ UTCTime -> UTCTime -> NominalDiffTime
diffUTCTime UTCTime
currentTime UTCTime
startTime

-- | Watchdog command
data WatchdogCommand
  = Kick !Int -- ^ Add another delay in seconds
  | PoisonPill -- ^ Stop the watchdog

-- | Enqueue a kick for the watchdog. It will extend the timeout by another one defined in the watchdog
-- configuration.
kickWatchdog :: MonadIO m => Watchdog -> m ()
kickWatchdog :: forall (m :: * -> *). MonadIO m => Watchdog -> m ()
kickWatchdog Watchdog{watchdogConfig :: Watchdog -> WatchdogConfig
watchdogConfig=WatchdogConfig{Int
watchdogTimeout :: WatchdogConfig -> Int
watchdogTimeout :: Int
watchdogTimeout}, TChan WatchdogCommand
kickChan :: Watchdog -> TChan WatchdogCommand
kickChan :: TChan WatchdogCommand
kickChan} = IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$
  STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TChan WatchdogCommand -> WatchdogCommand -> STM ()
forall a. TChan a -> a -> STM ()
writeTChan TChan WatchdogCommand
kickChan (Int -> WatchdogCommand
Kick Int
watchdogTimeout)

-- | Enqueue a poison pill for the watchdog. It will stop the watchdog after all timeouts.
poisonWatchdog :: MonadIO m => Watchdog -> m ()
poisonWatchdog :: forall (m :: * -> *). MonadIO m => Watchdog -> m ()
poisonWatchdog Watchdog{TChan WatchdogCommand
kickChan :: Watchdog -> TChan WatchdogCommand
kickChan :: TChan WatchdogCommand
kickChan} = IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$
  STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TChan WatchdogCommand -> WatchdogCommand -> STM ()
forall a. TChan a -> a -> STM ()
writeTChan TChan WatchdogCommand
kickChan WatchdogCommand
PoisonPill


-- | Execute a test case with a watchdog.
runWithWatchdog :: HasCallStack
                => MonadBaseControl IO m
                => WatchdogConfig -- ^ configuration
                -> (HasCallStack => Watchdog -> m a) -- ^ a test case to be wrapped in watchdog
                -> m a
runWithWatchdog :: forall (m :: * -> *) a.
(HasCallStack, MonadBaseControl IO m) =>
WatchdogConfig -> (HasCallStack => Watchdog -> m a) -> m a
runWithWatchdog WatchdogConfig
config HasCallStack => Watchdog -> m a
testCase = do
  ThreadId
watchedThreadId <- IO ThreadId -> m ThreadId
forall α. IO α -> m α
forall (b :: * -> *) (m :: * -> *) α. MonadBase b m => b α -> m α
liftBase IO ThreadId
myThreadId
  Watchdog
watchdog <- IO Watchdog -> m Watchdog
forall α. IO α -> m α
forall (b :: * -> *) (m :: * -> *) α. MonadBase b m => b α -> m α
liftBase (IO Watchdog -> m Watchdog) -> IO Watchdog -> m Watchdog
forall a b. (a -> b) -> a -> b
$ WatchdogConfig -> ThreadId -> IO Watchdog
forall (m :: * -> *).
MonadBase IO m =>
WatchdogConfig -> ThreadId -> m Watchdog
makeWatchdog WatchdogConfig
config ThreadId
watchedThreadId
  m () -> (Async (StM m ()) -> m a) -> m a
forall (m :: * -> *) a b.
MonadBaseControl IO m =>
m a -> (Async (StM m a) -> m b) -> m b
H.withAsync (Watchdog -> m ()
forall (m :: * -> *). MonadBase IO m => Watchdog -> m ()
runWatchdog Watchdog
watchdog) ((Async (StM m ()) -> m a) -> m a)
-> (Async (StM m ()) -> m a) -> m a
forall a b. (a -> b) -> a -> b
$
    \Async (StM m ())
_ -> HasCallStack => Watchdog -> m a
Watchdog -> m a
testCase Watchdog
watchdog

-- | Execute a test case with a watchdog.
runWithWatchdog_ :: HasCallStack
                 => MonadBaseControl IO m
                 => WatchdogConfig -- ^ configuration
                 -> (HasCallStack => m a) -- ^ a test case to be wrapped in watchdog
                 -> m a
runWithWatchdog_ :: forall (m :: * -> *) a.
(HasCallStack, MonadBaseControl IO m) =>
WatchdogConfig -> (HasCallStack => m a) -> m a
runWithWatchdog_ WatchdogConfig
config HasCallStack => m a
testCase = WatchdogConfig -> (HasCallStack => Watchdog -> m a) -> m a
forall (m :: * -> *) a.
(HasCallStack, MonadBaseControl IO m) =>
WatchdogConfig -> (HasCallStack => Watchdog -> m a) -> m a
runWithWatchdog WatchdogConfig
config (m a -> Watchdog -> m a
forall a b. a -> b -> a
const m a
HasCallStack => m a
testCase)

-- | Execute a test case with watchdog with default config.
runWithDefaultWatchdog :: HasCallStack
                       => MonadBaseControl IO m
                       => (HasCallStack => Watchdog -> m a) -- ^ a test case to be wrapped in watchdog
                       -> m a
runWithDefaultWatchdog :: forall (m :: * -> *) a.
(HasCallStack, MonadBaseControl IO m) =>
(HasCallStack => Watchdog -> m a) -> m a
runWithDefaultWatchdog = WatchdogConfig -> (HasCallStack => Watchdog -> m a) -> m a
forall (m :: * -> *) a.
(HasCallStack, MonadBaseControl IO m) =>
WatchdogConfig -> (HasCallStack => Watchdog -> m a) -> m a
runWithWatchdog WatchdogConfig
defaultWatchdogConfig

-- | Execute a test case with watchdog with default config.
runWithDefaultWatchdog_ :: HasCallStack
                        => MonadBaseControl IO m
                        => (HasCallStack => m a) -- ^ a test case to be wrapped in watchdog
                        -> m a
runWithDefaultWatchdog_ :: forall (m :: * -> *) a.
(HasCallStack, MonadBaseControl IO m) =>
(HasCallStack => m a) -> m a
runWithDefaultWatchdog_ HasCallStack => m a
testCase = (HasCallStack => Watchdog -> m a) -> m a
forall (m :: * -> *) a.
(HasCallStack, MonadBaseControl IO m) =>
(HasCallStack => Watchdog -> m a) -> m a
runWithDefaultWatchdog (m a -> Watchdog -> m a
forall a b. a -> b -> a
const m a
HasCallStack => m a
testCase)

-- | An exception thrown to the test case thread.
newtype WatchdogException = WatchdogException { WatchdogException -> NominalDiffTime
timeElapsed :: NominalDiffTime }

instance Show WatchdogException where
  show :: WatchdogException -> String
show WatchdogException{NominalDiffTime
timeElapsed :: WatchdogException -> NominalDiffTime
timeElapsed :: NominalDiffTime
timeElapsed} =
    String
"WatchdogException: Test watchdog killed test case thread after " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> String
show @Int (Pico -> Int
forall b. Integral b => Pico -> b
forall a b. (RealFrac a, Integral b) => a -> b
round (Pico -> Int) -> Pico -> Int
forall a b. (a -> b) -> a -> b
$ NominalDiffTime -> Pico
nominalDiffTimeToSeconds NominalDiffTime
timeElapsed) String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
" seconds."

instance Exception WatchdogException