Initial watchdog test

hedgehog-extras.cabal

@@ -131,6 +131,7 @@ library
                         Hedgehog.Extras.Test.MonadAssertion
                         Hedgehog.Extras.Test.Network
                         Hedgehog.Extras.Test.Process
+                        Hedgehog.Extras.Test.TestWatchdog
 
 test-suite hedgehog-extras-test
   import:               base, project-config,
@@ -139,10 +140,13 @@ test-suite hedgehog-extras-test
                         network,
                         tasty,
                         tasty-hedgehog,
+                        transformers,
   hs-source-dirs:       test
   main-is:              hedgehog-extras-test.hs
   type:                 exitcode-stdio-1.0
 
   other-modules:        Hedgehog.Extras.Stock.IO.Network.PortSpec
+                        Hedgehog.Extras.Test.TestWatchdogSpec
 
   build-tool-depends:   tasty-discover:tasty-discover
+  ghc-options:          -threaded -rtsopts "-with-rtsopts=-N -T"

src/Hedgehog/Extras/Test.hs

@@ -8,3 +8,4 @@ import Hedgehog.Extras.Test.File as X
 import Hedgehog.Extras.Test.MonadAssertion as X
 import Hedgehog.Extras.Test.Network as X
 import Hedgehog.Extras.Test.Process as X
+import Hedgehog.Extras.Test.TestWatchdog as X

src/Hedgehog/Extras/Test/TestWatchdog.hs

@@ -0,0 +1,146 @@
+{-# 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. To wrap an 'H.Integration' test case in a watchdog just use
+-- @
+-- runWithWatchdog watchdogConfig $ \watchdog -> do
+--   -- body of your test case
+-- @
+module Hedgehog.Extras.Test.TestWatchdog
+  ( runWithWatchdog_
+  , runWithWatchdog
+  , runWithDefaultWatchdog_
+  , runWithDefaultWatchdog
+  , Watchdog
+  , WatchdogConfig(..)
+  , kickWatchdog
+  , poisonWatchdog
+  ) 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 Prelude
+import Control.Monad.Trans.Control (MonadBaseControl)
+import Control.Monad.Base (MonadBase(..))
+import qualified Hedgehog.Extras.Test.Concurrent as H
+
+-- | Configuration for the watchdog.
+newtype WatchdogConfig = WatchdogConfig
+  { watchdogTimeout :: Int -- ^ Timeout in seconds after which watchdog will kill the test case
+  }
+
+-- | Default watchdog config with 10 minutes timeout.
+defaultWatchdogConfig :: WatchdogConfig
+defaultWatchdogConfig = WatchdogConfig
+  { watchdogTimeout = 600
+  }
+
+-- | A watchdog
+data Watchdog = Watchdog
+  { watchdogConfig :: !WatchdogConfig
+  , watchedThreadId :: !ThreadId -- ^ monitored thread id
+  , startTime :: !UTCTime -- ^ watchdog creation time
+  , kickChan :: TChan WatchdogCommand -- ^ a queue of watchdog commands
+  }
+
+-- | Create a new watchdog
+makeWatchdog :: MonadBase IO m
+             => WatchdogConfig
+             -> ThreadId -- ^ thread id which will get killed after timeouts expire
+             -> m Watchdog
+makeWatchdog config watchedThreadId' = liftBase $ do
+  watchdog <- Watchdog config watchedThreadId' <$> getCurrentTime <*> newTChanIO
+  kickWatchdog watchdog
+  pure watchdog
+
+-- | Run watchdog in a loop
+runWatchdog :: MonadBase IO m
+            => Watchdog
+            -> m ()
+runWatchdog w@Watchdog{watchedThreadId, startTime, kickChan} = liftBase $ do
+  atomically (tryReadTChan kickChan) >>= \case
+    Just PoisonPill ->
+      -- deactivate watchdog
+      pure ()
+    Just (Kick timeout) -> do
+      -- got a kick, wait for another period
+      threadDelay $ timeout * 1_000_000
+      runWatchdog w
+    Nothing -> do
+      -- we are out of scheduled timeouts, kill the monitored thread
+      currentTime <- getCurrentTime
+      throwTo watchedThreadId . WatchdogException $ diffUTCTime currentTime 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 Watchdog{watchdogConfig=WatchdogConfig{watchdogTimeout}, kickChan} = liftIO $
+  atomically $ writeTChan kickChan (Kick watchdogTimeout)
+
+-- | Enqueue a poison pill for the watchdog. It will stop the watchdog after all timeouts.
+poisonWatchdog :: MonadIO m => Watchdog -> m ()
+poisonWatchdog Watchdog{kickChan} = liftIO $
+  atomically $ writeTChan kickChan 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 config testCase = do
+  watchedThreadId <- liftBase myThreadId
+  watchdog <- liftBase $ makeWatchdog config watchedThreadId
+  H.withAsync (runWatchdog watchdog) $
+    \_ -> testCase watchdog
+
+-- | Execuate 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_ config testCase = runWithWatchdog config (const 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 = runWithWatchdog 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_ testCase = runWithDefaultWatchdog (const testCase)
+
+-- | An exception thrown to the test case thread.
+newtype WatchdogException = WatchdogException { timeElapsed :: NominalDiffTime }
+
+instance Show WatchdogException where
+  show WatchdogException{timeElapsed} =
+    "WatchdogException: Test watchdog killed test case thread after " <> show @Int (round $ nominalDiffTimeToSeconds timeElapsed) <> " seconds."
+
+instance Exception WatchdogException

test/Hedgehog/Extras/Test/TestWatchdogSpec.hs

@@ -0,0 +1,49 @@
+{-# LANGUAGE NumericUnderscores #-}
+module Hedgehog.Extras.Test.TestWatchdogSpec where
+
+import           Data.Function
+import           Data.Semigroup
+import           Hedgehog (Property)
+import qualified Hedgehog as H
+import qualified Hedgehog.Extras as H
+import qualified Hedgehog.Extras.Stock.IO.Network.Port as IO
+import Hedgehog.Extras.Test.TestWatchdog
+import qualified Hedgehog.Extras.Test.Concurrent as H
+import qualified Hedgehog.Internal.Property as H
+import           Control.Monad.Trans.Writer.Lazy (runWriterT)
+import           Control.Monad.Trans.Except (runExceptT)
+import qualified Network.Socket as N
+import           Text.Show
+import Control.Monad
+import Prelude
+import Control.Monad.IO.Class
+import GHC.Stack
+import Control.Concurrent
+import GHC.Conc
+
+hprop_check_watchdog_kills_hanged_thread :: Property
+hprop_check_watchdog_kills_hanged_thread = H.propertyOnce $ do
+  let watchdogCfg = WatchdogConfig 1
+  childTid <- liftIO $ newEmptyMVar
+  (res, log') <- runTestT $ runWithWatchdog_ watchdogCfg $ do
+    liftIO $ myThreadId >>= putMVar childTid 
+    void . forever $ H.threadDelay 1_000_000
+    H.success
+  H.threadDelay 1_000_000
+  H.noteShowM_ . liftIO $ (readMVar childTid >>= threadStatus)
+  H.failure
+
+runTestT :: HasCallStack
+         => H.MonadTest m
+         => MonadIO m
+         => Show a
+         => H.TestT IO a
+         -> m (Either H.Failure a, H.Journal)
+runTestT testt = withFrozenCallStack $ do
+  (res, log') <- H.evalIO . runWriterT . runExceptT . H.unTest $ testt
+  H.noteShow_ res
+  H.noteShow_ log'
+  pure (res, log')
+
+
+