sidekiq-worker-killer 0.2.0 → 1.0.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
-SHA1:
-  metadata.gz: 643766850e2b525cc9458efaeb259b2525654235
-  data.tar.gz: 25cc9a461b70258f9ce09ab71fa4d59f48150744
+SHA256:
+  metadata.gz: f40045c5805f35f8792f66c5dac896042c3b50b7b8c42f0f95e14827d4184a5a
+  data.tar.gz: 5721815ad342ec06d6231c005b60cb0bbda5300343a4bfb99b9b138fff8f6627
 SHA512:
-  metadata.gz: 13e54d178b6eaf6de4c18bff3cb7bf5a0f583a2a511880514cb0b50c4158a8acf93dad622f3d906304e6b4a581c1d4a9b34a9e100354e0d3984891fb09ce34ec
-  data.tar.gz: e538edc1f1d17a2ffcba3121551bde5234375eb1648e98c1790efd973e1f8934131ad3ec9ecc8ace4101e2ec0ca97c62908e04b62b7b3243c6676a3511458d70
+  metadata.gz: 4db165333c1a8465b15e514af6b178aec3091bdb90a2ced38b8131b98706922518b25632f45a65fa33b923ac6d2a42c3379dee4e4fde1dde9dd20d9a49aab1e1
+  data.tar.gz: dd3e9e16d4bfbc5d753559e6b94c95b334ef2855864bcc7ef53f1577e40f6bc5122ad7040133c04d2c019ab37d532f012339fb18180e8efe74dc15274d427f49

data/README.md

