{-# LANGUAGE TemplateHaskell #-}

-- | Common types and operations for cgroups (v2)
module System.CGroup.V2.CGroup (
  CGroup (..),
  resolveCGroup,
  resolveCGroup',
) where

import Data.Foldable (find)
import qualified Data.Text as Text
import Path
import System.CGroup.Types (Mount (..), RawCGroup (..), parseCGroups, parseFile, parseMountInfo)

-- | A cgroup (under cgroups v2)
data CGroup = CGroup
  { -- | The root of the cgroup hierarchy
    CGroup -> Path Abs Dir
cgroupRoot :: Path Abs Dir
  , -- | A specific cgroup's relative path from the cgroup hierarchy root
    CGroup -> Path Rel Dir
cgroupLeaf :: Path Rel Dir
  }
  deriving (CGroup -> CGroup -> Bool
(CGroup -> CGroup -> Bool)
-> (CGroup -> CGroup -> Bool) -> Eq CGroup
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: CGroup -> CGroup -> Bool
$c/= :: CGroup -> CGroup -> Bool
== :: CGroup -> CGroup -> Bool
$c== :: CGroup -> CGroup -> Bool
Eq, Eq CGroup
Eq CGroup
-> (CGroup -> CGroup -> Ordering)
-> (CGroup -> CGroup -> Bool)
-> (CGroup -> CGroup -> Bool)
-> (CGroup -> CGroup -> Bool)
-> (CGroup -> CGroup -> Bool)
-> (CGroup -> CGroup -> CGroup)
-> (CGroup -> CGroup -> CGroup)
-> Ord CGroup
CGroup -> CGroup -> Bool
CGroup -> CGroup -> Ordering
CGroup -> CGroup -> CGroup
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: CGroup -> CGroup -> CGroup
$cmin :: CGroup -> CGroup -> CGroup
max :: CGroup -> CGroup -> CGroup
$cmax :: CGroup -> CGroup -> CGroup
>= :: CGroup -> CGroup -> Bool
$c>= :: CGroup -> CGroup -> Bool
> :: CGroup -> CGroup -> Bool
$c> :: CGroup -> CGroup -> Bool
<= :: CGroup -> CGroup -> Bool
$c<= :: CGroup -> CGroup -> Bool
< :: CGroup -> CGroup -> Bool
$c< :: CGroup -> CGroup -> Bool
compare :: CGroup -> CGroup -> Ordering
$ccompare :: CGroup -> CGroup -> Ordering
Ord, Int -> CGroup -> ShowS
[CGroup] -> ShowS
CGroup -> FilePath
(Int -> CGroup -> ShowS)
-> (CGroup -> FilePath) -> ([CGroup] -> ShowS) -> Show CGroup
forall a.
(Int -> a -> ShowS) -> (a -> FilePath) -> ([a] -> ShowS) -> Show a
showList :: [CGroup] -> ShowS
$cshowList :: [CGroup] -> ShowS
show :: CGroup -> FilePath
$cshow :: CGroup -> FilePath
showsPrec :: Int -> CGroup -> ShowS
$cshowsPrec :: Int -> CGroup -> ShowS
Show)

-- | Resolve the cgroup (v2) used for the current process
--
-- see cgroups(7): \/proc\/self\/cgroup is a file that contains information
-- about control groups applied to this process
--
-- see proc(5): \/proc\/self\/mountinfo is a file that contains information
-- about mounts available to this process
--
-- Throws an Exception when the cgroup is unable to be found, or when the
-- current process is not running under cgroups v2
resolveCGroup :: IO CGroup
resolveCGroup :: IO CGroup
resolveCGroup = do
  Path Abs File
cgroupPath <- FilePath -> IO (Path Abs File)
forall (m :: * -> *). MonadThrow m => FilePath -> m (Path Abs File)
parseAbsFile FilePath
"/proc/self/cgroup"
  Path Abs File
mountinfoPath <- FilePath -> IO (Path Abs File)
forall (m :: * -> *). MonadThrow m => FilePath -> m (Path Abs File)
parseAbsFile FilePath
"/proc/self/mountinfo"
  Path Abs File -> Path Abs File -> IO CGroup
resolveCGroup' Path Abs File
cgroupPath Path Abs File
mountinfoPath

-- | Resolve a cgroup (v2) from the given cgroup and mountinfo files
--
-- Throws an Exception when the cgroup is unable to be found, or when the
-- provided paths do not construct a valid cgroup
resolveCGroup' :: Path Abs File -> Path Abs File -> IO CGroup
resolveCGroup' :: Path Abs File -> Path Abs File -> IO CGroup
resolveCGroup' Path Abs File
cgroupPath Path Abs File
mountinfoPath = do
  [RawCGroup]
cgroups <- Parser [RawCGroup] -> Path Abs File -> IO [RawCGroup]
forall a b. Parser a -> Path b File -> IO a
parseFile Parser [RawCGroup]
parseCGroups Path Abs File
cgroupPath
  case [RawCGroup]
cgroups of
    -- expect to find a cgroup with hierarchy ID 0 and an empty list of controllers
    [RawCGroup Text
"0" [] Path Abs Dir
cgroupLeafAbs] -> do
      [Mount]
mounts <- Parser [Mount] -> Path Abs File -> IO [Mount]
forall a b. Parser a -> Path b File -> IO a
parseFile Parser [Mount]
parseMountInfo Path Abs File
mountinfoPath
      Mount
cgroupRootMount <- IO Mount -> (Mount -> IO Mount) -> Maybe Mount -> IO Mount
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (FilePath -> IO Mount
forall (m :: * -> *) a. MonadFail m => FilePath -> m a
fail FilePath
"Couldn't find cgroup hierarchy root mount") Mount -> IO Mount
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Mount] -> Maybe Mount
findCGroupHierarchyRootMount [Mount]
mounts)
      Path Abs Dir
