sidekiq-worker-killer 1.0.0 → 1.1.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: f40045c5805f35f8792f66c5dac896042c3b50b7b8c42f0f95e14827d4184a5a
-  data.tar.gz: 5721815ad342ec06d6231c005b60cb0bbda5300343a4bfb99b9b138fff8f6627
+  metadata.gz: 1335f7577f6a507808c59651cc9f479a49a9d54e144906b39b06355af0fcfcba
+  data.tar.gz: f61243439575d5ec2f75e0e7e803ff6588279b349a4ab9ebabc62e58fb3fcb7e
 SHA512:
-  metadata.gz: 4db165333c1a8465b15e514af6b178aec3091bdb90a2ced38b8131b98706922518b25632f45a65fa33b923ac6d2a42c3379dee4e4fde1dde9dd20d9a49aab1e1
-  data.tar.gz: dd3e9e16d4bfbc5d753559e6b94c95b334ef2855864bcc7ef53f1577e40f6bc5122ad7040133c04d2c019ab37d532f012339fb18180e8efe74dc15274d427f49
+  metadata.gz: 97742129013c4374090ed58f35d9b407f6dc2b9ff44a23d4de187b7d3035c11961ebb089b0753bf7b1688d386d77b2aef7cceba4344dd614a2dd3181fc28a0c2
+  data.tar.gz: cbf8b55a9bbf3171f988c0eabb4a137bbe3e458f9aa8368cb2c1cbe3759a14ce4b6dd919b9568458e48ac473c7780f01ad72be437503ee8e3695e1b67613f59d

data/README.md

