paralleln: add internal lanes so test timer profiles are cleaner

sandwich/CHANGELOG.md

@@ -3,6 +3,10 @@
 ## Unreleased
 
 * TUI: be able to press 'S' to open the speedscope profile in a browser.
+* `parallelN` now limits with a pool of lanes instead of a semaphore, so a run produces N test timer profiles rather than one per test. The bound still covers the whole subtree, including nested `parallel` nodes.
+* BREAKING CHANGE: `parallelN` introduces `parallelLanes` instead of `parallelSemaphore`, and requires `HasBaseContext`. Claim a lane with `withParallelLane` rather than taking the `QSem` yourself; unlike the semaphore, it's re-entrant.
+* Add `withParallelLanes`, `withParallelLanesFromArgs` and `takeParallelLane`, for limiting a spec tree you can't wrap with `parallelN` (such as one from `getSpecFromFolder`).
+* Don't leave children running when a `parallel` node stops waiting on them early.
 
 ## 0.3.1.0
 

sandwich/src/Test/Sandwich.hs

@@ -46,7 +46,21 @@ module Test.Sandwich (
   , aroundEach
 
   -- * Parallel nodes
-  , module Test.Sandwich.ParallelN
+  , parallelN
+  , parallelN'
+  , parallelNFromArgs
+  , parallelNFromArgs'
+  , defaultParallelNodeOptions
+
+  -- ** Lower-level
+  , withParallelLanes
+  , withParallelLanesFromArgs
+  , withParallelLane
+  , takeParallelLane
+
+  , ParallelLanes
+  , parallelLanes
+  , HasParallelLanes
 
   -- * Timing
   --

sandwich/src/Test/Sandwich/Internal/Running.hs

@@ -114,13 +114,16 @@ baseContextFromOptions options@(Options {..}) = do
 
   let errorSymlinksDir = (</> "errors") <$> runRoot
   whenJust errorSymlinksDir $ createDirectoryIfMissing True
+
+  timingProfile <- newTVarIO (TimingProfile defaultProfileName [])
+
   return $ BaseContext {
     baseContextPath = mempty
     , baseContextOptions = options
     , baseContextRunRoot = runRoot
     , baseContextErrorSymlinksDir = errorSymlinksDir
     , baseContextOnlyRunIds = Nothing
-    , baseContextTestTimerProfile = defaultProfileName
+    , baseContextTimingProfile = timingProfile
     , baseContextTestTimer = testTimer
     , baseContextRunId = runId
     }

sandwich/src/Test/Sandwich/Interpreters/StartTree.hs

@@ -135,9 +135,13 @@ startTree node@(RunNodeIntroduce {..}) ctx' = do
             Right intro -> do
               -- Special hack to modify the test timer profile via an introduce, without needing to track it everywhere.
               -- It would be better to track the profile at the type level
-              let ctxFinal = case cast intro of
-                    Just (TestTimerProfile t) -> modifyBaseContext ctx (\bc -> bc { baseContextTestTimerProfile = t })
-                    Nothing -> ctx
+              ctxFinal <- case cast intro of
+                Just (TestTimerProfile t) -> do
+                  -- A fresh cell, so we rename the profile for our children only.
+                  current <- readTVarIO (baseContextTimingProfile (getBaseContext ctx))
+                  timingProfile <- newTVarIO (current { timingProfileName = t })
+                  return $ modifyBaseContext ctx (\bc -> bc { baseContextTimingProfile = timingProfile })
+                Nothing -> return ctx
 
               void $ runNodesSequentially runNodeChildrenAugmented ((LabelValue intro) :> ctxFinal)
       )
@@ -278,13 +282,19 @@ runInAsync :: (HasBaseContext context, MonadIO m) => RunNode context -> context
 runInAsync node ctx action = do
   let RunNodeCommonWithStatus {..} = runNodeCommon node
   let bc@(BaseContext {..}) = getBaseContext ctx
