content: -- | 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
  
  newtype SchedulerHandle = SchedulerHandle (MVar SchedulerState)
  
  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
    mvar <- newMVar SchedulerState
      { processes = Map.empty
      , stopRequested = Set.empty
      , queue = [0 .. sizeofArray flavors - 1]
      }
    let
      waitForNext :: IO ()
      waitForNext = do
        waitQSem sem
        done <- modifyMVar mvar \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.update (\_ -> 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 }
          case exitCode of
            ExitFailure _ -> doneWithFlavor
            ExitSuccess -> modifyMVar_ mvar \state
              -- 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.
              -> if flavorIndex `Set.member` state.stopRequested
                then do
                  doneWithFlavor
                  pure state
                    { processes = Map.delete flavorIndex state.processes
                    , stopRequested = Set.delete flavorIndex state.stopRequested
                    }
                else do
                  mProcess <- startSteps flavorIndex nextSteps
                  pure state
                    { processes
                      = Map.update (\_ -> mProcess) flavorIndex state.processes
                    }
  
      waitForDone :: IO ()
      waitForDone = do
        -- Once the queue has become empty, wait for all running jobs to finish
        replicateM_ input.jobs $ waitQSem sem
        cb OnDone
  
    _ <- forkIO waitForNext
    pure $ SchedulerHandle mvar
  
  signalScheduler :: SchedulerHandle -> SchedulerSignal -> IO ()
  signalScheduler (SchedulerHandle mvar) = \case
    InterruptFlavor{ flavorIndex } -> modifyMVar_ mvar \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_ mvar \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_ mvar \state -> if
      | flavorIndex `elem` state.queue
      -> pure state { queue = flavorIndex : delete flavorIndex state.queue }
        -- ^ TODO: leak? we accumulate 'delete' thunks?
      | otherwise
      -> pure state
  