Add TestWatchdog and Tripwire with their tests

hedgehog-extras.cabal

@@ -131,18 +131,26 @@ library
                         Hedgehog.Extras.Test.MonadAssertion
                         Hedgehog.Extras.Test.Network
                         Hedgehog.Extras.Test.Process
+                        Hedgehog.Extras.Test.TestWatchdog
+                        Hedgehog.Extras.Test.Tripwire
 
 test-suite hedgehog-extras-test
   import:               base, project-config,
                         hedgehog,
                         hedgehog-extras,
                         network,
+                        process,
+                        resourcet,
                         tasty,
                         tasty-hedgehog,
+                        transformers,
+                        time,
   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

@@ -2,9 +2,11 @@ module Hedgehog.Extras.Test
   ( module X
   ) where
 
-import Hedgehog.Extras.Test.Base as X
-import Hedgehog.Extras.Test.Concurrent as X
-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.Base as X
+import           Hedgehog.Extras.Test.Concurrent as X
+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
+import           Hedgehog.Extras.Test.Tripwire as X

src/Hedgehog/Extras/Test/Concurrent.hs

@@ -66,14 +66,17 @@ __Don't use concurrency abstractions from this module, when you need to aggregat
 -}
 module Hedgehog.Extras.Test.Concurrent
   ( threadDelay
+  , asyncRegister_
   -- * Re-exports of concurrency abstractions from @lifted-base@
   , module Control.Concurrent.Async.Lifted
+  , module Control.Concurrent.MVar.Lifted
   , module System.Timeout.Lifted
   ) where
 
 import           Control.Applicative
 import           Control.Concurrent.Async.Lifted
 import qualified Control.Concurrent.Lifted as IO
+import           Control.Concurrent.MVar.Lifted
 import           Control.Monad.Base
 import           Control.Monad.IO.Class
 import           Control.Monad.Trans.Control
@@ -85,13 +88,28 @@ import           System.IO (IO)
 import           System.Timeout.Lifted
 import qualified UnliftIO
 
+import           Control.Monad
+import           Control.Monad.Catch (MonadCatch)
+import           GHC.Stack
 import           Hedgehog
 import qualified Hedgehog as H
 
 -- | Delay the thread by 'n' milliseconds.
-threadDelay :: (MonadTest m, MonadIO m) => Int -> m ()
+threadDelay :: (HasCallStack, MonadTest m, MonadIO m) => Int -> m ()
 threadDelay n = GHC.withFrozenCallStack . H.evalIO $ IO.threadDelay n
 
+-- | Runs an action in background, and registers its cancellation to 'MonadResource'.
+asyncRegister_ :: HasCallStack
+               => MonadTest m
+               => MonadResource m
+               => MonadCatch m
+               => IO a -- ^ Action to run in background
+               -> m ()
+asyncRegister_ act = void . H.evalM $ allocate (async act) cleanUp
+  where
+    cleanUp :: Async a -> IO ()
+    cleanUp a = cancel a >> void (link a)
+
 instance MonadBase IO (ResourceT IO) where
   liftBase = liftIO
 

src/Hedgehog/Extras/Test/TestWatchdog.hs

@@ -0,0 +1,147 @@
+{-# 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 a 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
+  , WatchdogException(..)
+  ) 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
+  { 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

src/Hedgehog/Extras/Test/Tripwire.hs

@@ -0,0 +1,61 @@
+{-# LANGUAGE RankNTypes #-}
+module Hedgehog.Extras.Test.Tripwire
+  ( Tripwire
+  , makeTripwire
+  , triggerTripwire
+  , resetTripwire
+  , checkTripwire
+  , isTriggered
+  ) where
+
+import           Control.Monad.IO.Class
+import           GHC.Stack
+
+import           Control.Concurrent.MVar
+import           Control.Monad
+import           Hedgehog (MonadTest)
+import qualified Hedgehog.Extras.Test.Base as H
+import qualified Hedgehog.Internal.Property as H
+import           Prelude
+
+-- | Represents a tripwire which can be triggered only once
+newtype Tripwire = Tripwire (MVar CallStack)
+
+-- | Creates a new tripwire
+makeTripwire :: MonadIO m => m Tripwire
+makeTripwire = Tripwire <$> liftIO newEmptyMVar
+
+-- | Triggers the tripwire and registers the place of the first trigger. Idempotent.
+-- Does not do do anything besides just registering the place where this function is called.
+triggerTripwire :: HasCallStack
+                => MonadIO m
+                => Tripwire
+                -> m ()
+triggerTripwire (Tripwire mv) = withFrozenCallStack $
+  void . liftIO $ tryPutMVar mv callStack
+
+-- | Restore tripwire to initial non triggered state
+resetTripwire :: MonadIO m
+              => Tripwire
+              -> m ()
+resetTripwire (Tripwire mv) = liftIO $ void $ tryTakeMVar mv
+
+-- | Check if the tripwire is triggered. Return the first trigger location.
+isTriggered :: MonadIO m
+            => Tripwire
+            -> m (Maybe CallStack)
+isTriggered (Tripwire mv) = liftIO $ tryReadMVar mv
+
+-- | Fails the test if the tripwire was triggered. Prints the callstack where the tripwire was triggered.
+checkTripwire :: HasCallStack
+              => MonadTest m
+              => MonadIO m
+              => Tripwire
+              -> m ()
+checkTripwire = withFrozenCallStack $ do
+  isTriggered >=> void . mapM
+    (\cs -> do
+      H.note_ $ "Tripwire has been tripped at: " <> prettyCallStack cs
+      H.failure
+    )
+