@@ -5,7 +5,18 @@
 
 [Sidekiq](https://github.com/mperham/sidekiq) is probably the best background processing framework today. At the same time, memory leaks are very hard to tackle in Ruby and we often find ourselves with growing memory consumption. Instead of spending herculean effort fixing leaks, why not kill your processes when they got to be too large?
 
-Highly inspired by [Gitlab Sidekiq MemoryKiller](https://gitlab.com/gitlab-org/gitlab-ce/blob/master/lib/gitlab/sidekiq_middleware/shutdown.rb) and [Noxa Sidekiq killer](https://github.com/Noxa/sidekiq-killer).
+Highly inspired by [Gitlab Sidekiq MemoryKiller](https://gitlab.com/gitlab-org/gitlab-foss/-/blob/39c1731a53d1014eab7c876d70632b1abf738712/lib/gitlab/sidekiq_middleware/shutdown.rb) and [Noxa Sidekiq killer](https://github.com/Noxa/sidekiq-killer).
+
+---
+**NOTE**
+
+This gem needs to get the used memory of the Sidekiq process. For
+this we use [GetProcessGem](https://github.com/schneems/get_process_mem),
+but be aware that if you are running Sidekiq in Heroku(or any container) the
+memory usage will
+[not be accurate](https://github.com/schneems/get_process_mem/issues/7).
+
+---
 
 quick-refs: [install](#install) | [usage](#usage) | [available options](#available-options) | [development](#development)
 
@@ -31,16 +42,17 @@ end
 
 ## Available options
 
-The following options can be overrided.
+The following options can be overridden.
 
 | Option | Defaults | Description |
 | ------- | ------- | ----------- |
-| max_rss | 0 MB (disabled) | max RSS in megabytes. Above this, shutdown will be triggered. |
-| grace_time | 900 seconds | when shutdown is triggered, the Sidekiq process will not accept new job and wait at most 15 minutes for running jobs to finish. If Float::INFINITY specified, will wait forever. |
-| shutdown_wait | 30 seconds | when the grace time expires, still running jobs get 30 seconds to stop. After that, kill signal is triggered. |
+| max_rss | 0 MB (disabled) | Max RSS in megabytes used by the Sidekiq process. Above this, shutdown will be triggered. |
+| grace_time | 900 seconds | When shutdown is triggered, the Sidekiq process will not accept new job and wait at most 15 minutes for running jobs to finish. If Float::INFINITY specified, will wait forever. |
+| shutdown_wait | 30 seconds | When the grace time expires, still running jobs get 30 seconds to stop. After that, kill signal is triggered. |
 | kill_signal | SIGKILL | Signal to use to kill Sidekiq process if it doesn't stop. |
 | gc | true | Try to run garbage collection before Sidekiq process stops in case of exceeded max_rss. |
 | skip_shutdown_if | proc {false} | Executes a block of code after max_rss exceeds but before requesting shutdown. |
+| on_shutdown | nil | Executes a block of code before just before requesting shutdown. This can be used to send custom logs or metrics to external services. |
 
 *skip_shutdown_if* is expected to return anything other than `false` or `nil` to skip shutdown.
 
@@ -75,4 +87,4 @@ See the list of [contributors](https://github.com/klaxit/sidekiq-worker-killer/c
 
 ## License
 
-Please see LICENSE
+Please see [LICENSE](https://github.com/klaxit/sidekiq-worker-killer/blob/master/LICENSE)

data/lib/sidekiq/worker_killer/version.rb

@@ -1,7 +1,7 @@
 # rubocop:disable Style/ClassAndModuleChildren
 module Sidekiq
   class WorkerKiller
-    VERSION = "1.0.0".freeze
+    VERSION = "1.1.0".freeze
   end
 end
 # rubocop:enable Style/ClassAndModuleChildren

data/lib/sidekiq/worker_killer.rb

@@ -1,12 +1,18 @@
 require "get_process_mem"
 require "sidekiq"
-require "sidekiq/util"
+begin
+  require "sidekiq/util"
+  SidekiqComponent = Sidekiq::Util
+rescue LoadError
+  require "sidekiq/middleware/modules"
+  SidekiqComponent = Sidekiq::ServerMiddleware
+end
 require "sidekiq/api"
 
 # Sidekiq server middleware. Kill worker when the RSS memory exceeds limit
 # after a given grace time.
 class Sidekiq::WorkerKiller
-  include Sidekiq::Util
+  include SidekiqComponent
 
   MUTEX = Mutex.new
 
@@ -30,6 +36,9 @@ class Sidekiq::WorkerKiller
   # @option options [Proc] skip_shutdown_if
   #   Executes a block of code after max_rss exceeds but before requesting
   #   shutdown. (default: `proc {false}`)
+  # @option options [Proc] on_shutdown
+  #   Executes a block of code right before a shutdown happens.
+  #   (default: `nil`)
   def initialize(options = {})
     @max_rss         = options.fetch(:max_rss, 0)
     @grace_time      = options.fetch(:grace_time, 15 * 60)
@@ -37,6 +46,7 @@ class Sidekiq::WorkerKiller
     @kill_signal     = options.fetch(:kill_signal, "SIGKILL")
     @gc              = options.fetch(:gc, true)
     @skip_shutdown   = options.fetch(:skip_shutdown_if, proc { false })
+    @on_shutdown     = options.fetch(:on_shutdown, nil)
   end
 
   # @param [String, Class] worker_class
@@ -62,17 +72,23 @@ class Sidekiq::WorkerKiller
 
     warn "current RSS #{current_rss} of #{identity} exceeds " \
          "maximum RSS #{@max_rss}"
+
+    run_shutdown_hook(worker, job, queue)
     request_shutdown
   end
 
   private
 
+  def run_shutdown_hook(worker, job, queue)
+    @on_shutdown.respond_to?(:call) && @on_shutdown.call(worker, job, queue)
+  end
+
   def skip_shutdown?(worker, job, queue)
     @skip_shutdown.respond_to?(:call) && @skip_shutdown.call(worker, job, queue)
   end
 
   def request_shutdown
-    # In another thread to allow undelying job to finish
+    # In another thread to allow underlying job to finish
     Thread.new do
       # Only if another thread is not already
       # shutting down the Sidekiq process
@@ -125,6 +141,10 @@ class Sidekiq::WorkerKiller
     end || raise("No sidekiq worker with identity #{identity} found")
   end
 
+  def identity
+    config[:identity] || config["identity"]
+  end unless method_defined?(:identity)
+
   def warn(msg)
     Sidekiq.logger.warn(msg)
   end

data/spec/sidekiq/worker_killer_spec.rb

@@ -7,156 +7,235 @@ describe Sidekiq::WorkerKiller do
   before do
     allow(subject).to receive(:warn) # silence "warn" logs
     allow(subject).to receive(:sleep) # reduces tests running time
-    allow(Sidekiq::ProcessSet).to receive(:new) { sidekiq_process_set }
-    allow(sidekiq_process_set).to receive(:find) { sidekiq_process }
     allow(sidekiq_process).to receive(:quiet!)
     allow(sidekiq_process).to receive(:stop!)
     allow(sidekiq_process).to receive(:stopping?)
   end
 
-  describe "#call" do
-    let(:worker){ double("worker") }
-    let(:job){ double("job") }
-    let(:queue){ double("queue") }
+  describe "#sidekiq_process" do
+    let(:identity) do
+      "foobar"
+    end
+    let(:mock_processes) do
+      [{
+        "identity" => identity,
+        "tag" => "baz",
+        "started_at" => Time.now,
+        "concurrency" => 5,
+        "busy" => 2,
+        "queues" => %w[low medium high]
+      }]
+    end
+    it "finds the process by identity" do
+      if subject.respond_to?(:config=)
+        rspec_info "variant 6.5+: Sidekiq::ServerMiddleware"
+        subject.config = mock_processes.first
+      else
+        rspec_info "variant pre 6.5: Sidekiq::Util"
+        allow(subject).to receive(:identity).and_return(identity)
+      end
+      allow(Sidekiq).to receive(:redis)
+      sidekiq_process_set = Sidekiq::ProcessSet.new
+      sidekiq_mock_processes = mock_processes.map {|mock_process|
+        Sidekiq::Process.new(mock_process)
+      }
+      allow(sidekiq_process_set).to receive(:each) {|&block|
+        sidekiq_mock_processes.each(&block)
+      }
+      expect(sidekiq_process_set.to_a.map(&:identity)).to eq([identity])
+      allow(Sidekiq::ProcessSet).to receive(:new).and_return(sidekiq_process_set)
+      expect(subject.send(:identity)).to eq(identity)
+      expect(subject.send(:sidekiq_process)).to be_a(Sidekiq::Process)
+      expect(subject.send(:sidekiq_process).identity).to eq(identity)
+    end
+  end
 
-    it "should yield" do
-      expect { |b|
-        subject.call(worker, job, queue, &b)
-      }.to yield_with_no_args
+  context "with stubbed sidekiq_process" do
+    before do
+      allow(Sidekiq::ProcessSet).to receive(:new) { sidekiq_process_set }
+      allow(sidekiq_process_set).to receive(:find) { sidekiq_process }
+      identity = "some-identity"
+      if subject.respond_to?(:config=)
+        rspec_info "variant 6.5+: Sidekiq::ServerMiddleware"
+        subject.config = { "identity" => identity }
+      else
+        rspec_info "variant pre 6.5: Sidekiq::Util"
+        allow(subject).to receive(:identity).and_return(identity)
+      end
     end
 
-    context "when current rss is over max rss" do
-      subject{ described_class.new(max_rss: 2) }
+    describe "#call" do
+      let(:worker){ double("worker") }
+      let(:job){ double("job") }
+      let(:queue){ double("queue") }
 
-      before do
-        allow(subject).to receive(:current_rss).and_return(3)
+      it "should yield" do
+        expect { |b|
+          subject.call(worker, job, queue, &b)
+        }.to yield_with_no_args
       end
 
-      it "should request shutdown" do
-        expect(subject).to receive(:request_shutdown)
-        subject.call(worker, job, queue){}
-      end
+      context "when current rss is over max rss" do
+        subject{ described_class.new(max_rss: 2) }
 
-      it "should call garbage collect" do
-        allow(subject).to receive(:request_shutdown)
-        expect(GC).to receive(:start).with(full_mark: true, immediate_sweep: true)
-        subject.call(worker, job, queue){}
-      end
+        before do
+          allow(subject).to receive(:current_rss).and_return(3)
+        end
 
-      context "and skip_shutdown_if is given" do
-        subject{ described_class.new(max_rss: 2, skip_shutdown_if: skip_shutdown_proc) }
+        it "should request shutdown" do
+          expect(subject).to receive(:request_shutdown)
+          subject.call(worker, job, queue){}
+        end
 
-        context "and skip_shutdown_if is a proc" do
-          let(:skip_shutdown_proc) { proc { |worker| true } }
-          it "should NOT request shutdown" do
-            expect(subject).not_to receive(:request_shutdown)
-            subject.call(worker, job, queue){}
+        it "should call garbage collect" do
+          allow(subject).to receive(:request_shutdown)
+          expect(GC).to receive(:start).with(full_mark: true, immediate_sweep: true)
+          subject.call(worker, job, queue){}
+        end
+
+        context "and skip_shutdown_if is given" do
+          subject{ described_class.new(max_rss: 2, skip_shutdown_if: skip_shutdown_proc) }
+
+          context "and skip_shutdown_if is a proc" do
+            let(:skip_shutdown_proc) { proc { |worker| true } }
+            it "should NOT request shutdown" do
+              expect(subject).not_to receive(:request_shutdown)
+              subject.call(worker, job, queue){}
+            end
+          end
+
+          context "and skip_shutdown_if is a lambda" do
+            let(:skip_shutdown_proc) { ->(worker, job, queue) { true } }
+            it "should NOT request shutdown" do
+              expect(subject).not_to receive(:request_shutdown)
+              subject.call(worker, job, queue){}
+            end
+          end
+
+          context "and skip_shutdown_if returns false" do
+            let(:skip_shutdown_proc) { proc { |worker, job, queue| false } }
+            it "should still request shutdown" do
+              expect(subject).to receive(:request_shutdown)
+              subject.call(worker, job, queue){}
+            end
+          end
+
+          context "and skip_shutdown_if returns nil" do
+            let(:skip_shutdown_proc) { proc { |worker, job, queue| nil } }
+            it "should still request shutdown" do
+              expect(subject).to receive(:request_shutdown)
+              subject.call(worker, job, queue){}
+            end
           end
         end
 
-        context "and skip_shutdown_if is a lambda" do
-          let(:skip_shutdown_proc) { ->(worker, job, queue) { true } }
-          it "should NOT request shutdown" do
-            expect(subject).not_to receive(:request_shutdown)
-            subject.call(worker, job, queue){}
+        context "and on_shutdown is given" do
+          subject{ described_class.new(max_rss: 2, on_shutdown: on_shutdown_proc) }
+
+          context "and on_shutdown is a proc" do
+            let(:on_shutdown_proc) { proc { |worker, job, queue| nil } }
+            it "should execute on_shutdown hook" do
+              expect(subject).to receive(:request_shutdown).once
+              expect(on_shutdown_proc).to receive(:call).once
+              subject.call(worker, job, queue){}
+            end
+          end
+
+          context "and on_shutdown is a lambda" do
+            let(:on_shutdown_proc) { ->(worker, job, queue) { nil } }
+            it "should execute on_shutdown hook" do
+              expect(subject).to receive(:request_shutdown).once
+              expect(on_shutdown_proc).to receive(:call).once
+              subject.call(worker, job, queue){}
+            end
           end
         end
 
-        context "and skip_shutdown_if returns false" do
-          let(:skip_shutdown_proc) { proc { |worker, job, queue| false } }
-          it "should still request shutdown" do
-            expect(subject).to receive(:request_shutdown)
+        context "when gc is false" do
+          subject{ described_class.new(max_rss: 2, gc: false) }
+          it "should not call garbage collect" do
+            allow(subject).to receive(:request_shutdown)
+            expect(GC).not_to receive(:start)
             subject.call(worker, job, queue){}
           end
         end
 
-        context "and skip_shutdown_if returns nil" do
-          let(:skip_shutdown_proc) { proc { |worker, job, queue| nil } }
-          it "should still request shutdown" do
-            expect(subject).to receive(:request_shutdown)
+        context "but max rss is 0" do
+          subject{ described_class.new(max_rss: 0) }
+          it "should not request shutdown" do
+            expect(subject).to_not receive(:request_shutdown)
             subject.call(worker, job, queue){}
           end
         end
       end
+    end
 
-      context "when gc is false" do
-        subject{ described_class.new(max_rss: 2, gc: false) }
-        it "should not call garbage collect" do
-          allow(subject).to receive(:request_shutdown)
-          expect(GC).not_to receive(:start)
-          subject.call(worker, job, queue){}
+    describe "#request_shutdown" do
+      context "grace time is default" do
+        before { allow(subject).to receive(:shutdown){ sleep 0.01 } }
+        it "should call shutdown" do
+          expect(subject).to receive(:shutdown)
+          subject.send(:request_shutdown).join
         end
-      end
-
-      context "but max rss is 0" do
-        subject{ described_class.new(max_rss: 0) }
-        it "should not request shutdown" do
-          expect(subject).to_not receive(:request_shutdown)
-          subject.call(worker, job, queue){}
+        it "should not call shutdown twice when called concurrently" do
+          expect(subject).to receive(:shutdown).once
+          2.times.map{ subject.send(:request_shutdown) }.each(&:join)
         end
       end
-    end
-  end
-
-  describe "#request_shutdown" do
-    context "grace time is default" do
-      before { allow(subject).to receive(:shutdown){ sleep 0.01 } }
-      it "should call shutdown" do
-        expect(subject).to receive(:shutdown)
-        subject.send(:request_shutdown).join
-      end
-      it "should not call shutdown twice when called concurrently" do
-        expect(subject).to receive(:shutdown).once
-        2.times.map{ subject.send(:request_shutdown) }.each(&:join)
-      end
-    end
 
-    context "grace time is 5 seconds" do
-      subject{ described_class.new(max_rss: 2, grace_time: 5.0, shutdown_wait: 0) }
-      it "should wait the specified grace time before calling shutdown" do
-        # there are busy jobs that will not terminate within the grace time
-        allow(subject).to receive(:no_jobs_on_quiet_processes?).and_return(false)
-
-        shutdown_request_time     = nil
-        shutdown_time             = nil
-
-        # replace the original #request_shutdown to track
-        # when the shutdown is requested
-        original_request_shutdown = subject.method(:request_shutdown)
-        allow(subject).to receive(:request_shutdown) do
-          shutdown_request_time= Time.now
-          original_request_shutdown.call
-        end
+      context "grace time is 5 seconds" do
+        subject{ described_class.new(max_rss: 2, grace_time: 5.0, shutdown_wait: 0) }
+        it "should wait the specified grace time before calling shutdown" do
+          # there are busy jobs that will not terminate within the grace time
+          allow(subject).to receive(:no_jobs_on_quiet_processes?).and_return(false)
+
+          shutdown_request_time     = nil
+          shutdown_time             = nil
+
+          # replace the original #request_shutdown to track
+          # when the shutdown is requested
+          original_request_shutdown = subject.method(:request_shutdown)
+          allow(subject).to receive(:request_shutdown) do
+            shutdown_request_time= Time.now
+            original_request_shutdown.call
+          end
 
-        # track when the process has been required to stop
-        expect(sidekiq_process).to receive(:stop!) do |*args|
-          shutdown_time = Time.now
-        end
+          # track when the process has been required to stop
+          expect(sidekiq_process).to receive(:stop!) do |*args|
+            shutdown_time = Time.now
+          end
 
-        allow(Process).to receive(:kill)
-        allow(Process).to receive(:pid).and_return(99)
+          allow(Process).to receive(:kill)
+          allow(Process).to receive(:pid).and_return(99)
 
-        subject.send(:request_shutdown).join
+          subject.send(:request_shutdown).join
 
-        elapsed_time = shutdown_time - shutdown_request_time
+          elapsed_time = shutdown_time - shutdown_request_time
 
-        # the elapsed time beetween shutdown request and the actual
-        # shutdown signal should be greater than the specificed grace_time
-        expect(elapsed_time).to be >= 5.0
+          # the elapsed time between shutdown request and the actual
+          # shutdown signal should be greater than the specified grace_time
+          expect(elapsed_time).to be >= 5.0
+        end
       end
-    end
 
-    context "grace time is Float::INFINITY" do
-      subject{ described_class.new(max_rss: 2, grace_time: Float::INFINITY, shutdown_wait: 0) }
-      it "call signal only on jobs" do
-        allow(subject).to receive(:jobs_finished?).and_return(true)
-        allow(Process).to receive(:pid).and_return(99)
-        expect(sidekiq_process).to receive(:quiet!)
-        expect(sidekiq_process).to receive(:stop!)
-        expect(Process).to receive(:kill).with('SIGKILL', 99)
+      context "grace time is Float::INFINITY" do
+        subject{ described_class.new(max_rss: 2, grace_time: Float::INFINITY, shutdown_wait: 0) }
+        it "call signal only on jobs" do
+          allow(subject).to receive(:jobs_finished?).and_return(true)
+          allow(Process).to receive(:pid).and_return(99)
+          expect(sidekiq_process).to receive(:quiet!)
+          expect(sidekiq_process).to receive(:stop!)
+          expect(Process).to receive(:kill).with('SIGKILL', 99)
 
-        subject.send(:request_shutdown).join
+          subject.send(:request_shutdown).join
+        end
       end
     end
   end
+
+  def rspec_info(msg)
+    is_printing_progress = RSpec.configuration.formatters.any? { |f| /ProgressFormatter/.match?(f.class.name) }
+    return if is_printing_progress
+    RSpec.configuration.reporter.message msg
+  end
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: sidekiq-worker-killer
 version: !ruby/object:Gem::Version
-  version: 1.0.0
+  version: 1.1.0
 platform: ruby
 authors:
 - Cyrille Courtiere
-autorequire: 
+autorequire:
 bindir: bin
 cert_chain: []
-date: 2020-05-28 00:00:00.000000000 Z
+date: 2023-01-20 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: get_process_mem
@@ -66,7 +66,21 @@ dependencies:
     - - "~>"
       - !ruby/object:Gem::Version
         version: 0.49.1
-description: 
+- !ruby/object:Gem::Dependency
+  name: appraisal
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+description:
 email:
 - dev@klaxit.com
 executables: []
@@ -81,7 +95,7 @@ files:
 homepage: http://github.com/klaxit/sidekiq-worker-killer
 licenses: []
 metadata: {}
-post_install_message: 
+post_install_message:
 rdoc_options: []
 require_paths:
 - lib
@@ -96,8 +110,8 @@ required_rubygems_version: !ruby/object:Gem::Requirement
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 3.1.0.pre3
-signing_key: 
+rubygems_version: 3.3.26
+signing_key:
 specification_version: 4
 summary: Sidekiq worker killer
 test_files: