content: {-# LANGUAGE CPP #-}
  module GHC.Driver.MakeAction
    ( MakeAction(..)
    , MakeEnv(..)
    , RunMakeM
    -- * Running the pipelines
    , runAllPipelines
    , runParPipelines
    , runSeqPipelines
    , runPipelines
    -- * Worker limit
    , WorkerLimit(..)
    , mkWorkerLimit
    , runWorkerLimit
    -- * Utility
    , withLoggerHsc
    , withParLog
    , withLocalTmpFS
    , withLocalTmpFSMake
    ) where
  
  import GHC.Prelude
  import GHC.Driver.DynFlags
  
  import GHC.Driver.Monad
  import GHC.Driver.Env
  import GHC.Driver.Errors.Types
  import GHC.Driver.Messager
  import GHC.Driver.MakeSem
  
  import GHC.Utils.Logger
  import GHC.Utils.TmpFs
  
  import Control.Concurrent ( newQSem, waitQSem, signalQSem, ThreadId, killThread, forkIOWithUnmask )
  import qualified GHC.Conc as CC
  import Control.Concurrent.MVar
  import Control.Monad
  import qualified Control.Monad.Catch as MC
  
  import GHC.Conc ( getNumProcessors, getNumCapabilities, setNumCapabilities )
  import Control.Monad.Trans.Reader
  import GHC.Driver.Pipeline.LogQueue
  import Control.Concurrent.STM
  import Control.Monad.Trans.Maybe
  
  -- Executing the pipelines
  
  mkWorkerLimit :: DynFlags -> IO WorkerLimit
  mkWorkerLimit dflags =
    case parMakeCount dflags of
      Nothing -> pure $ num_procs 1
      Just (ParMakeSemaphore h) -> pure (JSemLimit (SemaphoreName h))
      Just ParMakeNumProcessors -> num_procs <$> getNumProcessors
      Just (ParMakeThisMany n) -> pure $ num_procs n
    where
      num_procs x = NumProcessorsLimit (max 1 x)
  
  isWorkerLimitSequential :: WorkerLimit -> Bool
  isWorkerLimitSequential (NumProcessorsLimit x) = x <= 1
  isWorkerLimitSequential (JSemLimit {})         = False
  
  -- | This describes what we use to limit the number of jobs, either we limit it
  -- ourselves to a specific number or we have an external parallelism semaphore
  -- limit it for us.
  data WorkerLimit
    = NumProcessorsLimit Int
    | JSemLimit
      SemaphoreName
        -- ^ Semaphore name to use
    deriving Eq
  
  -- | Environment used when compiling a module
  data MakeEnv = MakeEnv { hsc_env :: !HscEnv -- The basic HscEnv which will be augmented for each module
                         , compile_sem :: !AbstractSem
                         -- Modify the environment for module k, with the supplied logger modification function.
                         -- For -j1, this wrapper doesn't do anything
                         -- For -jn, the wrapper initialised a log queue and then modifies the logger to pipe its output
                         --          into the log queue.
                         , withLogger :: forall a . Int -> ((Logger -> Logger) -> IO a) -> IO a
                         , env_messager :: !(Maybe Messager)
                         , diag_wrapper :: GhcMessage -> AnyGhcDiagnostic
                         }
  
  
  label_self :: String -> IO ()
  label_self thread_name = do
      self_tid <- CC.myThreadId
      CC.labelThread self_tid thread_name
  
  
  runPipelines :: WorkerLimit -> HscEnv -> (GhcMessage -> AnyGhcDiagnostic) -> Maybe Messager -> [MakeAction] -> IO ()
  -- Don't even initialise plugins if there are no pipelines
  runPipelines n_job hsc_env diag_wrapper mHscMessager all_pipelines = do
    liftIO $ label_self "main --make thread"
    case n_job of
      NumProcessorsLimit n | n <= 1 -> runSeqPipelines hsc_env diag_wrapper mHscMessager all_pipelines
      _n -> runParPipelines n_job hsc_env diag_wrapper mHscMessager all_pipelines
  
  runSeqPipelines :: HscEnv -> (GhcMessage -> AnyGhcDiagnostic) -> Maybe Messager -> [MakeAction] -> IO ()
  runSeqPipelines plugin_hsc_env diag_wrapper mHscMessager all_pipelines =
    let env = MakeEnv { hsc_env = plugin_hsc_env
                      , withLogger = \_ k -> k id
                      , compile_sem = AbstractSem (return ()) (return ())
                      , env_messager = mHscMessager
                      , diag_wrapper = diag_wrapper
                      }
    in runAllPipelines (NumProcessorsLimit 1) env all_pipelines
  
  runNjobsAbstractSem :: Int -> (AbstractSem -> IO a) -> IO a
  runNjobsAbstractSem n_jobs action = do
    compile_sem <- newQSem n_jobs
    n_capabilities <- getNumCapabilities
    n_cpus <- getNumProcessors
    let
      asem = AbstractSem (waitQSem compile_sem) (signalQSem compile_sem)
      set_num_caps n = unless (n_capabilities /= 1) $ setNumCapabilities n
      updNumCapabilities =  do
        -- Setting number of capabilities more than
        -- CPU count usually leads to high userspace
        -- lock contention. #9221
        set_num_caps $ min n_jobs n_cpus
      resetNumCapabilities = set_num_caps n_capabilities
    MC.bracket_ updNumCapabilities resetNumCapabilities $ action asem
  
  runWorkerLimit :: WorkerLimit -> (AbstractSem -> IO a) -> IO a
  #if defined(wasm32_HOST_ARCH)
  runWorkerLimit _ action = do
    lock <- newMVar ()
    action $ AbstractSem (takeMVar lock) (putMVar lock ())
  #else
  runWorkerLimit worker_limit action = case worker_limit of
      NumProcessorsLimit n_jobs ->
        runNjobsAbstractSem n_jobs action
      JSemLimit sem ->
        runJSemAbstractSem sem action
  #endif
  
  -- | Build and run a pipeline
  runParPipelines :: WorkerLimit -- ^ How to limit work parallelism
               -> HscEnv         -- ^ The basic HscEnv which is augmented with specific info for each module
               -> (GhcMessage -> AnyGhcDiagnostic)
               -> Maybe Messager   -- ^ Optional custom messager to use to report progress
               -> [MakeAction]  -- ^ The build plan for all the module nodes
               -> IO ()
  runParPipelines worker_limit plugin_hsc_env diag_wrapper mHscMessager all_pipelines = do
  
  
    -- A variable which we write to when an error has happened and we have to tell the
    -- logging thread to gracefully shut down.
    stopped_var <- newTVarIO False
    -- The queue of LogQueues which actions are able to write to. When an action starts it
    -- will add it's LogQueue into this queue.
    log_queue_queue_var <- newTVarIO newLogQueueQueue
    -- Thread which coordinates the printing of logs
    wait_log_thread <- logThread (hsc_logger plugin_hsc_env) stopped_var log_queue_queue_var
  
  
    -- Make the logger thread-safe, in case there is some output which isn't sent via the LogQueue.
    thread_safe_logger <- liftIO $ makeThreadSafe (hsc_logger plugin_hsc_env)
    let thread_safe_hsc_env = plugin_hsc_env { hsc_logger = thread_safe_logger }
  
    runWorkerLimit worker_limit $ \abstract_sem -> do
      let env = MakeEnv { hsc_env = thread_safe_hsc_env
                        , withLogger = withParLog log_queue_queue_var
                        , compile_sem = abstract_sem
                        , env_messager = mHscMessager
                        , diag_wrapper = diag_wrapper
                        }
      -- Reset the number of capabilities once the upsweep ends.
      runAllPipelines worker_limit env all_pipelines
      atomically $ writeTVar stopped_var True
      wait_log_thread
  
  withLoggerHsc :: Int -> MakeEnv -> (HscEnv -> IO a) -> IO a
  withLoggerHsc k MakeEnv{withLogger, hsc_env} cont = do
    withLogger k $ \modifyLogger -> do
      let lcl_logger = modifyLogger (hsc_logger hsc_env)
          hsc_env' = hsc_env { hsc_logger = lcl_logger }
      -- Run continuation with modified logger
      cont hsc_env'
  
  withParLog :: TVar LogQueueQueue -> Int -> ((Logger -> Logger) -> IO b) -> IO b
  withParLog lqq_var k cont = do
    let init_log = do
          -- Make a new log queue
          lq <- newLogQueue k
          -- Add it into the LogQueueQueue
          atomically $ initLogQueue lqq_var lq
          return lq
        finish_log lq = liftIO (finishLogQueue lq)
    MC.bracket init_log finish_log $ \lq -> cont (pushLogHook (const (parLogAction lq)))
  
  withLocalTmpFS :: TmpFs -> (TmpFs -> IO a) -> IO a
  withLocalTmpFS tmpfs act = do
    let initialiser = do
          liftIO $ forkTmpFsFrom tmpfs
        finaliser tmpfs_local = do
          liftIO $ mergeTmpFsInto tmpfs_local tmpfs
         -- Add remaining files which weren't cleaned up into local tmp fs for
         -- clean-up later.
         -- Clear the logQueue if this node had it's own log queue
    MC.bracket initialiser finaliser act
  
  withLocalTmpFSMake :: MakeEnv -> (MakeEnv -> IO a) -> IO a
  withLocalTmpFSMake env k =
    withLocalTmpFS (hsc_tmpfs (hsc_env env)) $ \lcl_tmpfs
      -> k (env { hsc_env = (hsc_env env) { hsc_tmpfs = lcl_tmpfs }})
  
  
  -- | Run the given actions and then wait for them all to finish.
  runAllPipelines :: WorkerLimit -> MakeEnv -> [MakeAction] -> IO ()
  runAllPipelines worker_limit env acts = do
    let single_worker = isWorkerLimitSequential worker_limit
        spawn_actions :: IO [ThreadId]
        spawn_actions = if single_worker
          then (:[]) <$> (forkIOWithUnmask $ \unmask -> void $ runLoop (\io -> io unmask) env acts)
          else runLoop forkIOWithUnmask env acts
  
        kill_actions :: [ThreadId] -> IO ()
        kill_actions tids = mapM_ killThread tids
  
    MC.bracket spawn_actions kill_actions $ \_ -> do
      mapM_ waitMakeAction acts
  
  -- | Execute each action in order, limiting the amount of parallelism by the given
  -- semaphore.
  runLoop :: (((forall a. IO a -> IO a) -> IO ()) -> IO a) -> MakeEnv -> [MakeAction] -> IO [a]
  runLoop _ _env [] = return []
  runLoop fork_thread env (MakeAction act res_var :acts) = do
  
    -- withLocalTmpFs has to occur outside of fork to remain deterministic
    new_thread <- withLocalTmpFSMake env $ \lcl_env ->
      fork_thread $ \unmask -> (do
              mres <- (unmask $ run_pipeline lcl_env act)
                        `MC.onException` (putMVar res_var Nothing) -- Defensive: If there's an unhandled exception then still signal the failure.
              putMVar res_var mres)
    threads <- runLoop fork_thread env acts
    return (new_thread : threads)
    where
        run_pipeline :: MakeEnv -> RunMakeM a -> IO (Maybe a)
        run_pipeline env p = runMaybeT (runReaderT p env)
  
  type RunMakeM a = ReaderT MakeEnv (MaybeT IO) a
  
  data MakeAction = forall a . MakeAction !(RunMakeM a) !(MVar (Maybe a))
  
  waitMakeAction :: MakeAction -> IO ()
  waitMakeAction (MakeAction _ mvar) = () <$ readMVar mvar