content: {-# LANGUAGE BlockArguments #-}
  {-# LANGUAGE NamedFieldPuns #-}
  {-# LANGUAGE RecordWildCards #-}
  {-# LANGUAGE TupleSections #-}
  {-# LANGUAGE NumericUnderscores #-}
  
  -- | Implementation of a jobserver using system semaphores.
  --
  --
  module GHC.Driver.MakeSem
    ( -- * JSem: parallelism semaphore backed
      -- by a system semaphore (Posix/Windows)
      runJSemAbstractSem
  
    -- * System semaphores
    , Semaphore, SemaphoreName(..)
  
    -- * Abstract semaphores
    , AbstractSem(..)
    , withAbstractSem
    )
    where
  
  import GHC.Prelude
  import GHC.Conc
  import GHC.Data.OrdList
  import GHC.IO.Exception
  import GHC.Utils.Outputable
  import GHC.Utils.Panic
  import GHC.Utils.Json
  
  import System.Semaphore
  
  import Control.Monad
  import qualified Control.Monad.Catch as MC
  import Control.Concurrent.MVar
  import Control.Concurrent.STM
  import Data.Foldable
  import Data.Functor
  import GHC.Stack
  import Debug.Trace
  
  ---------------------------------------
  -- Semaphore jobserver
  
  -- | A jobserver based off a system 'Semaphore'.
  --
  -- Keeps track of the pending jobs and resources
  -- available from the semaphore.
  data Jobserver
    = Jobserver
    { jSemaphore :: !Semaphore
      -- ^ The semaphore which controls available resources
    , jobs :: !(TVar JobResources)
      -- ^ The currently pending jobs, and the resources
      -- obtained from the semaphore
    }
  
  data JobserverOptions
    = JobserverOptions
    { releaseDebounce    :: !Int
       -- ^ Minimum delay, in milliseconds, between acquiring a token
       -- and releasing a token.
    , setNumCapsDebounce :: !Int
      -- ^ Minimum delay, in milliseconds, between two consecutive
      -- calls of 'setNumCapabilities'.
    }
  
  defaultJobserverOptions :: JobserverOptions
  defaultJobserverOptions =
    JobserverOptions
      { releaseDebounce    = 1000 -- 1 second
      , setNumCapsDebounce = 1000 -- 1 second
      }
  
  -- | Resources available for running jobs, i.e.
  -- tokens obtained from the parallelism semaphore.
  data JobResources
    = Jobs
    { tokensOwned :: !Int
      -- ^ How many tokens have been claimed from the semaphore
    , tokensFree  :: !Int
      -- ^ How many tokens are not currently being used
    , jobsWaiting :: !(OrdList (TMVar ()))
      -- ^ Pending jobs waiting on a token, the job will be blocked on the TMVar so putting into
      -- the TMVar will allow the job to continue.
    }
  
  instance Outputable JobResources where
    ppr Jobs{..}
      = text "JobResources" <+>
          ( braces $ hsep
            [ text "owned=" <> ppr tokensOwned
            , text "free=" <> ppr tokensFree
            , text "num_waiting=" <> ppr (length jobsWaiting)
            ] )
  
  -- | Add one new token.
  addToken :: JobResources -> JobResources
  addToken jobs@( Jobs { tokensOwned = owned, tokensFree = free })
    = jobs { tokensOwned = owned + 1, tokensFree = free + 1 }
  
  -- | Free one token.
  addFreeToken :: JobResources -> JobResources
  addFreeToken jobs@( Jobs { tokensFree = free })
    = assertPpr (tokensOwned jobs > free)
        (text "addFreeToken:" <+> ppr (tokensOwned jobs) <+> ppr free)
    $ jobs { tokensFree = free + 1 }
  
  -- | Use up one token.
  removeFreeToken :: JobResources -> JobResources
  removeFreeToken jobs@( Jobs { tokensFree = free })
    = assertPpr (free > 0)
        (text "removeFreeToken:" <+> ppr free)
    $ jobs { tokensFree = free - 1 }
  
  -- | Return one owned token.
  removeOwnedToken :: JobResources -> JobResources
  removeOwnedToken jobs@( Jobs { tokensOwned = owned })
    = assertPpr (owned > 1)
        (text "removeOwnedToken:" <+> ppr owned)
    $ jobs { tokensOwned = owned - 1 }
  
  -- | Add one new job to the end of the list of pending jobs.
  addJob :: TMVar () -> JobResources -> JobResources
  addJob job jobs@( Jobs { jobsWaiting = wait })
    = jobs { jobsWaiting = wait `SnocOL` job }
  
  -- | The state of the semaphore job server.
  data JobserverState
    = JobserverState
      { jobserverAction  :: !JobserverAction
        -- ^ The current action being performed by the
        -- job server.
      , canChangeNumCaps :: !(TVar Bool)
        -- ^ A TVar that signals whether it has been long
        -- enough since we last changed 'numCapabilities'.
      , canReleaseToken  :: !(TVar Bool)
        -- ^ A TVar that signals whether we last acquired
        -- a token long enough ago that we can now release
        -- a token.
      }
  data JobserverAction
    -- | The jobserver is idle: no thread is currently
    -- interacting with the semaphore.
    = Idle
    -- | A thread is waiting for a token on the semaphore.
    | Acquiring
      { activeWaitId   :: WaitId
      , threadFinished :: TMVar (Maybe MC.SomeException) }
  
  -- | Retrieve the 'TMVar' that signals if the current thread has finished,
  -- if any thread is currently active in the jobserver.
  activeThread_maybe :: JobserverAction -> Maybe (TMVar (Maybe MC.SomeException))
  activeThread_maybe Idle                                   = Nothing
  activeThread_maybe (Acquiring { threadFinished = tmvar }) = Just tmvar
  
  -- | Whether we should try to acquire a new token from the semaphore:
  -- there is a pending job and no free tokens.
  guardAcquire :: JobResources -> Bool
  guardAcquire ( Jobs { tokensFree, jobsWaiting } )
    = tokensFree == 0 && not (null jobsWaiting)
  
  -- | Whether we should release a token from the semaphore:
  -- there are no pending jobs and we can release a token.
  guardRelease :: JobResources -> Bool
  guardRelease ( Jobs { tokensFree, tokensOwned, jobsWaiting } )
    = null jobsWaiting && tokensFree > 0 && tokensOwned > 1
  
  ---------------------------------------
  -- Semaphore jobserver implementation
  
  -- | Add one pending job to the jobserver.
  --
  -- Blocks, waiting on the jobserver to supply a free token.
  acquireJob :: TVar JobResources -> IO ()
  acquireJob jobs_tvar = do
    (job_tmvar, _jobs0) <- tracedAtomically "acquire" $
      modifyJobResources jobs_tvar \ jobs -> do
        job_tmvar <- newEmptyTMVar
        return ((job_tmvar, jobs), addJob job_tmvar jobs)
    atomically $ takeTMVar job_tmvar
  
  -- | Signal to the job server that one job has completed,
  -- releasing its corresponding token.
  releaseJob :: TVar JobResources -> IO ()
  releaseJob jobs_tvar = do
    tracedAtomically "release" do
      modifyJobResources jobs_tvar \ jobs -> do
        massertPpr (tokensFree jobs < tokensOwned jobs)
          (text "releaseJob: more free jobs than owned jobs!")
        return ((), addFreeToken jobs)
  
  
  -- | Release all tokens owned from the semaphore (to clean up
  -- the jobserver at the end).
  cleanupJobserver :: Jobserver -> IO ()
  cleanupJobserver (Jobserver { jSemaphore = sem
                              , jobs       = jobs_tvar })
    = do
      Jobs { tokensOwned = owned } <- readTVarIO jobs_tvar
      let toks_to_release = owned - 1
        -- Subtract off the implicit token: whoever spawned the ghc process
        -- in the first place is responsible for that token.
      releaseSemaphore sem toks_to_release
  
  -- | Dispatch the available tokens acquired from the semaphore
  -- to the pending jobs in the job server.
  dispatchTokens :: JobResources -> STM JobResources
  dispatchTokens jobs@( Jobs { tokensFree = toks_free, jobsWaiting = wait } )
    | toks_free > 0
    , next `ConsOL` rest <- wait
    -- There's a pending job and a free token:
    -- pass on the token to that job, and recur.
    = do
        putTMVar next ()
        let jobs' = jobs { tokensFree = toks_free - 1, jobsWaiting = rest }
        dispatchTokens jobs'
    | otherwise
    = return jobs
  
  -- | Update the available resources used from a semaphore, dispatching
  -- any newly acquired resources.
  --
  -- Invariant: if the number of available resources decreases, there
  -- must be no pending jobs.
  --
  -- All modifications should go through this function to ensure the contents
  -- of the 'TVar' remains in normal form.
  modifyJobResources :: HasCallStack => TVar JobResources
                     -> (JobResources -> STM (a, JobResources))
                     -> STM (a, Maybe JobResources)
  modifyJobResources jobs_tvar action = do
    old_jobs  <- readTVar jobs_tvar
    (a, jobs) <- action old_jobs
  
    -- Check the invariant: if the number of free tokens has decreased,
    -- there must be no pending jobs.
    massertPpr (null (jobsWaiting jobs) || tokensFree jobs >= tokensFree old_jobs) $
      vcat [ text "modiyJobResources: pending jobs but fewer free tokens" ]
    dispatched_jobs <- dispatchTokens jobs
    writeTVar jobs_tvar dispatched_jobs
    return (a, Just dispatched_jobs)
  
  
  tracedAtomically_ :: String -> STM (Maybe JobResources) -> IO ()
  tracedAtomically_ s act = tracedAtomically s (((),) <$> act)
  
  tracedAtomically :: String -> STM (a, Maybe JobResources) -> IO a
  tracedAtomically origin act = do
    (a, mjr) <- atomically act
    forM_ mjr $ \ jr -> do
      -- Use the "jsem:" prefix to identify where the write traces are
      traceEventIO ("jsem:" ++ renderJobResources origin jr)
    return a
  
  renderJobResources :: String -> JobResources -> String
  renderJobResources origin (Jobs own free pending) = showSDocUnsafe $ renderJSON $
    JSObject [ ("name", JSString origin)
             , ("owned", JSInt own)
             , ("free", JSInt free)
             , ("pending", JSInt (length pending) )
             ]
  
  
  -- | Spawn a new thread that waits on the semaphore in order to acquire
  -- an additional token.
  acquireThread :: Jobserver -> IO JobserverAction
  acquireThread (Jobserver { jSemaphore = sem, jobs = jobs_tvar }) = do
      threadFinished_tmvar <- newEmptyTMVarIO
      let
        wait_result_action :: Either MC.SomeException Bool -> IO ()
        wait_result_action wait_res =
          tracedAtomically_ "acquire_thread" do
            (r, jb) <- case wait_res of
              Left (e :: MC.SomeException) -> do
                return $ (Just e, Nothing)
              Right success -> do
                if success
                  then do
                    modifyJobResources jobs_tvar \ jobs ->
                      return (Nothing, addToken jobs)
                  else
                    return (Nothing, Nothing)
            putTMVar threadFinished_tmvar r
            return jb
      wait_id <- forkWaitOnSemaphoreInterruptible sem wait_result_action
      labelThread (waitingThreadId wait_id) "acquire_thread"
      return $ Acquiring { activeWaitId   = wait_id
                         , threadFinished = threadFinished_tmvar }
  
  -- | Spawn a thread to release ownership of one resource from the semaphore,
  -- provided we have spare resources and no pending jobs.
  releaseThread :: Jobserver -> IO JobserverAction
  releaseThread (Jobserver { jSemaphore = sem, jobs = jobs_tvar }) = do
    threadFinished_tmvar <- newEmptyTMVarIO
    MC.mask_ do
      -- Pre-release the resource so that another thread doesn't take control of it
      -- just as we release the lock on the semaphore.
      still_ok_to_release
        <- tracedAtomically "pre_release" $
           modifyJobResources jobs_tvar \ jobs ->
             if guardRelease jobs
                 -- TODO: should this also debounce?
             then return (True , removeOwnedToken $ removeFreeToken jobs)
             else return (False, jobs)
      if not still_ok_to_release
      then return Idle
      else do
        tid <- forkIO $ do
          x <- MC.try $ releaseSemaphore sem 1
          tracedAtomically_ "post-release" $ do
            (r, jobs) <- case x of
              Left (e :: MC.SomeException) -> do
                modifyJobResources jobs_tvar \ jobs ->
                  return (Just e, addToken jobs)
              Right _ -> do
                return (Nothing, Nothing)
            putTMVar threadFinished_tmvar r
            return jobs
        labelThread tid "release_thread"
        return Idle
  
  -- | When there are pending jobs but no free tokens,
  -- spawn a thread to acquire a new token from the semaphore.
  --
  -- See 'acquireThread'.
  tryAcquire :: JobserverOptions
             -> Jobserver
             -> JobserverState
             -> STM (IO JobserverState)
  tryAcquire opts js@( Jobserver { jobs = jobs_tvar })
    st@( JobserverState { jobserverAction = Idle } )
    = do
      jobs <- readTVar jobs_tvar
      guard $ guardAcquire jobs
      return do
        action           <- acquireThread js
        -- Set a debounce after acquiring a token.
        can_release_tvar <- registerDelay $ (releaseDebounce opts * 1000)
        return $ st { jobserverAction = action
                    , canReleaseToken = can_release_tvar }
  tryAcquire _ _ _ = retry
  
  -- | When there are free tokens and no pending jobs,
  -- spawn a thread to release a token from the semaphore.
  --
  -- See 'releaseThread'.
  tryRelease :: Jobserver
             -> JobserverState
             -> STM (IO JobserverState)
  tryRelease sjs@( Jobserver { jobs = jobs_tvar } )
    st@( JobserverState
        { jobserverAction = Idle
        , canReleaseToken = can_release_tvar } )
    = do
      jobs <- readTVar jobs_tvar
      guard  $ guardRelease jobs
      can_release <- readTVar can_release_tvar
      guard can_release
      return do
        action <- releaseThread sjs
        return $ st { jobserverAction = action }
  tryRelease _ _ = retry
  
  -- | Wait for an active thread to finish. Once it finishes:
  --
  --  - set the 'JobserverAction' to 'Idle',
  --  - update the number of capabilities to reflect the number
  --    of owned tokens from the semaphore.
  tryNoticeIdle :: JobserverOptions
                -> TVar JobResources
                -> JobserverState
                -> STM (IO JobserverState)
  tryNoticeIdle opts jobs_tvar jobserver_state
    | Just threadFinished_tmvar <- activeThread_maybe $ jobserverAction jobserver_state
    = sync_num_caps (canChangeNumCaps jobserver_state) threadFinished_tmvar
    | otherwise
    = retry -- no active thread: wait until jobserver isn't idle
    where
      sync_num_caps :: TVar Bool
                    -> TMVar (Maybe MC.SomeException)
                    -> STM (IO JobserverState)
      sync_num_caps can_change_numcaps_tvar threadFinished_tmvar = do
        mb_ex <- takeTMVar threadFinished_tmvar
        for_ mb_ex MC.throwM
        Jobs { tokensOwned } <- readTVar jobs_tvar
        can_change_numcaps <- readTVar can_change_numcaps_tvar
        guard can_change_numcaps
        return do
          x <- getNumCapabilities
          can_change_numcaps_tvar_2 <-
            if x == tokensOwned
            then return can_change_numcaps_tvar
            else do
              setNumCapabilities tokensOwned
              registerDelay $ (setNumCapsDebounce opts * 1000)
          return $
            jobserver_state
              { jobserverAction  = Idle
              , canChangeNumCaps = can_change_numcaps_tvar_2 }
  
  -- | Try to stop the current thread which is acquiring/releasing resources
  -- if that operation is no longer relevant.
  tryStopThread :: TVar JobResources
                -> JobserverState
                -> STM (IO JobserverState)
  tryStopThread jobs_tvar jsj = do
    case jobserverAction jsj of
      Acquiring { activeWaitId = wait_id } -> do
       jobs <- readTVar jobs_tvar
       guard $ null (jobsWaiting jobs)
       return do
         interruptWaitOnSemaphore wait_id
         return $ jsj { jobserverAction = Idle }
      _ -> retry
  
  -- | Main jobserver loop: acquire/release resources as
  -- needed for the pending jobs and available semaphore tokens.
  jobserverLoop :: JobserverOptions -> Jobserver -> IO ()
  jobserverLoop opts sjs@(Jobserver { jobs = jobs_tvar })
    = do
        true_tvar <- newTVarIO True
        let init_state :: JobserverState
            init_state =
              JobserverState
                { jobserverAction  = Idle
                , canChangeNumCaps = true_tvar
                , canReleaseToken  = true_tvar }
        loop init_state
    where
      loop s = do
        action <- atomically $ asum $ (\x -> x s) <$>
          [ tryRelease    sjs
          , tryAcquire    opts sjs
          , tryNoticeIdle opts jobs_tvar
          , tryStopThread jobs_tvar
          ]
        s <- action
        loop s
  
  -- | Create a new jobserver using the given semaphore handle.
  makeJobserver :: SemaphoreName -> IO (AbstractSem, IO ())
  makeJobserver sem_name = do
    semaphore <- openSemaphore sem_name
    let
      init_jobs =
        Jobs { tokensOwned = 1
             , tokensFree  = 1
             , jobsWaiting = NilOL
             }
    jobs_tvar <- newTVarIO init_jobs
    let
      opts = defaultJobserverOptions -- TODO: allow this to be configured
      sjs = Jobserver { jSemaphore = semaphore
                      , jobs       = jobs_tvar }
    loop_finished_mvar <- newEmptyMVar
    loop_tid <- forkIOWithUnmask \ unmask -> do
      r <- try $ unmask $ jobserverLoop opts sjs
      putMVar loop_finished_mvar $
        case r of
          Left e
            | Just ThreadKilled <- fromException e
            -> Nothing
            | otherwise
            -> Just e
          Right () -> Nothing
    labelThread loop_tid "job_server"
    let
      acquireSem = acquireJob jobs_tvar
      releaseSem = releaseJob jobs_tvar
      cleanupSem = do
        -- this is interruptible
        cleanupJobserver sjs
        killThread loop_tid
        mb_ex <- takeMVar loop_finished_mvar
        for_ mb_ex MC.throwM
  
    return (AbstractSem{..}, cleanupSem)
  
  -- | Implement an abstract semaphore using a semaphore 'Jobserver'
  -- which queries the system semaphore of the given name for resources.
  runJSemAbstractSem :: SemaphoreName         -- ^ the system semaphore to use
                     -> (AbstractSem -> IO a) -- ^ the operation to run
                                              -- which requires a semaphore
                     -> IO a
  runJSemAbstractSem sem action = MC.mask \ unmask -> do
    (abs, cleanup) <- makeJobserver sem
    r <- try $ unmask $ action abs
    case r of
      Left (e1 :: MC.SomeException) -> do
        (_ :: Either MC.SomeException ()) <- MC.try cleanup
        MC.throwM e1
      Right x -> cleanup $> x
  
  {- Note [Architecture of the Job Server]
  ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  In `-jsem` mode, the amount of parallelism that GHC can use is controlled by a
  system semaphore. We take resources from the semaphore when we need them, and
  give them back if we don't have enough to do.
  
  A naive implementation would just take and release the semaphore around performing
  the action, but this leads to two issues:
  
  * When taking a token in the semaphore, we must call `setNumCapabilities` in order
    to adjust how many capabilities are available for parallel garbage collection.
    This causes unnecessary synchronisations.
  * We want to implement a debounce, so that whilst there is pending work in the
    current process we prefer to keep hold of resources from the semaphore.
    This reduces overall memory usage, as there are fewer live GHC processes at once.
  
  Therefore, the obtention of semaphore resources is separated away from the
  request for the resource in the driver.
  
  A token from the semaphore is requested using `acquireJob`. This creates a pending
  job, which is a MVar that can be filled in to signal that the requested token is ready.
  
  When the job is finished, the token is released by calling `releaseJob`, which just
  increases the number of `free` jobs. If there are more pending jobs when the free count
  is increased, the token is immediately reused (see `modifyJobResources`).
  
  The `jobServerLoop` interacts with the system semaphore: when there are pending
  jobs, `acquireThread` blocks, waiting for a token from the semaphore. Once a
  token is obtained, it increases the owned count.
  
  When GHC has free tokens (tokens from the semaphore that it is not using),
  no pending jobs, and the debounce has expired, then `releaseThread` will
  release tokens back to the global semaphore.
  
  `tryStopThread` attempts to kill threads which are waiting to acquire a resource
  when we no longer need it. For example, consider that we attempt to acquire two
  tokens, but the first job finishes before we acquire the second token.
  This second token is no longer needed, so we should cancel the wait
  (as it would not be used to do any work, and not be returned until the debounce).
  We only need to kill `acquireJob`, because `releaseJob` never blocks.
  
  Note [Eventlog Messages for jsem]
  ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  It can be tricky to verify that the work is shared adequately across different
  processes. To help debug this, we output the values of `JobResource` to the
  eventlog whenever the global state changes. There are some scripts which can be used
  to analyse this output and report statistics about core saturation in the
  GitHub repo (https://github.com/mpickering/ghc-jsem-analyse).
  
  -}
  