mountPointAsPath <- FilePath -> IO (Path Abs Dir)
forall (m :: * -> *). MonadThrow m => FilePath -> m (Path Abs Dir)
parseAbsDir (Text -> FilePath
Text.unpack (Mount -> Text
mountPoint Mount
cgroupRootMount))

      case Path Abs Dir -> FilePath
fromAbsDir Path Abs Dir
cgroupLeafAbs of
        FilePath
"/" ->
          CGroup -> IO CGroup
forall (f :: * -> *) a. Applicative f => a -> f a
pure
            ( CGroup
                { cgroupRoot :: Path Abs Dir
cgroupRoot = Path Abs Dir
mountPointAsPath
                , cgroupLeaf :: Path Rel Dir
cgroupLeaf = $(mkRelDir ".")
                }
            )
        FilePath
_ -> do
          -- Drop the leading '/' from the cgroup path
          Path Rel Dir
cgroupLeafRel <- FilePath -> IO (Path Rel Dir)
forall (m :: * -> *). MonadThrow m => FilePath -> m (Path Rel Dir)
parseRelDir (Int -> ShowS
forall a. Int -> [a] -> [a]
drop Int
1 (Path Abs Dir -> FilePath
fromAbsDir Path Abs Dir
cgroupLeafAbs))
          CGroup -> IO CGroup
forall (f :: * -> *) a. Applicative f => a -> f a
pure
            ( CGroup
                { cgroupRoot :: Path Abs Dir
cgroupRoot = Path Abs Dir
mountPointAsPath
                , cgroupLeaf :: Path Rel Dir
cgroupLeaf = Path Rel Dir
cgroupLeafRel
                }
            )
    [RawCGroup]
_ -> FilePath -> IO CGroup
forall (m :: * -> *) a. MonadFail m => FilePath -> m a
fail (FilePath
"Found incompatible cgroups: " FilePath -> ShowS
forall a. Semigroup a => a -> a -> a
<> [RawCGroup] -> FilePath
forall a. Show a => a -> FilePath
show [RawCGroup]
cgroups)

-- | Find the cgroups v2 hierarchy root.
--
-- We expect to find a mount with the filesystem type "cgroup2"
findCGroupHierarchyRootMount :: [Mount] -> Maybe Mount
findCGroupHierarchyRootMount :: [Mount] -> Maybe Mount
findCGroupHierarchyRootMount = (Mount -> Bool) -> [Mount] -> Maybe Mount
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
"cgroup2") (Text -> Bool) -> (Mount -> Text) -> Mount -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Mount -> Text
mountFilesystemType)