@@ -3,13 +3,15 @@
 [![Gem Version](https://badge.fury.io/rb/sidekiq-worker-killer.svg)](https://badge.fury.io/rb/sidekiq-worker-killer)
 [![CircleCI](https://circleci.com/gh/klaxit/sidekiq-worker-killer.svg?style=shield&circle-token=:circle-token)](https://circleci.com/gh/klaxit/sidekiq-worker-killer)
 
-[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.
+[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).
 
+quick-refs: [install](#install) | [usage](#usage) | [available options](#available-options) | [development](#development)
+
 ## Install
 Use [Bundler](http://bundler.io/)
-```
+```ruby
 gem "sidekiq-worker-killer"
 ```
 
@@ -17,7 +19,9 @@ gem "sidekiq-worker-killer"
 
 Add this to your Sidekiq configuration.
 
-```
+```ruby
+require 'sidekiq/worker_killer'
+
 Sidekiq.configure_server do |config|
   config.server_middleware do |chain|
     chain.add Sidekiq::WorkerKiller, max_rss: 480
@@ -25,16 +29,45 @@ Sidekiq.configure_server do |config|
 end
 ```
 
-# Available options
+## Available options
 
 The following options can be overrided.
 
 | 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 but wait 15 minutes for running jobs to finish.  |
-| shutdown_wait | 30 seconds | when the grace time expires, still running jobs get 30 seconds to terminate. After that, kill signal is triggered.  |
-| kill_signal | SIGKILL | Signal to use kill Sidekiq process if it doesn't terminate.  |
+| 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. |
+
+*skip_shutdown_if* is expected to return anything other than `false` or `nil` to skip shutdown.
+
+```ruby
+require 'sidekiq/worker_killer'
+
+Sidekiq.configure_server do |config|
+  config.server_middleware do |chain|
+    chain.add Sidekiq::WorkerKiller, max_rss: 480, skip_shutdown_if: ->(worker, job, queue) do
+      worker.to_s == 'LongWorker'
+    end
+  end
+end
+```
+
+## Development
+
+Pull Requests are very welcome!
+
+There are tasks that may help you along the way in a makefile:
+
+```bash
+make console # Loads the whole stack in an IRB session.
+make test # Run tests.
+make lint # Run rubocop linter.
+```
+Please make sure that you have tested your code carefully before opening a PR, and make sure as well that you have no style issues.
 
 ## Authors
 

data/lib/sidekiq/worker_killer.rb

@@ -1,87 +1,131 @@
 require "get_process_mem"
 require "sidekiq"
 require "sidekiq/util"
+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
+
+  MUTEX = Mutex.new
+
+  # @param [Hash] options
+  # @option options [Integer] max_rss
+  #   Max RSS in MB. Above this, shutdown will be triggered.
+  #   (default: `0` (disabled))
+  # @option options [Integer] grace_time
+  #   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 is specified, will wait forever. (default: `900`)
+  # @option options [Integer] shutdown_wait
+  #   when the grace time expires, still running jobs get 30 seconds to
+  #   stop. After that, kill signal is triggered. (default: `30`)
+  # @option options [String] kill_signal
+  #   Signal to use to kill Sidekiq process if it doesn't stop.
+  #   (default: `"SIGKILL"`)
+  # @option options [Boolean] gc
+  #   Try to run garbage collection before Sidekiq process stops in case
+  #   of exceeded max_rss. (default: `true`)
+  # @option options [Proc] skip_shutdown_if
+  #   Executes a block of code after max_rss exceeds but before requesting
+  #   shutdown. (default: `proc {false}`)
+  def initialize(options = {})
+    @max_rss         = options.fetch(:max_rss, 0)
+    @grace_time      = options.fetch(:grace_time, 15 * 60)
+    @shutdown_wait   = options.fetch(:shutdown_wait, 30)
+    @kill_signal     = options.fetch(:kill_signal, "SIGKILL")
+    @gc              = options.fetch(:gc, true)
+    @skip_shutdown   = options.fetch(:skip_shutdown_if, proc { false })
+  end
+
+  # @param [String, Class] worker_class
+  #   the string or class of the worker class being enqueued
+  # @param [Hash] job
+  #   the full job payload
+  #   @see https://github.com/mperham/sidekiq/wiki/Job-Format
+  # @param [String] queue
+  #   the name of the queue the job was pulled from
+  # @yield the next middleware in the chain or the enqueuing of the job
+  def call(worker, job, queue)
+    yield
+    # Skip if the max RSS is not exceeded
+    return unless @max_rss > 0
+    return unless current_rss > @max_rss
+    GC.start(full_mark: true, immediate_sweep: true) if @gc
+    return unless current_rss > @max_rss
+    if skip_shutdown?(worker, job, queue)
+      warn "current RSS #{current_rss} exceeds maximum RSS #{@max_rss}, " \
+           "however shutdown will be ignored"
+      return
+    end
 
-module Sidekiq
-  # Sidekiq server middleware. Kill worker when the RSS memory exceeds limit
-  # after a given grace time.
-  class WorkerKiller
-    include Sidekiq::Util
+    warn "current RSS #{current_rss} of #{identity} exceeds " \
+         "maximum RSS #{@max_rss}"
+    request_shutdown
+  end
 
-    MUTEX = Mutex.new
+  private
 
-    def initialize(options = {})
-      @max_rss         = (options[:max_rss]         || 0)
-      @grace_time      = (options[:grace_time]      || 15 * 60)
-      @shutdown_wait   = (options[:shutdown_wait]   || 30)
-      @kill_signal     = (options[:kill_signal]     || "SIGKILL")
-    end
+  def skip_shutdown?(worker, job, queue)
+    @skip_shutdown.respond_to?(:call) && @skip_shutdown.call(worker, job, queue)
+  end
 
-    def call(_worker, _job, _queue)
-      yield
-      # Skip if the max RSS is not exceeded
-      return unless @max_rss > 0 && current_rss > @max_rss
-      # Launch the shutdown process
-      warn "current RSS #{current_rss} of #{identity} exceeds " \
-           "maximum RSS #{@max_rss}"
-      request_shutdown
+  def request_shutdown
+    # In another thread to allow undelying job to finish
+    Thread.new do
+      # Only if another thread is not already
+      # shutting down the Sidekiq process
+      shutdown if MUTEX.try_lock
     end
+  end
 
-    private
+  def shutdown
+    warn "sending quiet to #{identity}"
+    sidekiq_process.quiet!
 
-    def request_shutdown
-      # In another thread to allow undelying job to finish
-      Thread.new do
-        # Only if another thread is not already
-        # shutting down the Sidekiq process
-        shutdown if MUTEX.try_lock
-      end
-    end
+    sleep(5) # gives Sidekiq API 5 seconds to update ProcessSet
 
-    def shutdown
-      warn "sending #{quiet_signal} to #{identity}"
-      signal(quiet_signal, pid)
+    warn "shutting down #{identity} in #{@grace_time} seconds"
+    wait_job_finish_in_grace_time
 
-      warn "shutting down #{identity} in #{@grace_time} seconds"
-      sleep(@grace_time)
+    warn "stopping #{identity}"
+    sidekiq_process.stop!
 
-      warn "sending SIGTERM to #{identity}"
-      signal("SIGTERM", pid)
+    warn "waiting #{@shutdown_wait} seconds before sending " \
+          "#{@kill_signal} to #{identity}"
+    sleep(@shutdown_wait)
 
-      warn "waiting #{@shutdown_wait} seconds before sending " \
-            "#{@kill_signal} to #{identity}"
-      sleep(@kill_signal)
+    warn "sending #{@kill_signal} to #{identity}"
+    ::Process.kill(@kill_signal, ::Process.pid)
+  end
 
-      warn "sending #{@kill_signal} to #{identity}"
-      signal(@kill_signal, pid)
-    end
+  def wait_job_finish_in_grace_time
+    start = Time.now
+    sleep(1) until grace_time_exceeded?(start) || jobs_finished?
+  end
 
-    def current_rss
-      ::GetProcessMem.new.mb
-    end
+  def grace_time_exceeded?(start)
+    return false if @grace_time == Float::INFINITY
 
-    def signal(signal, pid)
-      ::Process.kill(signal, pid)
-    end
+    start + @grace_time < Time.now
+  end
 
-    def pid
-      ::Process.pid
-    end
+  def jobs_finished?
+    sidekiq_process.stopping? && sidekiq_process["busy"] == 0
+  end
 
-    def identity
-      "#{hostname}:#{pid}"
-    end
+  def current_rss
+    ::GetProcessMem.new.mb
+  end
 
-    def quiet_signal
-      if Gem::Version.new(Sidekiq::VERSION) >= Gem::Version.new("5.0")
-        "TSTP"
-      else
-        "USR1"
-      end
-    end
+  def sidekiq_process
+    Sidekiq::ProcessSet.new.find do |process|
+      process["identity"] == identity
+    end || raise("No sidekiq worker with identity #{identity} found")
+  end
 
-    def warn(msg)
-      Sidekiq.logger.warn(msg)
-    end
+  def warn(msg)
+    Sidekiq.logger.warn(msg)
   end
 end

data/lib/sidekiq/worker_killer/version.rb

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

data/spec/sidekiq/worker_killer_spec.rb

@@ -1,29 +1,93 @@
 require "spec_helper"
 
 describe Sidekiq::WorkerKiller do
+  let(:sidekiq_process_set) { instance_double(Sidekiq::ProcessSet) }
+  let(:sidekiq_process) { instance_double(Sidekiq::Process) }
 
   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") }
+
     it "should yield" do
       expect { |b|
         subject.call(worker, job, queue, &b)
       }.to yield_with_no_args
     end
+
     context "when current rss is over max rss" do
       subject{ described_class.new(max_rss: 2) }
+
       before do
         allow(subject).to receive(:current_rss).and_return(3)
       end
+
       it "should request shutdown" do
         expect(subject).to receive(:request_shutdown)
         subject.call(worker, job, queue){}
       end
+
+      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 "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 "but max rss is 0" do
         subject{ described_class.new(max_rss: 0) }
         it "should not request shutdown" do
@@ -35,29 +99,64 @@ describe Sidekiq::WorkerKiller do
   end
 
   describe "#request_shutdown" 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)
+    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
-  end
 
-  describe "#quiet_signal" do
-    it "should give TSTP if Sidekiq version is > 5.0" do
-      stub_const("Sidekiq::VERSION", "5.0")
-      expect(subject.send :quiet_signal).to eq "TSTP"
-      stub_const("Sidekiq::VERSION", "5.2.1")
-      expect(subject.send :quiet_signal).to eq "TSTP"
+    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
+
+        allow(Process).to receive(:kill)
+        allow(Process).to receive(:pid).and_return(99)
+
+        subject.send(:request_shutdown).join
+
+        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
+      end
     end
-    it "should give USR1 if Sidekiq version is < 5.0" do
-      stub_const("Sidekiq::VERSION", "3.0")
-      expect(subject.send :quiet_signal).to eq "USR1"
-      stub_const("Sidekiq::VERSION", "4.6.7")
-      expect(subject.send :quiet_signal).to eq "USR1"
+
+    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
+      end
     end
   end
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: sidekiq-worker-killer
 version: !ruby/object:Gem::Version
-  version: 0.2.0
+  version: 1.0.0
 platform: ruby
 authors:
 - Cyrille Courtiere
 autorequire: 
 bindir: bin
 cert_chain: []
-date: 2018-04-18 00:00:00.000000000 Z
+date: 2020-05-28 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: get_process_mem
@@ -30,42 +30,42 @@ dependencies:
     requirements:
     - - ">="
       - !ruby/object:Gem::Version
-        version: '3'
+        version: '5'
   type: :runtime
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - ">="
       - !ruby/object:Gem::Version
-        version: '3'
+        version: '5'
 - !ruby/object:Gem::Dependency
-  name: rubocop
+  name: rspec
   requirement: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: 0.49.1
+        version: '3.5'
   type: :development
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: 0.49.1
+        version: '3.5'
 - !ruby/object:Gem::Dependency
-  name: rspec
+  name: rubocop
   requirement: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '3.5'
+        version: 0.49.1
   type: :development
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '3.5'
+        version: 0.49.1
 description: 
 email:
 - dev@klaxit.com
@@ -96,8 +96,7 @@ required_rubygems_version: !ruby/object:Gem::Requirement
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubyforge_project: 
-rubygems_version: 2.6.14
+rubygems_version: 3.1.0.pre3
 signing_key: 
 specification_version: 4
 summary: Sidekiq worker killer