Add processor_count and processor_quota APIs

etc.gemspec

@@ -31,6 +31,7 @@ Gem::Specification.new do |spec|
     ext/etc/constdefs.h
     ext/etc/etc.c
     ext/etc/extconf.rb
+    ext/etc/lib/etc.rb
     ext/etc/mkconstants.rb
     test/etc/test_etc.rb
   ] + changelogs

ext/etc/lib/etc.rb

@@ -0,0 +1,173 @@
+# frozen_string_literal: true
+
+require "etc.so"
+
+module Etc
+  # Internal support for determining the CPU capacity available through Linux
+  # control groups.
+  module Cgroup
+    CGROUP_PATH = "/proc/self/cgroup"
+    MOUNTINFO_PATH = "/proc/self/mountinfo"
+
+    module_function
+
+    def cpu_quota(cgroup_path = CGROUP_PATH, mountinfo_path = MOUNTINFO_PATH)
+      memberships = memberships(cgroup_path)
+      return if memberships.empty?
+
+      limits, matched = limits_from_mounts(memberships, mountinfo_path)
+      limits = limits_from_conventional_paths(memberships) unless matched
+      limits.min
+    rescue Errno::EACCES, Errno::EINVAL, Errno::ENOENT, Errno::ENOTDIR, ArgumentError
+      nil
+    end
+
+    def memberships(path)
+      File.readlines(path, chomp: true).filter_map do |line|
+        _id, controllers, cgroup = line.split(":", 3)
+        next unless controllers && cgroup
+
+        if controllers.empty?
+          [:v2, cgroup]
+        elsif controllers.split(",").include?("cpu")
+          [:v1, cgroup]
+        end
+      end
+    end
+
+    def limits_from_mounts(memberships, mountinfo_path)
+      mounts = File.readlines(mountinfo_path, chomp: true).filter_map do |line|
+        fields = line.split
+        separator = fields.index("-")
+        next unless separator && separator >= 6
+
+        type = case fields[separator + 1]
+               when "cgroup2"
+                 :v2
+               when "cgroup"
+                 options = fields[5...separator] + fields[(separator + 2)..]
+                 :v1 if options.any? { |option| option.split(",").include?("cpu") }
+               end
+        next unless type
+
+        [type, unescape_path(fields[3]), unescape_path(fields[4])]
+      end
+
+      limits = []
+      matched = false
+      memberships.each do |type, path|
+        mounts.each do |mount_type, root, mountpoint|
+          next unless type == mount_type
+
+          directory = resolve_directory(root, mountpoint, path)
+          next unless directory
+
+          matched = true
+          limits.concat(limits_for_ancestors(type, directory, mountpoint))
+        end
+      end
+      [limits, matched]
+    end
+
+    def limits_from_conventional_paths(memberships)
+      limits = []
+      limits << read_limit(:v2, "/sys/fs/cgroup") if memberships.any? { |type,| type == :v2 }
+
+      if memberships.any? { |type,| type == :v1 }
+        limits << read_limit(:v1, "/sys/fs/cgroup/cpu")
+        limits << read_limit(:v1, "/sys/fs/cgroup/cpu,cpuacct")
+      end
+
+      limits.compact
+    end
+
+    def resolve_directory(root, mountpoint, path)
+      root = File.expand_path(root)
+      path = File.expand_path(path)
+
+      relative = if path == root || path == "/"
+                   ""
+                 elsif root == "/"
+                   path.delete_prefix("/")
+                 elsif path.start_with?("#{root}/")
+                   path.delete_prefix("#{root}/")
+                 else
+                   path.delete_prefix("/")
+                 end
+
+      File.expand_path(relative, mountpoint)
+    end
+
+    def limits_for_ancestors(type, directory, mountpoint)
+      mountpoint = File.expand_path(mountpoint)
+      directory = File.expand_path(directory)
+      return [] unless directory == mountpoint || directory.start_with?("#{mountpoint}/")
+
+      limits = []
+      loop do
+        limit = read_limit(type, directory)
+        limits << limit if limit
+        break if directory == mountpoint
+
+        directory = File.dirname(directory)
+      end
+      limits
+    end
+
+    def read_limit(type, directory)
+      case type
+      when :v2
+        maximum, period = File.read(File.join(directory, "cpu.max")).split
+        return if maximum == "max" || !maximum || !period
+
+        ratio(maximum, period)
+      when :v1
+        maximum = File.read(File.join(directory, "cpu.cfs_quota_us"))
+        return if Integer(maximum) < 0
+
+        period = File.read(File.join(directory, "cpu.cfs_period_us"))
+        ratio(maximum, period)
+      end
+    rescue Errno::EACCES, Errno::EINVAL, Errno::ENOENT, Errno::ENOTDIR, ArgumentError
+      nil
+    end
+
+    def ratio(maximum, period)
+      maximum = Float(maximum)
+      period = Float(period)
+      return unless maximum.positive? && period.positive?
+
+      maximum / period
+    end
+
+    def unescape_path(path)
+      path.gsub(/\\([0-7]{3})/) { $1.to_i(8).chr }
+    end
+  end
+  private_constant :Cgroup
+
+  # call-seq:
+  #   processor_count -> Integer
+  #
+  # Returns the number of processors available to the process as an Integer.
+  #
+  # This is an alias for Etc.nprocessors.
+  alias processor_count nprocessors
+  module_function :processor_count
+
+  # call-seq:
+  #   processor_quota -> Float
+  #
+  # Returns the processor capacity available to the process as a Float.
+  #
+  # On Linux, this accounts for CPU bandwidth limits in cgroup v1 and v2. The
+  # result will not exceed Etc.processor_count, which already accounts for CPU
+  # affinity. On other platforms, or when no cgroup limit can be determined,
+  # it returns Etc.processor_count converted to a Float.
+  def processor_quota
+    count = processor_count.to_f
+    quota = Cgroup.cpu_quota if RUBY_PLATFORM.include?("linux")
+    quota ? [count, quota].min : count
+  end
+  module_function :processor_quota
+end

