{-# LANGUAGE TemplateHaskell #-}

-- | Types and operations for CPU-related data within a cgroup.
module System.CGroup.V2.CPU (
  getProcessEffectiveCPUQuota,
  getEffectiveCPUQuota,
  getCPUQuota,
  CPUQuota (..),
) where

import Data.Ratio
import Path
import System.CGroup.Types (CPUQuota (..))
import System.CGroup.V2.CGroup
import System.Directory (doesFileExist)

-- | Get the current process' effective CPU quota
--
-- See 'getEffectiveCPUQuota'
getProcessEffectiveCPUQuota :: IO CPUQuota
getProcessEffectiveCPUQuota :: IO CPUQuota
getProcessEffectiveCPUQuota = CGroup -> IO CPUQuota
getEffectiveCPUQuota (CGroup -> IO CPUQuota) -> IO CGroup -> IO CPUQuota
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< IO CGroup
resolveCGroup

-- | Compute the "effective CPU quota" fr a cgroup, which may be smaller than
-- the given cgroup's individual quota.
--
-- When a parent (or grandparent, etc) of this cgroup has a lower cpu quota,
-- the lower quota is returned instead.
getEffectiveCPUQuota :: CGroup -> IO CPUQuota
getEffectiveCPUQuota :: CGroup -> IO CPUQuota
getEffectiveCPUQuota CGroup
cgroup = do
  [CPUQuota]
quotas <- CGroup -> IO [CPUQuota]
cpuQuotasUntilRoot CGroup
cgroup
  CPUQuota -> IO CPUQuota
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((CPUQuota -> CPUQuota -> CPUQuota)
-> CPUQuota -> [CPUQuota] -> CPUQuota
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr CPUQuota -> CPUQuota -> CPUQuota
forall a. Ord a => a -> a -> a
min CPUQuota
NoQuota [CPUQuota]
quotas)

-- | Read this specific cgroup's "cpu.max" file into a CPUQuota.
--
-- For example:
--
-- @
-- | cpu.max        | quota            |
-- | -------------- | ---------------- |
-- |  100000 100000 | CPUQuota (1 % 1) |
-- |  200000 100000 | CPUQuota (2 % 1) |
-- |   50000 100000 | CPUQuota (1 % 2) |
-- |     max 100000 | NoQuota          |
-- @
--
--
-- Returns NoQuota for the root cgroup, or when the cpu controller is not enabled in this cgroup.
--
-- __Most often, you'll want to use 'getEffectiveCPUQuota' instead.__
getCPUQuota :: CGroup -> IO CPUQuota
getCPUQuota :: CGroup -> IO CPUQuota
getCPUQuota CGroup
cgroup = do
  let path :: Path Abs File
path = CGroup -> Path Abs Dir
cgroupRoot CGroup
cgroup Path Abs Dir -> Path Rel File -> Path Abs File
forall b t. Path b Dir -> Path Rel t -> Path b t
</> CGroup -> Path Rel Dir
cgroupLeaf CGroup
cgroup Path Rel Dir -> Path Rel File -> Path Rel File
forall b t. Path b Dir -> Path Rel t -> Path b t
</> Path Rel File
cpuMaxPath
  Bool
exists <- String -> IO Bool
doesFileExist (Path Abs File -> String
forall b t. Path b t -> String
toFilePath Path Abs File
path)
  if Bool -> Bool
not Bool
exists
    then CPUQuota -> IO CPUQuota
forall (f :: * -> *) a. Applicative f => a -> f a
pure CPUQuota
NoQuota
    else do
      String
content <- String -> IO String
readFile (Path Abs File -> String
forall b t. Path b t -> String
toFilePath Path Abs File
path)
      case String -> [String]
words String
content of
        [String
"max", String
_] -> CPUQuota -> IO CPUQuota
forall (f :: * -> *) a. Applicative f => a -> f a
pure CPUQuota
NoQuota
        [String
numText, String
denText] -> do
          Int
num <- String -> IO Int
forall a. Read a => String -> IO a
readIO String
numText
          Int
den <- String -> IO Int
forall a. Read a => String -> IO a
readIO String
denText
          CPUQuota -> IO CPUQuota
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Ratio Int -> CPUQuota
CPUQuota (Int
num Int -> Int -> Ratio Int
forall a. Integral a => a -> a -> Ratio a
% Int
den))
        [String]
_ -> String -> IO CPUQuota
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Couldn't parse cpu.max"

-- | Get the parent cgroup. Returns Nothing when the provided cgroup is the root
-- cgroup.
getParentCGroup :: CGroup -> Maybe CGroup
getParentCGroup :: CGroup -> Maybe CGroup
getParentCGroup CGroup
cgroup =
  if CGroup -> Path Rel Dir
cgroupLeaf CGroup
cgroup Path Rel Dir -> Path Rel Dir -> Bool
forall a. Eq a => a -> a -> Bool
== Path Rel Dir -> Path Rel Dir
forall b t. Path b t -> Path b Dir
parent (CGroup -> Path Rel Dir
cgroupLeaf CGroup
cgroup)
    then Maybe CGroup
forall a. Maybe a
Nothing
    else CGroup -> Maybe CGroup
forall a. a -> Maybe a
Just CGroup
cgroup{cgroupLeaf :: Path Rel Dir
cgroupLeaf = Path Rel Dir -> Path Rel Dir
forall b t. Path b t -> Path b Dir
parent (CGroup -> Path Rel Dir
cgroupLeaf CGroup
cgroup)}

-- | Return all CPU quotas from all parent cgroups
cpuQuotasUntilRoot :: CGroup -> IO [CPUQuota]
cpuQuotasUntilRoot :: CGroup -> IO [CPUQuota]
cpuQuotasUntilRoot = (CGroup -> IO CPUQuota) -> [CGroup] -> IO [CPUQuota]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse CGroup -> IO CPUQuota
getCPUQuota ([CGroup] -> IO [CPUQuota])
-> (CGroup -> [CGroup]) -> CGroup -> IO [CPUQuota]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CGroup -> [CGroup]
cgroupsUntilRoot

-- | Return the list of cgroups up to but excluding the root cgroup
cgroupsUntilRoot :: CGroup -> [CGroup]
cgroupsUntilRoot :: CGroup -> [CGroup]
cgroupsUntilRoot = (CGroup -> Maybe CGroup) -> CGroup -> [CGroup]
forall a. (a -> Maybe a) -> a -> [a]
iterateMaybe CGroup -> Maybe CGroup
getParentCGroup

-- | like 'iterate', but terminated when @Nothing@ is returned by the provided
-- function
iterateMaybe :: (a -> Maybe a) -> a -> [a]
iterateMaybe :: forall a. (a -> Maybe a) -> a -> [a]
iterateMaybe a -> Maybe a
f a
x = a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a] -> (a -> [a]) -> Maybe a -> [a]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe [] ((a -> Maybe a) -> a -> [a]
forall a. (a -> Maybe a) -> a -> [a]
iterateMaybe a -> Maybe a
f) (a -> Maybe a
f a
x)

cpuMaxPath :: Path Rel File
cpuMaxPath :: Path Rel File
cpuMaxPath = $(mkRelFile "cpu.max")