Scheduler.hs

-- | Execute multiple build steps on multiple 'Flavor's of a build, with live
-- progress and output updates.
--
-- The scheduler maintains a job queue of flavors (not steps), so that when it
-- starts work on a flavor, it goes through all of its steps. The scheduler can
-- be told to reprioritize a flavor that hasn't been started yet.
--
-- The scheduler is given an 'Array' of 'Flavor's, and subsequent communication
-- refers to flavors from this array by their index.
module Cabal.Matrix.Scheduler
  ( SchedulerConfig(..)
  , mkCabalArgs
  , FlavorIndex
  , startScheduler
  , SchedulerMessage(..)
  , SchedulerHandle
  , signalScheduler
  , SchedulerSignal(..)
  ) where

import Cabal.Matrix.CabalArgs
import Cabal.Matrix.ProcessRunner
import Control.Concurrent
import Control.Exception.Safe
import Control.Monad
import Data.ByteString (ByteString)
import Data.Foldable
import Data.List
import Data.Map.Strict (Map)
import Data.Map.Strict qualified as Map
import Data.Primitive
import Data.Set (Set)
import Data.Set qualified as Set
import Data.Text (Text)
import System.Directory
import System.Exit


data SchedulerConfig = SchedulerConfig
  { jobs :: Int
    -- ^ Number of flavors that could be building in parallel at a time.
  , options :: [Text]
    -- ^ Options to use in all builds.
  , targets :: [Text]
  , cabalExecutable :: FilePath
  , mode :: CabalMode
  , steps :: PerCabalStep Bool
    -- ^ Which build steps to run or skip.
  }

-- | Index into 'flavors' of 'SchedulerInput'
type FlavorIndex = Int

data SchedulerMessage
  = OnStepStarted
    { flavorIndex :: FlavorIndex
    , step :: CabalStep
    } -- ^ Either the first message for this flavor, or sequenced after
      -- 'OnStepFinished' for this flavor.
  | OnStepFinished
    { flavorIndex :: FlavorIndex
    , step :: CabalStep
    , exitCode :: ExitCode
    } -- ^ Always sequenced after 'OnOutput' for this flavor.
  | OnOutput
    { flavorIndex :: FlavorIndex
    , step :: CabalStep
    , channel :: OutputChannel
    , output :: ByteString
    } -- ^ Sequenced after 'OnStepStarted' for this flavor.
  | OnDone -- ^ Sequenced after all events.
  deriving stock (Show)

data SchedulerSignal
  = InterruptFlavor
    { flavorIndex :: FlavorIndex
    } -- ^ Try to interrupt the current step (if any) with SIGINT and don't
      -- start subsequent steps (if any).
  | TerminateFlavor
    { flavorIndex :: FlavorIndex
    } -- ^ Try to interrupt the current step (if any) with SIGTERM and don't
      -- start subsequent steps (if any).
  | PrioritizeFlavor
    { flavorIndex :: FlavorIndex
    } -- ^ Move the given flavor the front of the priority queue, if it hasn't
      -- been started yet
  | RestartFlavor
    { flavorIndex :: FlavorIndex
    }

data SchedulerHandle = SchedulerHandle
  { state :: MVar SchedulerState
  , queueFilled :: MVar ()
  }

data SchedulerState = SchedulerState
  { processes :: Map FlavorIndex ProcessHandle
  , stopRequested :: Set FlavorIndex
  , queue :: [FlavorIndex]
    -- ^ INVARIANT: disjoint from 'stopRequested' and from keys of 'processes'
  }

mkCabalArgs :: SchedulerConfig -> CabalStep -> Flavor -> CabalArgs
mkCabalArgs input step flavor = CabalArgs
  { cabalExecutable = input.cabalExecutable
  , step
  , mode = input.mode
  , options = input.options
  , targets = input.targets
  , flavor
  }

startScheduler
  :: SchedulerConfig
  -> Array Flavor
  -> (SchedulerMessage -> IO ())
  -> IO SchedulerHandle
