CGroup.hs

-- | This module provides a container-/cgroup-aware substitute for GHC's RTS
-- @-N@ flag. See 'initRTSThreads'.
module Control.Concurrent.CGroup (
  initRTSThreads,
) where

import Control.Exception (Exception (..), SomeAsyncException (SomeAsyncException), SomeException, catch, throwIO)
import qualified Data.Ratio as Ratio
import GHC.Conc (getNumProcessors, setNumCapabilities)
import System.CGroup.Types (CPUQuota (..))
import qualified System.CGroup.V1.CPU as V1
import qualified System.CGroup.V2.CPU as V2

-- | A container-/cgroup-aware substitute for GHC's RTS @-N@ flag.
--
-- On most platforms, this sets the number of runtime threads to match the
-- number of physical processors (see 'GHC.Conc.getNumProcessors'), which is the
-- default behavior of the GHC @-N@ flag.
--
-- When running within a cgroup on linux (most often within a container), this
-- observes the current process' cgroup cpu quota to constrain the number of
-- runtime threads.
--
-- See 'CPUQuota'
initRTSThreads :: IO ()
initRTSThreads = do
  quota <-
    V1.getProcessCPUQuota
      `fallback` V2.getProcessEffectiveCPUQuota
      `fallback` pure NoQuota
  initRTSThreadsFromQuota quota

-- | Use a CPU quota to set the number of runtime threads.
initRTSThreadsFromQuota :: CPUQuota -> IO ()
initRTSThreadsFromQuota NoQuota = defaultInitRTSThreads
initRTSThreadsFromQuota (CPUQuota ratio) = do
  procs <- getNumProcessors
  let capabilities = clamp 1 procs (Ratio.numerator ratio `div` Ratio.denominator ratio)
  setNumCapabilities capabilities

-- | Set number of runtime threads to the number of available processors. This
-- matches the behavior of GHC's RTS @-N@ flag.
defaultInitRTSThreads :: IO ()
defaultInitRTSThreads = setNumCapabilities =<< getNumProcessors

-- | Clamp a value within a range
clamp :: Int -> Int -> Int -> Int
clamp lower upper = max lower . min upper

-- | Catch non-async exceptions
safeCatch :: Exception e => IO a -> (e -> IO a) -> IO a
safeCatch act hdl = act `catch` (\e -> if isSyncException e then hdl e else throwIO e)

isSyncException :: Exception e => e -> Bool
isSyncException e =
  case fromException (toException e) of
    Just (SomeAsyncException _) -> False
    Nothing -> True

-- | Return the result of the first successful action
fallback :: IO a -> IO a -> IO a
fallback a b = a `safeCatch` (\(_ :: SomeException) -> b)