From 947aebf201cf77db45d0937eed6093e907bcb17c Mon Sep 17 00:00:00 2001
From: heka1024 <heka1024@gmail.com>
Date: Wed, 31 Jul 2024 02:22:29 +0900
Subject: [PATCH] Add `Concurrent.cpu_shares` that is cgroups aware.

---
 .../concurrent/utility/processor_counter.rb   | 27 +++++++++++++++++++
 .../utility/processor_count_spec.rb           | 22 +++++++++++++++
 2 files changed, 49 insertions(+)

diff --git a/lib/concurrent-ruby/concurrent/utility/processor_counter.rb b/lib/concurrent-ruby/concurrent/utility/processor_counter.rb
index 007668ae6..2dd91b59a 100644
--- a/lib/concurrent-ruby/concurrent/utility/processor_counter.rb
+++ b/lib/concurrent-ruby/concurrent/utility/processor_counter.rb
@@ -12,6 +12,7 @@ def initialize
         @processor_count          = Delay.new { compute_processor_count }
         @physical_processor_count = Delay.new { compute_physical_processor_count }
         @cpu_quota                = Delay.new { compute_cpu_quota }
+        @cpu_shares               = Delay.new { compute_cpu_shares }
       end
 
       def processor_count
@@ -41,6 +42,10 @@ def cpu_quota
         @cpu_quota.value
       end
 
+      def cpu_shares
+        @cpu_shares.value
+      end
+
       private
 
       def compute_processor_count
@@ -113,6 +118,20 @@ def compute_cpu_quota
           end
         end
       end
+
+      def compute_cpu_shares
+        if RbConfig::CONFIG["target_os"].include?("linux")
+          if File.exist?("/sys/fs/cgroup/cpu.weight")
+            # cgroups v2: https://docs.kernel.org/admin-guide/cgroup-v2.html#cpu-interface-files
+            # Ref: https://github.com/kubernetes/enhancements/tree/master/keps/sig-node/2254-cgroup-v2#phase-1-convert-from-cgroups-v1-settings-to-v2
+            weight = File.read("/sys/fs/cgroup/cpu.weight").to_f
+            ((((weight - 1) * 262142) / 9999) + 2) / 1024
+          elsif File.exist?("/sys/fs/cgroup/cpu/cpu.shares")
+            # cgroups v1: https://kernel.googlesource.com/pub/scm/linux/kernel/git/glommer/memcg/+/cpu_stat/Documentation/cgroups/cpu.txt
+            File.read("/sys/fs/cgroup/cpu/cpu.shares").to_f / 1024
+          end
+        end
+      end
     end
   end
 
@@ -187,4 +206,12 @@ def self.available_processor_count
   def self.cpu_quota
     processor_counter.cpu_quota
   end
+
+  # The CPU shares requested by the process. For performance reasons the calculated
+  # value will be memoized on the first call.
+  #
+  # @return [Float, nil] CPU shares requested by the process, or nil if not set
+  def self.cpu_shares
+    processor_counter.cpu_shares
+  end
 end
diff --git a/spec/concurrent/utility/processor_count_spec.rb b/spec/concurrent/utility/processor_count_spec.rb
index fdc44b0ae..7a636640c 100644
--- a/spec/concurrent/utility/processor_count_spec.rb
+++ b/spec/concurrent/utility/processor_count_spec.rb
@@ -92,4 +92,26 @@ module Concurrent
     end
 
   end
+
+  RSpec.describe '#cpu_shares' do
+    let(:counter) { Concurrent::Utility::ProcessorCounter.new }
+
+    it 'returns a float when cgroups v2 sets a cpu.weight' do
+      expect(RbConfig::CONFIG).to receive(:[]).with("target_os").and_return("linux")
+      expect(File).to receive(:exist?).with("/sys/fs/cgroup/cpu.weight").and_return(true)
+
+      expect(File).to receive(:read).with("/sys/fs/cgroup/cpu.weight").and_return("10000\n")
+      expect(counter.cpu_shares).to be == 256.0
+    end
+
+    it 'returns a float if cgroups v1 sets a cpu.shares' do
+      expect(RbConfig::CONFIG).to receive(:[]).with("target_os").and_return("linux")
+      expect(File).to receive(:exist?).with("/sys/fs/cgroup/cpu.weight").and_return(false)
+      expect(File).to receive(:exist?).with("/sys/fs/cgroup/cpu/cpu.shares").and_return(true)
+
+      expect(File).to receive(:read).with("/sys/fs/cgroup/cpu/cpu.shares").and_return("512\n")
+      expect(counter.cpu_shares).to be == 0.5
+    end
+
+  end
 end