test-framework: Fix exception handling

1 files changed, 57 insertions, 33 deletions

diff --git a/test-framework/Test/Framework.hs b/test-framework/Test/Framework.hs
index 5ceb866..80711b2 100644
--- a/test-framework/Test/Framework.hs
+++ b/test-framework/Test/Framework.hs
@@ -4,6 +4,7 @@
 {-# LANGUAGE ImplicitParams #-}
 {-# LANGUAGE ImportQualifiedPost #-}
 {-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE TupleSections #-}
 module Test.Framework (
   TestTree,
@@ -20,11 +21,11 @@ module Test.Framework (
   TestName,
 ) where
 
-import Control.Concurrent (setNumCapabilities, forkIO, killThread, forkOn)
+import Control.Concurrent (setNumCapabilities, forkIO, forkOn, killThread)
 import Control.Concurrent.MVar
 import Control.Concurrent.STM
-import Control.Exception (finally)
-import Control.Monad (forM, when, forM_, replicateM_)
+import Control.Exception (SomeException, throw, try, throwIO)
+import Control.Monad (forM, when, forM_)
 import Control.Monad.IO.Class
 import Data.IORef
 import Data.List (isInfixOf, intercalate)
@@ -49,6 +50,9 @@ import Hedgehog.Internal.Runner qualified as H
 import Hedgehog.Internal.Seed qualified as H.Seed
 
 
+-- TODO: with GHC 9.12 we have tryWithContext and rethrowIO, which is better for rethrowing exceptions
+
+
 type TestName = String
 
 data TestTree
@@ -107,6 +111,11 @@ filterTree (Options { optsPattern = pat }) = go []
 
     renderPath comps = "^" ++ intercalate "/" (reverse comps) ++ "$"
 
+treeNumTests :: TestTree -> Int
+treeNumTests (Group _ _ ts) = sum (map treeNumTests ts)
+treeNumTests (Resource _ _ _ fun) = treeNumTests (fun undefined)
+treeNumTests HP{} = 1
+
 computeMaxLen :: TestTree -> Int
 computeMaxLen = go 0
   where
@@ -220,7 +229,7 @@ runTests options = \tree' ->
              else isJust <$> runTreeSeq 0 [] tree
       stats <- readIORef statsRef
       endtm <- getCurrentTime
-      let ?istty = isterm in printStats stats (diffUTCTime endtm starttm)
+      let ?istty = isterm in printStats (treeNumTests tree) stats (diffUTCTime endtm starttm)
       return (if success then ExitSuccess else ExitFailure 1)
 
 -- | Returns when all jobs in this tree have been scheduled. When all jobs are
@@ -275,31 +284,34 @@ runTreePar topmparregion revidxlist revpath toptree@Resource{} topoutvar = runRe
       let pathitem = '[' : show depth ++ "](" ++ inhname ++ ")"
           path = intercalate "/" (reverse (pathitem : revpath))
           idxlist = reverse revidxlist
-      -- outputConcurrent $ "! " ++ path ++ ": R Submitting\n"
       submitOrRunIn mparregion idxlist Nothing $ \makeRegion -> do
         setConsoleRegion makeRegion ('|' : path ++ " [R] making...")
 
-        -- outputConcurrent $ "! " ++ path ++ ": R Making\n"
-        value <- make  -- TODO: catch exceptions
-        -- outputConcurrent $ "! " ++ path ++ ": R Made\n"
-
-        -- outputConcurrent $ "! " ++ path ++ ": R Running subtree\n"
-        suboutvar <- newEmptyMVar
-        runResource (Just makeRegion) (depth + 1) (fun value) suboutvar  -- will consume makeRegion
-        -- outputConcurrent $ "! " ++ path ++ ": R Scheduled subtree\n"
+        evalue <- try make
+        case evalue of
+          Left (err :: SomeException) -> do
+            finishConsoleRegion makeRegion $
+              ansiRed ++ "Exception building resource at " ++ path ++ ":" ++ ansiReset ++ "\n" ++ show err
+            putMVar outvar False
+          Right value -> do
+            suboutvar <- newEmptyMVar
+            runResource (Just makeRegion) (depth + 1) (fun value) suboutvar  -- will consume makeRegion
 
-        _ <- forkIO $ do
-          success <- readMVar suboutvar
-          -- outputConcurrent $ "! " ++ path ++ ": R Subtree done, scheduling cleanup\n"
-          poolSubmit ?pool idxlist (Just outvar) $ do
-            cleanupRegion <- openConsoleRegion Linear
-            setConsoleRegion cleanupRegion ('|' : path ++ " [R] cleanup...")
-            -- outputConcurrent $ "! " ++ path ++ ": R Cleaning up\n"
-            cleanup value  -- TODO: catch exceptions
-            -- outputConcurrent $ "! " ++ path ++ ": R Cleanup done\n"
-            closeConsoleRegion cleanupRegion
-            return success
-        return ()
+            _ <- forkIO $ do
+              success <- readMVar suboutvar
+              poolSubmit ?pool idxlist Nothing $ do
+                cleanupRegion <- openConsoleRegion Linear
+                setConsoleRegion cleanupRegion ('|' : path ++ " [R] cleanup...")
+                eres <- try (cleanup value)
+                case eres of
+                  Left (err :: SomeException) -> do
+                    finishConsoleRegion cleanupRegion $
+                      ansiRed ++ "Exception cleaning up resource at " ++ path ++ ":" ++ ansiReset ++ "\n" ++ show err
+                    putMVar outvar False
+                  Right () -> do
+                    closeConsoleRegion cleanupRegion
+                    putMVar outvar success
+            return ()
     runResource mparregion _ tree outvar = runTreePar mparregion revidxlist revpath tree outvar
 
 runTreePar mparregion revidxlist revpath (HP name prop) outvar = do
@@ -407,15 +419,19 @@ renderResult report path timeTaken = do
       return (str ++ suffix)
     _ -> return str
 
-printStats :: (?istty :: Bool) => Stats -> NominalDiffTime -> IO ()
-printStats stats timeTaken
-  | statsOK stats == statsTotal stats = do
+printStats :: (?istty :: Bool) => Int -> Stats -> NominalDiffTime -> IO ()
+printStats numTests stats timeTaken
+  | statsOK stats == numTests = do
       putStrLn $ ansiGreen ++ "All " ++ show (statsTotal stats) ++
                  " tests passed." ++ prettyDuration True (realToFrac timeTaken) ++ ansiReset
+  | statsOK stats == statsTotal stats =
+      putStrLn $ ansiRed ++ "Failed (" ++ show (numTests - statsTotal stats) ++ " tests could not run)." ++
+                 prettyDuration True (realToFrac timeTaken) ++ ansiReset
   | otherwise =
       let nfailed = statsTotal stats - statsOK stats
-      in putStrLn $ ansiRed ++ "Failed " ++ show nfailed ++ " out of " ++ show (statsTotal stats) ++
-                    " tests." ++ prettyDuration True (realToFrac timeTaken) ++ ansiReset
+      in putStrLn $ ansiRed ++ "Failed " ++ show nfailed ++ " out of " ++ show (statsTotal stats) ++ " tests" ++
+                    (if statsTotal stats /= numTests then " (" ++ show (numTests - statsTotal stats) ++ " could not run)" else "") ++
+                    "." ++ prettyDuration True (realToFrac timeTaken) ++ ansiReset
 
 
 newtype WorkerPool k = WorkerPool (TVar (PQ.MinPQueue k (Terminate PoolJob)))
@@ -427,9 +443,17 @@ withWorkerPool :: Ord k => Int -> (WorkerPool k -> IO a) -> IO a
 withWorkerPool numWorkers k = do
   chan <- newTVarIO PQ.empty
   threads <- forM [0..numWorkers-1] (\i -> forkOn i (worker i chan))
-  k (WorkerPool chan) `finally` do
-    replicateM_ numWorkers (atomically $ writeTVar chan PQ.empty)
-    forM_ threads killThread
+  eres <- try (k (WorkerPool chan))
+  case eres of
+    Left (err :: SomeException) -> do
+      atomically $ writeTVar chan PQ.empty
+      forM_ threads killThread
+      throw err
+    Right res -> do
+      readTVarIO chan >>= \case
+        PQ.Empty -> return ()
+        _ -> throwIO (userError "withWorkerPool: computation exited before all jobs were handled")
+      return res
   where
     worker :: Ord k => Int -> TVar (PQ.MinPQueue k (Terminate PoolJob)) -> IO ()
     worker idx chan = do