test/etc/test_etc.rb

@@ -1,6 +1,8 @@
 # frozen_string_literal: true
 require "test/unit"
 require "etc"
+require "fileutils"
+require "tmpdir"
 
 class TestEtc < Test::Unit::TestCase
   def test_getlogin
@@ -169,6 +171,51 @@ def test_nprocessors
     assert_operator(1, :<=, n)
   end
 
+  def test_processor_count
+    assert_instance_of(Integer, Etc.processor_count)
+    assert_equal(Etc.nprocessors, Etc.processor_count)
+  end
+
+  def test_processor_quota
+    quota = Etc.processor_quota
+    assert_instance_of(Float, quota)
+    assert_operator(0.0, :<, quota)
+    assert_operator(quota, :<=, Etc.processor_count.to_f)
+  end
+
+  def test_cgroup_v2_processor_quota
+    with_cgroup_files("0::/workload\n", "cgroup2", "rw") do |mount, cgroup_path, mountinfo|
+      write_cpu_max(File.join(mount, "workload"), "150000 100000\n")
+      assert_equal(1.5, processor_cgroup.cpu_quota(cgroup_path, mountinfo))
+    end
+  end
+
+  def test_cgroup_v2_processor_quota_inherited_from_parent
+    with_cgroup_files("0::/parent/workload\n", "cgroup2", "rw") do |mount, cgroup_path, mountinfo|
+      write_cpu_max(File.join(mount, "parent"), "100000 100000\n")
+      write_cpu_max(File.join(mount, "parent", "workload"), "150000 100000\n")
+      assert_equal(1.0, processor_cgroup.cpu_quota(cgroup_path, mountinfo))
+    end
+  end
+
+  def test_cgroup_v1_processor_quota
+    with_cgroup_files("2:cpu,cpuacct:/workload\n", "cgroup", "rw,cpu,cpuacct") do |mount, cgroup_path, mountinfo|
+      write_cpu_cfs(mount, "-1\n", "100000\n")
+      write_cpu_cfs(File.join(mount, "workload"), "50000\n", "100000\n")
+      assert_equal(0.5, processor_cgroup.cpu_quota(cgroup_path, mountinfo))
+    end
+  end
+
+  def test_cgroup_processor_quota_ignores_unlimited_or_invalid_values
+    with_cgroup_files("0::/workload\n", "cgroup2", "rw") do |mount, cgroup_path, mountinfo|
+      write_cpu_max(File.join(mount, "workload"), "max 100000\n")
+      assert_nil(processor_cgroup.cpu_quota(cgroup_path, mountinfo))
+
+      write_cpu_max(File.join(mount, "workload"), "invalid\n")
+      assert_nil(processor_cgroup.cpu_quota(cgroup_path, mountinfo))
+    end
+  end
+
   def test_sysconfdir
     assert_operator(File, :absolute_path?, Etc.sysconfdir)
   end if File.method_defined?(:absolute_path?)
@@ -196,6 +243,8 @@ def test_ractor_parallel
               }
             end
             raise unless Integer === Etc.nprocessors
+            raise unless Integer === Etc.processor_count
+            raise unless Float === Etc.processor_quota
           end
         end
       end.each(&:join)
@@ -254,4 +303,34 @@ def test_ractor_getgrgid
       end.each(&:join)
     RUBY
   end
+
+  private
+
+  def processor_cgroup
+    Etc.const_get(:Cgroup, false)
+  end
+
+  def with_cgroup_files(membership, filesystem, options)
+    Dir.mktmpdir do |directory|
+      mount = File.join(directory, "cgroup mount")
+      Dir.mkdir(mount)
+      cgroup = File.join(directory, "cgroup")
+      mountinfo = File.join(directory, "mountinfo")
+      File.write(cgroup, membership)
+      escaped_mount = mount.gsub("\\") { "\\134" }.gsub(" ") { "\\040" }
+      File.write(mountinfo, "36 29 0:32 / #{escaped_mount} rw - #{filesystem} cgroup #{options}\n")
+      yield mount, cgroup, mountinfo
+    end
+  end
+
+  def write_cpu_max(directory, value)
+    FileUtils.mkdir_p(directory)
+    File.write(File.join(directory, "cpu.max"), value)
+  end
+
+  def write_cpu_cfs(directory, quota, period)
+    FileUtils.mkdir_p(directory)
+    File.write(File.join(directory, "cpu.cfs_quota_us"), quota)
+    File.write(File.join(directory, "cpu.cfs_period_us"), period)
+  end
 end