Add support for cgroups v2

CHANGELOG.md

@@ -1,5 +1,9 @@
 # Revision history for cgroup-rts-threads
 
+## 0.2.0.0
+
+- Adds support for cgroups v2
+
 ## 0.1.0.0
 
 - Initial release

cgroup-rts-threads.cabal

@@ -1,6 +1,6 @@
 cabal-version:      3.0
 name:               cgroup-rts-threads
-version:            0.1.0.0
+version:            0.2.0.0
 synopsis:
   A container-/cgroup-aware substitute for the GHC RTS `-N` flag
 
@@ -40,15 +40,17 @@ library
   -- cabal-fmt: expand src
   exposed-modules:
     Control.Concurrent.CGroup
-    System.CGroup
-    System.CGroup.CPU
-    System.CGroup.Controller
-    System.CGroup.Controller.Internal
+    System.CGroup.Types
+    System.CGroup.V1.CPU
+    System.CGroup.V1.Controller
+    System.CGroup.V2.CGroup
+    System.CGroup.V2.CPU
 
   build-depends:
-    , base        >=4.13 && <5
-    , megaparsec  >=8    && <10
-    , path        >=0.7  && <0.10
+    , base        >=4.13    && <5
+    , directory   ^>=1.3.4.0
+    , megaparsec  >=8       && <10
+    , path        >=0.7     && <0.10
     , text
 
   hs-source-dirs:  src
@@ -59,8 +61,10 @@ test-suite test
   hs-source-dirs: test
   main-is:        Main.hs
   other-modules:
-    System.CGroup.CPUSpec
-    System.CGroup.ControllerSpec
+    System.CGroup.V1.CPUSpec
+    System.CGroup.V1.ControllerSpec
+    System.CGroup.V2.CGroupSpec
+    System.CGroup.V2.CPUSpec
 
   build-depends:
     , base

src/Control/Concurrent/CGroup.hs

@@ -5,8 +5,11 @@ module Control.Concurrent.CGroup (
 ) where
 
 import Control.Exception (Exception (..), SomeAsyncException (SomeAsyncException), SomeException, catch, throwIO)
+import qualified Data.Ratio as Ratio
 import GHC.Conc (getNumProcessors, setNumCapabilities)
-import System.CGroup.CPU (CPUQuota (..), getCPUQuota, resolveCPUController)
+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.
 --
@@ -20,24 +23,20 @@ import System.CGroup.CPU (CPUQuota (..), getCPUQuota, resolveCPUController)
 --
 -- See 'CPUQuota'
 initRTSThreads :: IO ()
-initRTSThreads =
-  initRTSThreadsFromCGroup
-    `safeCatch` (\(_ :: SomeException) -> defaultInitRTSThreads)
+initRTSThreads = do
+  quota <-
+    V1.getProcessCPUQuota
+      `fallback` V2.getProcessEffectiveCPUQuota
+      `fallback` pure NoQuota
+  initRTSThreadsFromQuota quota
 
--- | Uses the current process' cgroup cpu quota to set the number of runtime
--- threads.
---
--- Throws an Exception when the current process is not running within a cgroup.
-initRTSThreadsFromCGroup :: IO ()
-initRTSThreadsFromCGroup = do
-  cpuController <- resolveCPUController
-  cgroupCpuQuota <- getCPUQuota cpuController
-  case cgroupCpuQuota of
-    NoQuota -> defaultInitRTSThreads
-    CPUQuota quota period -> do
-      procs <- getNumProcessors
-      let capabilities = clamp 1 procs (quota `div` period)
-      setNumCapabilities capabilities
+-- | 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.
@@ -57,3 +56,7 @@ 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)