Add support for measuring linux process shared memory

CHANGELOG.md

@@ -4,6 +4,11 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## 1.2.0
+
+### Added
+- Add support on Linux for measuring process memory using unique size, proportional size, or resident size.
+
 ## 1.1.2
 
 ### Added

README.md

@@ -46,6 +46,14 @@ For a Rails application, you would normally want to preboot the `config/boot.rb`
 
 You can also specify a maximum memory footprint that you want to allow for each child process. You can use this feature to automatically guard against poorly designed workers that bloat the Ruby memory heap. Note that you can also use an external process monitor to kill processes with memory bloat; the process manager will restart any process regardless of how it exits.
 
+You can also specify how to measure process memory on a Linux system depending on how you want to treat shared memory.
+
+- rss - measure resident memory usage ([resident set size](https://en.wikipedia.org/wiki/Resident_set_size))
+- uss - measure only private (non-shared) memory ([unique set size](https://en.wikipedia.org/wiki/Unique_set_size))
+- pss - measure shared memory proportionally for all processes sharing it ([proportional set size](https://en.wikipedia.org/wiki/Proportional_set_size))
+
+The default is to measure resident memory. Non-Linux systems will always measure resident memory.
+
 ## Usage
 
 Install the gem in your sidekiq process and run it with `bundle exec sidekiq-process-manager` or, if you use [bundle binstubs](https://bundler.io/man/bundle-binstubs.1.html), `bin/sidekiq-process-manager`. Command line arguments are passed through to `sidekiq`. If you want to supply on of the `sidekiq-process_manager` specific options, those options should come first and the `sidekiq` options should appear after a `--` flag
@@ -97,16 +105,16 @@ or
 SIDEKIQ_PREBOOT=config/boot.rb SIDEKIQ_PROCESSES=4 bundle exec sidekiq-process-manager
 ```
 
-You can set the maximum memory allowed per sidekiq process with the `--max-memory` argument or with the `SIDEKIQ_MAX_MEMORY` environment variable. You can suffix the value with "m" to specify megabytes or "g" to specify gigabytes.
+You can set the maximum memory allowed per sidekiq process with the `--max-memory` argument or with the `SIDEKIQ_MAX_MEMORY` environment variable. You can suffix the value with "m" to specify megabytes or "g" to specify gigabytes. You can specify the method for measuring process memory on Linux with `--memory-measurement` or the `SIDEKIQ_MEMORY_MEASUREMENT` environment variable.
 
 ```bash
-bundle exec sidekiq-process-manager --processes 4 --max-memory 2g
+bundle exec sidekiq-process-manager --processes 4 --max-memory 2g --memory-measurement rss
 ```
 
 or
 
 ```bash
-SIDEKIQ_MAX_MEMORY=2000m SIDEKIQ_PROCESSES=4 bundle exec sidekiq-process-manager
+SIDEKIQ_MAX_MEMORY=2000m SIDEKIQ_MEMORY_MEASUREMENT=rss SIDEKIQ_PROCESSES=4 bundle exec sidekiq-process-manager
 ```
 
 ## Alternatives

VERSION

@@ -1 +1 @@
-1.1.2
+1.2.0

bin/sidekiq-process-manager

@@ -27,6 +27,7 @@ options = {
   prefork: !ENV.fetch("SIDEKIQ_PREFORK", "").empty?,
   preboot: ENV["SIDEKIQ_PREBOOT"],
   max_memory: parse_max_memory(ENV["SIDEKIQ_MAX_MEMORY"]),
+  memory_measurement: ENV.fetch("SIDEKIQ_MEMORY_MEASUREMENT", "rss"),
   mode: nil,
 }
 
@@ -49,6 +50,10 @@ parser = OptionParser.new do |opts|
     options[:max_memory] = parse_max_memory(max_memory)
   end
 
+  opts.on('--memory-measurement rss|uss|pss', "How to measure process memory on a Linux system; this can be rss (resident), uss (unique/non-shared), or pss (proportional)") do |measurement|
+    options[:memory_measurement] = measurement
+  end
+
   opts.on('--testing', "Enable test mode") do |testing|
     options[:mode] = :testing if testing
   end

lib/sidekiq/process_manager/manager.rb

@@ -2,6 +2,7 @@
 
 require "sidekiq"
 require "get_process_mem"
+require "linux_process_memory"
 
 module Sidekiq
   module ProcessManager
@@ -16,9 +17,12 @@ class Manager
       # @param process_count [Integer] The number of sidekiq processes to start.
       # @param prefork [Boolean] If true, the process manager will load the application before forking.
       # @param preboot [String] If set, the process manager will require the specified file before forking the child processes.
+      # @param max_memory [Integer] If set, the process manager will kill any child processes that exceed the specified memory limit in bytes.
+      # @param memory_measurement [String] If set, the process manager will use the specified memory measurement to determine memory usage.
+      #   Valid values are "rss", "uss", and "pss". Defaults to "rss". This setting only works on Linux. Other platforms will use rss.
       # @param mode [Symbol] If set to :testing, the process manager will use a mock CLI.
       # @param silent [Boolean] If true, the process manager will not output any messages.
-      def initialize(process_count: 1, prefork: false, preboot: nil, max_memory: nil, mode: nil, silent: false)
+      def initialize(process_count: 1, prefork: false, preboot: nil, max_memory: nil, memory_measurement: nil, mode: nil, silent: false)
         require "sidekiq/cli"
 
         # Get the number of processes to fork
@@ -151,6 +155,27 @@ def started?
         @started
       end
 
+      # Get the amount of memory being used by a process.
+      #
+      # @param pid [Integer] The process id.
+      # @param measurement [String] The type of memory measurement to use (rss, uss, or pss).
+      # @return [Integer] The amount of memory used in bytes.
+      def get_process_memory(pid, measurement)
+        if LinuxProcessMemory.supported?
+          memory = LinuxProcessMemory.new(pid)
+          case measurement.to_s.downcase
+          when "uss"
+            memory.uss
+          when "pss"
+            memory.pss
+          else
+            memory.rss
+          end
+        else
+          GetProcessMem.new(pid).bytes
+        end
+      end
+
       private
 
       def log_info(message)
@@ -271,9 +296,10 @@ def start_memory_monitor
 
             pids.each do |pid|
               begin
-                memory = GetProcessMem.new(pid)
+                memory = get_process_memory(pid, @memory_measurement)
                 if memory.bytes > @max_memory
-                  log_warning("Killing bloated sidekiq process #{pid}: #{memory.mb.round}mb used")
+                  megabytes = (memory.to_f / (1024**2)).round
+                  log_warning("Killing bloated sidekiq process #{pid}: #{megabytes}mb used")
                   send_signal_to_pid(:TERM, pid)
                   break
                 end

sidekiq-process_manager.gemspec

@@ -33,6 +33,7 @@ Gem::Specification.new do |spec|
 
   spec.add_dependency "sidekiq", ">= 5.0"
   spec.add_dependency "get_process_mem"
+  spec.add_dependency "linux_process_memory"
 
   spec.add_development_dependency "bundler"
 end

spec/sidekiq/process_manager/manager_spec.rb

@@ -78,9 +78,39 @@
       expect(::Process).to receive(:kill).with(:TERM, pids.first).and_call_original
       sleep(2)
       manager.wait
-      # This check is flakey with Sidekiq 6.0 and below
-      if Gem::Version.new(Sidekiq::VERSION) >= Gem::Version.new("6.5")
-        expect(manager.pids).to_not include(pids.first)
+      sleep(2)
+    end
+  end
+
+  describe "memory measurement" do
+    context "on Linux" do
+      before do
+        allow(LinuxProcessMemory).to receive(:supported?).and_return(true)
+      end
+
+      it "uses resident memory by default" do
+        expect_any_instance_of(LinuxProcessMemory).to receive(:rss).and_call_original
+        expect(manager.get_process_memory(1, nil)).to be_a Integer
+      end
+
+      it "can use unique memory measurements" do
+        expect_any_instance_of(LinuxProcessMemory).to receive(:uss).and_call_original
+        expect(manager.get_process_memory(1, "uss")).to be_a Integer
+      end
+
+      it "can use proportional memory measurements" do
+        expect_any_instance_of(LinuxProcessMemory).to receive(:pss).and_call_original
+        expect(manager.get_process_memory(1, "pss")).to be_a Integer
+      end
+    end
+
+    context "on on-Linux systems" do
+      before do
+        allow(LinuxProcessMemory).to receive(:supported?).and_return(false)
+      end
+
+      it "uses resident memory" do
+        expect(manager.get_process_memory(Process.pid, "pss")).to be > 0
       end
     end
   end