startScheduler input flavors cb = do
  sem <- newQSem input.jobs
  stateVar <- newMVar SchedulerState
    { processes = Map.empty
    , stopRequested = Set.empty
    , queue = [0 .. sizeofArray flavors - 1]
    }
  queueVar <- newEmptyMVar
  let
    waitForNext :: IO ()
    waitForNext = do
      waitQSem sem
      done <- modifyMVar stateVar \state -> case state.queue of
        [] -> do
          doneWithFlavor
          pure (state, True)
        flavorIndex:queue' -> do
          mProcess <- startSteps flavorIndex
            (filter (indexCabalStep input.steps) [minBound..maxBound])
          pure
            ( state
              { processes
                = Map.alter (\_ -> mProcess) flavorIndex state.processes
              , queue = queue'
              }
            , False
            )

      if done
      then waitForDone
      else waitForNext

    doneWithFlavor :: IO ()
    doneWithFlavor = signalQSem sem

    startSteps :: FlavorIndex -> [CabalStep] -> IO (Maybe ProcessHandle)
    startSteps flavorIndex = \case
      [] -> Nothing <$ doneWithFlavor
      step:nextSteps -> Just <$> do
        stdoutClosed <- newEmptyMVar
        stderrClosed <- newEmptyMVar
        cb OnStepStarted { flavorIndex, step }
        let args = mkCabalArgs input step (indexArray flavors flavorIndex)
        for_ (environmentFilePath args) $ void . try @_ @IOError . removeFile
        startProcess (renderCabalArgs args)
          (reactStep flavorIndex step nextSteps stdoutClosed stderrClosed)

    reactStep
      :: FlavorIndex
      -> CabalStep
      -> [CabalStep]
      -> MVar ()
      -> MVar ()
      -> ProcessMessage
      -> IO ()
    reactStep flavorIndex step nextSteps stdoutClosed stderrClosed = \case
      OnProcessOutput channel output -> cb OnOutput
        { flavorIndex
        , step
        , channel
        , output
        }
      OnChannelClosed Stdout -> putMVar stdoutClosed ()
      OnChannelClosed Stderr -> putMVar stderrClosed ()
      OnProcessExit exitCode -> do
        -- Synchronize OnStepFinished to be sequenced after OnOutput.
        -- Perhaps ProcessRunner should do this.
        takeMVar stdoutClosed
        takeMVar stderrClosed
        cb OnStepFinished { flavorIndex, step, exitCode }
        modifyMVar_ stateVar \state -> case exitCode of
          ExitSuccess
            -- It's possible that an interrupt signal has been received after
            -- the process has already exited (successfully), but before we got
            -- the 'OnProcessExit' message. So we must check this flag to see if
            -- we shouldn't start subsequent steps.
            | flavorIndex `Set.notMember` state.stopRequested
            -> do
              mProcess <- startSteps flavorIndex nextSteps
              pure state
                { processes
                  = Map.alter (\_ -> mProcess) flavorIndex state.processes
                }
          _ -> do
            doneWithFlavor
            pure state
              { processes = Map.delete flavorIndex state.processes
              , stopRequested = Set.delete flavorIndex state.stopRequested
              }

    waitForDone :: IO ()
    waitForDone = do
      -- Once the queue has become empty, wait for all running jobs to finish
      replicateM_ input.jobs $ waitQSem sem
      cb OnDone
      -- Restart the process if someone touches the queueVar
      takeMVar queueVar
      replicateM_ input.jobs $ signalQSem sem
      waitForNext

  _ <- forkIO waitForNext
  pure $ SchedulerHandle
    { state = stateVar
    , queueFilled = queueVar
    }

signalScheduler :: SchedulerHandle -> SchedulerSignal -> IO ()
signalScheduler hdl = \case
  InterruptFlavor{ flavorIndex } -> modifyMVar_ hdl.state \state -> if
    | flavorIndex `elem` state.queue
    -> pure state { queue = delete flavorIndex state.queue }
    | Just processHdl <- Map.lookup flavorIndex state.processes
    -> do
      signalProcess processHdl SignalInterrupt
      pure state { stopRequested = Set.insert flavorIndex state.stopRequested }
    | otherwise
    -> pure state
  TerminateFlavor{ flavorIndex } -> modifyMVar_ hdl.state \state -> if
    | flavorIndex `elem` state.queue
    -> pure state { queue = delete flavorIndex state.queue }
    | Just processHdl <- Map.lookup flavorIndex state.processes
    -> do
      signalProcess processHdl SignalTerminate
      pure state { stopRequested = Set.insert flavorIndex state.stopRequested }
    | otherwise
    -> pure state
  PrioritizeFlavor{ flavorIndex } -> modifyMVar_ hdl.state \state -> if
    | flavorIndex `elem` state.queue
    -> pure state { queue = flavorIndex : delete flavorIndex state.queue }
      -- ^ TODO: leak? we accumulate 'delete' thunks?
    | otherwise
    -> pure state
  RestartFlavor{ flavorIndex } -> do
    modifyMVar_ hdl.state \state -> if
      | flavorIndex `notElem` state.queue
      , flavorIndex `Map.notMember` state.processes
      -> pure state { queue = flavorIndex : state.queue }
      | otherwise
      -> pure state
    void $ tryPutMVar hdl.queueFilled ()