Add Crystal::System.effective_cpu_count (#16148)

The method is meant to return how many logical CPUs are actually available to the process versus the total number of CPUs from `System.cpu_count`.

- Linux: uses `sched_getaffinity` syscall
- Darwin / DragonflyBSD: _unsupported_
- FreeBSD: uses `cpuset_getaffinity`
- NetBSD / OpenBSD: uses `sysctl(CTL_HW, HW_NCPUONLINE`
- Solaris: _unknown_
- Windows: uses `GetProcessAffinityMask`

Co-authored-by: Quinton Miller <[email protected]>

Author: ysbaddaden

spec/std/crystal/system/cpucount_spec.cr

@@ -0,0 +1,8 @@
+require "spec"
+
+describe Crystal::System do
+  it "#effective_cpu_count" do
+    # smoke test: must compile and must run
+    Crystal::System.effective_cpu_count
+  end
+end

src/crystal/system.cr

@@ -3,21 +3,29 @@ module Crystal::System
   # Returns the hostname
   # def self.hostname
 
-  # Returns the number of logical processors available to the system.
-  #
+  # Returns the number of logical processors available to the system. Returns -1
+  # on errors or when unknown.
   # def self.cpu_count
+
+  # Returns the number of logical processors available to the process. Should be
+  # less than or equal to `.cpu_count`. Returns -1 on errors or when unknown.
+  def self.effective_cpu_count
+    -1
+  end
 end
 
 {% if flag?(:wasi) %}
   require "./system/wasi/hostname"
   require "./system/wasi/cpucount"
 {% elsif flag?(:unix) %}
   require "./system/unix/hostname"
-
   {% if flag?(:bsd) %}
     require "./system/unix/sysctl_cpucount"
   {% else %}
     require "./system/unix/sysconf_cpucount"
+    {% if flag?(:linux) %}
+      require "./system/unix/linux_cpucount"
+    {% end %}
   {% end %}
 {% elsif flag?(:win32) %}
   require "./system/win32/hostname"

src/crystal/system/unix/linux_cpucount.cr

@@ -0,0 +1,21 @@
+{% skip_file unless flag?(:linux) %}
+
+require "./syscall"
+
+module Crystal::System
+  def self.effective_cpu_count
+    {% unless flag?(:interpreted) %}
+      # we use the syscall because it returns the number of bytes to check in
+      # the set, while glibc always returns 0 and would require to zero the
+      # buffer and check every byte
+      set = uninitialized UInt8[8192] # allows up to 65536 logical cpus
+      byte_count = Syscall.sched_getaffinity(0, LibC::SizeT.new(8192), set.to_unsafe)
+      if byte_count > 0
+        count = set.to_slice[0, byte_count].sum(&.popcount)
+        return count if count > 0
+      end
+    {% end %}
+
+    -1
+  end
+end

src/crystal/system/unix/syscall.cr

@@ -7,4 +7,5 @@ module Crystal::System::Syscall
   GRND_NONBLOCK = 1u32
 
   ::Syscall.def_syscall getrandom, LibC::SSizeT, buf : UInt8*, buflen : LibC::SizeT, flags : UInt32
+  ::Syscall.def_syscall sched_getaffinity, LibC::Int, pid : LibC::PidT, cpusetsize : LibC::SizeT, mask : Pointer(UInt8)
 end

src/crystal/system/unix/sysctl_cpucount.cr

@@ -2,16 +2,47 @@
 
 require "c/sysctl"
 
+{% if flag?(:freebsd) %}
+  require "c/sys/cpuset"
+{% end %}
+
 module Crystal::System
   def self.cpu_count
     mib = Int32[LibC::CTL_HW, LibC::HW_NCPU]
     ncpus = 0
-    size = sizeof(Int32).to_u64
+    size = LibC::SizeT.new(sizeof(Int32))
 
     if LibC.sysctl(mib, 2, pointerof(ncpus), pointerof(size), nil, 0) == 0
       ncpus
     else
       -1
     end
   end
+
+  def self.effective_cpu_count
+    {% if flag?(:freebsd) %}
+      buffer = uninitialized UInt8[8192]
+      maxcpus = 0
+      size = LibC::SizeT.new(sizeof(Int32))
+
+      if LibC.sysctlbyname("kern.smp.maxcpus", pointerof(maxcpus), pointerof(size), nil, 0) == 0
+        len = ((maxcpus + 7) // 8).clamp(..buffer.size)
+        set = buffer.to_slice[0, len]
+
+        if LibC.cpuset_getaffinity(LibC::CPU_LEVEL_WHICH, LibC::CPU_WHICH_PID, -1, len, set) == 0
+          return set.sum(&.popcount)
+        end
+      end
+    {% elsif flag?(:netbsd) || flag?(:openbsd) %}
+      mib = Int32[LibC::CTL_HW, LibC::HW_NCPUONLINE]
+      ncpus = 0
+      size = LibC::SizeT.new(sizeof(Int32))
+
+      if LibC.sysctl(mib, 2, pointerof(ncpus), pointerof(size), nil, 0) == 0
+        return ncpus
+      end
+    {% end %}
+
+    -1
+  end
 end

src/crystal/system/wasi/cpucount.cr

@@ -3,4 +3,8 @@ module Crystal::System
     # TODO: There isn't a good way to get the number of CPUs on WebAssembly
     1
   end
+
+  def self.effective_cpu_count
+    -1
+  end
 end

src/crystal/system/win32/cpucount.cr

@@ -5,4 +5,12 @@ module Crystal::System
     LibC.GetNativeSystemInfo(out system_info)
     system_info.dwNumberOfProcessors
   end
+
+  def self.effective_cpu_count
+    if LibC.GetProcessAffinityMask(LibC.GetCurrentProcess, out process_affinity, out _) == 0
+      -1
+    else
+      process_affinity.popcount
+    end
+  end
 end

src/lib_c/x86_64-freebsd/c/sys/cpuset.cr

@@ -0,0 +1,6 @@
+lib LibC
+  CPU_LEVEL_WHICH = 3
+  CPU_WHICH_PID   = 2
+
+  fun cpuset_getaffinity(Int32, Int32, IdT, SizeT, UInt8*)
+end

src/lib_c/x86_64-freebsd/c/sysctl.cr

@@ -7,4 +7,5 @@ lib LibC
   PATH_MAX           = 1024
 
   fun sysctl(name : Int*, namelen : UInt, oldp : Void*, oldlenp : SizeT*, newp : Void*, newlen : SizeT) : Int
+  fun sysctlbyname(name : UInt8*, oldp : Void*, oldlenp : SizeT*, newp : Void*, newlen : SizeT) : Int
 end

src/lib_c/x86_64-netbsd/c/sysctl.cr

@@ -1,6 +1,7 @@
 lib LibC
-  CTL_HW  = 6
-  HW_NCPU = 3
+  CTL_HW        =  6
+  HW_NCPU       =  3
+  HW_NCPUONLINE = 16
 
   CTL_KERN           =  1
   KERN_PROC          = 14