-  let timerFn = if runTreeRecordTime then timeAction' (getTestTimer bc) baseContextTestTimerProfile (T.pack runTreeLabel) else id
   let asyncName = T.pack [i|node #{runTreeId}, #{runTreeLabel}|]
   startTime <- liftIO getCurrentTime
   mvar <- liftIO newEmptyMVar
   myAsync <- liftIO $ managedAsyncWithUnmask baseContextRunId asyncName $ \unmask -> do
     flip withException (recordExceptionInStatus runTreeStatus) $ unmask $ do
       readMVar mvar
+      -- Resolve the timing profile now rather than when this node was created, since a lane may
+      -- have been claimed for us in the meantime.
+      timerFn <- case runTreeRecordTime of
+        False -> pure id
+        True -> do
+          profile <- currentTestTimerProfile bc
+          pure $ timeAction' (getTestTimer bc) profile (T.pack runTreeLabel)
       (result, extraTimingInfo) <- timerFn action
       endTime <- liftIO getCurrentTime
       liftIO $ atomically $ writeTVar runTreeStatus $ Done startTime (setupFinishTime extraTimingInfo) (teardownStartTime extraTimingInfo) endTime result
@@ -361,22 +371,29 @@ runNodesSequentially children ctx =
 -- | Run a list of children concurrently, cancelling everything on async exception
 runNodesConcurrently :: forall context. HasBaseContext context => RunNodeCommon -> [RunNode context] -> context -> IO [Result]
 runNodesConcurrently (RunNodeCommonWithStatus {runTreeLabel, runTreeId}) children ctx =
-  flip withException (\(e :: SomeAsyncException) -> cancelAllChildrenWith children e) $
-    mapM wait =<< sequence [startTree child (modifyTimingProfile ix ctx)
-                           | (child, ix) <- L.zip runnableChildren [0..]]
+  flip withException (\(e :: SomeAsyncException) -> cancelAllChildrenWith children e) $ do
+    asyncs <- sequence [startTree child =<< childContext ix
+                       | (child, ix) <- L.zip runnableChildren [0..]]
+    -- If we stop waiting early because one child threw, take the others down with us;
+    -- the handler above only fires for async exceptions.
+    mapM wait asyncs `onException` mapM_ cancel asyncs
   where
     runnableChildren = L.filter (shouldRunChild ctx) children
 
     leftPadWithZeros :: Int -> String
     leftPadWithZeros num = L.replicate (L.length (show (L.length runnableChildren)) - L.length (show num)) '0' <> show num
 
-    modifyTimingProfile :: Int -> context -> context
-    modifyTimingProfile n = flip modifyBaseContext (modifyTimingProfile' n)
+    -- Give each child its own test timer profile, since sharing one would mess up the nesting of the
+    -- profile's frames. Children keep any lanes we're already holding, so they don't claim a second
+    -- one and deadlock.
+    childContext :: Int -> IO context
+    childContext n = do
+      current <- readTVarIO (baseContextTimingProfile (getBaseContext ctx))
+      timingProfile <- newTVarIO (current { timingProfileName = timingProfileName current <> suffix n })
+      return $ modifyBaseContext ctx $ \bc -> bc { baseContextTimingProfile = timingProfile }
 
-    modifyTimingProfile' :: Int -> BaseContext -> BaseContext
-    modifyTimingProfile' n bc@(BaseContext {..}) = bc {
-      baseContextTestTimerProfile = baseContextTestTimerProfile <> [i|-#{runTreeLabel}-#{runTreeId}-#{leftPadWithZeros n}|]
-      }
+    suffix :: Int -> T.Text
+    suffix n = [i|-#{runTreeLabel}-#{runTreeId}-#{leftPadWithZeros n}|]
 
 markAllChildrenWithResult :: (MonadIO m, HasBaseContext context') => [RunNode context] -> context' -> Result -> m ()
 markAllChildrenWithResult children baseContext status = do
@@ -506,8 +523,12 @@ recordExceptionInStatus status e = do
     _ -> Done endTime Nothing Nothing endTime ret
 
 timed :: forall m a s l t. (MonadUnliftIO m) => RunNodeCommonWithStatus s l t -> Bool -> BaseContext -> String -> m a -> m (a, UTCTime, UTCTime)
-timed _rnc recordTime bc@(BaseContext {..}) label action = do
-  let timerFn = if recordTime then timeAction' (getTestTimer bc) baseContextTestTimerProfile (T.pack label) else id
+timed _rnc recordTime bc@(BaseContext {}) label action = do
+  timerFn <- case recordTime of
+    False -> pure id
+    True -> do
+      profile <- currentTestTimerProfile bc
+      pure $ timeAction' (getTestTimer bc) profile (T.pack label)
 
   startTime <- liftIO getCurrentTime
   ret <- timerFn action

sandwich/src/Test/Sandwich/ParallelN.hs

@@ -4,95 +4,223 @@
 {-# LANGUAGE TypeApplications #-}
 {-# LANGUAGE TypeOperators #-}
 
--- | Wrapper around 'parallel' for limiting the threads using a semaphore.
+-- | Limiting how much of a test tree runs at once.
 
 module Test.Sandwich.ParallelN (
+  -- * Limiting a parallel node
   parallelN
   , parallelN'
 
   , parallelNFromArgs
   , parallelNFromArgs'
 
+  -- * Lower-level
+  , withParallelLanes
+  , withParallelLanesFromArgs
+
+  , withParallelLane
+  , takeParallelLane
+
   , defaultParallelNodeOptions
 
   -- * Types
-  , parallelSemaphore
-  , HasParallelSemaphore
+  , ParallelLanes
+  , parallelLanes
+  , HasParallelLanes
   ) where
 
 import Control.Concurrent.QSem
+import Control.Concurrent.STM (retry)
 import Control.Monad
 import Control.Monad.IO.Class
 import Control.Monad.IO.Unlift
+import Control.Monad.Reader
+import qualified Data.List as L
+import Data.String.Interpolate
+import qualified Data.Text as T
 import Test.Sandwich.Contexts
+import Test.Sandwich.TestTimer
 import Test.Sandwich.Types.ArgParsing
 import Test.Sandwich.Types.RunTree
 import Test.Sandwich.Types.Spec
 import UnliftIO.Exception
+import UnliftIO.STM
 
 
 -- * Types
 
-parallelSemaphore :: Label "parallelSemaphore" QSem
-parallelSemaphore = Label
+-- | A pool of lanes. Each lane is held by one part of the tree at a time, and has a test timer
+-- profile of its own.
+data ParallelLanes = ParallelLanes {
+  parallelLanesSource :: TimingLaneSource
+  -- | The bound. Gating on this first hands lanes out in wait order.
+  , parallelLanesSem :: QSem
+  , parallelLanesFree :: TVar [Int]
+  , parallelLanesProfileNames :: [T.Text]
+  }
+
+parallelLanes :: Label "parallelLanes" ParallelLanes
+parallelLanes = Label
 
-type HasParallelSemaphore context = HasLabel context "parallelSemaphore" QSem
+type HasParallelLanes context = HasLabel context "parallelLanes" ParallelLanes
 
 defaultParallelNodeOptions :: NodeOptions
 defaultParallelNodeOptions = defaultNodeOptions { nodeOptionsVisibilityThreshold = 70 }
 
+-- Options for nodes that only introduce or claim lanes. Timing them would put a frame outside
+-- every lane, forcing a profile of its own.
+laneNodeOptions :: NodeOptions
+laneNodeOptions = defaultNodeOptions {
+  nodeOptionsRecordTime = False
+  , nodeOptionsCreateFolder = False
+  , nodeOptionsVisibilityThreshold = 125
+  }
+
 -- * Functions
 
--- | Wrapper around 'parallel'. Introduces a semaphore to limit the parallelism to N threads.
+-- | Like 'parallel', but limits the parallelism to N tests at a time.
+--
+-- Introduces a pool of N lanes, one of which each test claims while it runs. The pool is shared by
+-- the whole subtree, so nested 'parallel' nodes are limited too, and each lane is a test timer
+-- profile, so you get N profiles rather than one per test.
 parallelN :: (
-  MonadUnliftIO m
-  ) => Int -> SpecFree (LabelValue "parallelSemaphore" QSem :> context) m () -> SpecFree context m ()
+  MonadUnliftIO m, HasBaseContext context
+  )
+  -- | Number of lanes
+  => Int
+  -> SpecFree (LabelValue "parallelLanes" ParallelLanes :> context) m ()
+  -> SpecFree context m ()
 parallelN = parallelN' defaultParallelNodeOptions
 
 parallelN' :: (
-  MonadUnliftIO m
+  MonadUnliftIO m, HasBaseContext context
   )
   -- | Node options
   => NodeOptions
-  -- | Number of threads
+  -- | Number of lanes
   -> Int
-  -> SpecFree (LabelValue "parallelSemaphore" QSem :> context) m ()
+  -> SpecFree (LabelValue "parallelLanes" ParallelLanes :> context) m ()
   -> SpecFree context m ()
-parallelN' nodeOptions n = parallelN'' nodeOptions (liftIO $ newQSem n)
+parallelN' nodeOptions n = parallelN'' nodeOptions (pure n)
 
--- | Same as 'parallelN', but extracts the semaphore size from the command line options.
+-- | Same as 'parallelN', but extracts the number of lanes from the command line options.
 parallelNFromArgs :: forall context a m. (
-  MonadUnliftIO m, HasCommandLineOptions context a
+  MonadUnliftIO m, HasBaseContext context, HasCommandLineOptions context a
   )
-  -- | Callback to extract the semaphore size
+  -- | Callback to extract the number of lanes
   => (CommandLineOptions a -> Int)
-  -> SpecFree (LabelValue "parallelSemaphore" QSem :> context) m ()
+  -> SpecFree (LabelValue "parallelLanes" ParallelLanes :> context) m ()
   -> SpecFree context m ()
 parallelNFromArgs = parallelNFromArgs' @context @a defaultParallelNodeOptions
 
 parallelNFromArgs' :: forall context a m. (
-  MonadUnliftIO m, HasCommandLineOptions context a
+  MonadUnliftIO m, HasBaseContext context, HasCommandLineOptions context a
   )
   -- | Node options
   => NodeOptions
-  -- | Callback to extract the semaphore size
+  -- | Callback to extract the number of lanes
   -> (CommandLineOptions a -> Int)
-  -> SpecFree (LabelValue "parallelSemaphore" QSem :> context) m ()
+  -> SpecFree (LabelValue "parallelLanes" ParallelLanes :> context) m ()
   -> SpecFree context m ()
-parallelNFromArgs' nodeOptions getParallelism = parallelN'' nodeOptions f
-  where
-    f = getContext commandLineOptions >>= (liftIO . newQSem) . getParallelism
+parallelNFromArgs' nodeOptions getParallelism =
+  parallelN'' nodeOptions (getParallelism <$> getContext commandLineOptions)
 
 parallelN'' :: (
-  MonadUnliftIO m
+  MonadUnliftIO m, HasBaseContext context
   )
   => NodeOptions
-  -> ExampleT context m QSem
-  -> SpecFree (LabelValue "parallelSemaphore" QSem :> context) m ()
+  -> ExampleT context m Int
+  -> SpecFree (LabelValue "parallelLanes" ParallelLanes :> context) m ()
   -> SpecFree context m ()
-parallelN'' nodeOptions makeQSem children = introduce "Introduce parallel semaphore" parallelSemaphore makeQSem (const $ return ()) $
-  parallel' nodeOptions $ aroundEach "Take parallel semaphore" claimRunSlot children
+parallelN'' nodeOptions getLanes children =
+  withParallelLanes'' getLanes $
+    parallel' nodeOptions $
+      aroundEach' Nothing laneNodeOptions "Take parallel lane" (withParallelLane . void) children
+
+-- | Introduce a pool of N lanes, to be claimed with 'withParallelLane' or 'takeParallelLane'.
+--
+-- Use this when the tests to limit aren't in one place you can wrap with 'parallelN', such as ones
+-- from 'Test.Sandwich.TH.getSpecFromFolder'.
+withParallelLanes :: (
+  MonadIO m, HasBaseContext context
+  )
+  -- | Number of lanes
+  => Int
+  -> SpecFree (LabelValue "parallelLanes" ParallelLanes :> context) m ()
+  -> SpecFree context m ()
+withParallelLanes n = withParallelLanes'' (pure n)
+
+-- | Same as 'withParallelLanes', but extracts the number of lanes from the command line options.
+withParallelLanesFromArgs :: forall context a m. (
+  MonadIO m, HasBaseContext context, HasCommandLineOptions context a
+  )
+  -- | Callback to extract the number of lanes
+  => (CommandLineOptions a -> Int)
+  -> SpecFree (LabelValue "parallelLanes" ParallelLanes :> context) m ()
+  -> SpecFree context m ()
+withParallelLanesFromArgs getParallelism =
+  withParallelLanes'' (getParallelism <$> getContext commandLineOptions)
+
+withParallelLanes'' :: (
+  MonadIO m, HasBaseContext context
+  )
+  => ExampleT context m Int
+  -> SpecFree (LabelValue "parallelLanes" ParallelLanes :> context) m ()
+  -> SpecFree context m ()
+withParallelLanes'' getLanes =
+  introduce' laneNodeOptions "Introduce parallel lanes" parallelLanes alloc (const $ return ())
   where
-    claimRunSlot f = do
-      s <- getContext parallelSemaphore
-      bracket_ (liftIO $ waitQSem s) (liftIO $ signalQSem s) (void f)
+    alloc = do
+      numLanes <- max 1 <$> getLanes
+      source <- newTimingLaneSource
+      sem <- liftIO $ newQSem numLanes
+      free <- newTVarIO [0 .. numLanes - 1]
+      -- Name the lanes after the profile we're in, so lanes introduced inside another lane (or
+      -- inside one branch of a parallel node) don't collide with ones introduced elsewhere.
+      baseProfile <- asks getBaseContext >>= currentTestTimerProfile
+      let names = [baseProfile <> [i|-lane-#{leftPadWithZeros numLanes lane}|]
+                  | lane <- [0 .. numLanes - 1]]
+      return $ ParallelLanes source sem free names
+
+    leftPadWithZeros :: Int -> Int -> String
+    leftPadWithZeros total num =
+      L.replicate (L.length (show (total - 1)) - L.length (show num)) '0' <> show num
+
+-- | Claim a lane, run the action, and release it. Blocks until a lane is free. Everything the
+-- action runs, at any depth, is timed under the lane's profile.
+--
+-- A no-op if this part of the tree already holds a lane from the same pool, since claiming a
+-- second one could deadlock.
+withParallelLane :: (
+  MonadUnliftIO m, HasBaseContextMonad context m, HasParallelLanes context
+  ) => m a -> m a
+withParallelLane action = do
+  pool <- getContext parallelLanes
+  inTimingLane (parallelLanesSource pool) >>= \case
+    True -> action
+    False -> bracket (claimLane pool) (releaseLane pool) $ \lane ->
+      withTimingLane (parallelLanesSource pool) (parallelLanesProfileNames pool !! lane) action
+
+-- | 'withParallelLane' as a spec node, designed for use with
+-- 'Test.Sandwich.TH.getSpecIndividualSpecHooks'.
+takeParallelLane :: (
+  MonadUnliftIO m, HasBaseContext context, HasParallelLanes context
+  )
+  -- | Ignored
+  => FilePath
+  -> SpecFree context m ()
+  -> SpecFree context m ()
+takeParallelLane _ = around' laneNodeOptions "Take parallel lane" (withParallelLane . void)
+
+claimLane :: (MonadIO m) => ParallelLanes -> m Int
+claimLane (ParallelLanes {parallelLanesSem, parallelLanesFree}) = liftIO $ do
+  waitQSem parallelLanesSem
+  flip onException (signalQSem parallelLanesSem) $ atomically $ readTVar parallelLanesFree >>= \case
+    -- Can't happen: the semaphore already limits us to the number of lanes.
+    [] -> retry
+    (lane:rest) -> writeTVar parallelLanesFree rest >> return lane
+
+releaseLane :: (MonadIO m) => ParallelLanes -> Int -> m ()
+releaseLane (ParallelLanes {parallelLanesSem, parallelLanesFree}) lane = liftIO $ do
+  atomically $ modifyTVar' parallelLanesFree (lane :)
+  signalQSem parallelLanesSem