Merge pull request #51 from creditkudos/master

Adds `on_calculation` callback.

README.md

@@ -130,6 +130,12 @@ PumaWorkerKiller: Rolling Restart. 5 workers consuming total: 650mb mb. Sending
 
 However you may want to collect more data, such as sending an event to an error collection service like rollbar or airbrake. The `pre_term` lambda gets called before any worker is killed by PWK for any reason.
 
+### on_calculation
+
+`config.on_calculation` will be called every time Puma Worker Killer calculates memory usage (`config.frequency`). This may be useful for monitoring your total puma application memory usage, which can be contrasted with other application monitoring solutions.
+
+This callback lambda is given a single value for the amount of memory used.
+
 ## Attention
 
 If you start puma as a daemon, to add puma worker killer config into puma config file, rather than into initializers:

lib/puma_worker_killer.rb

@@ -3,20 +3,21 @@
 module PumaWorkerKiller
   extend self
 
-  attr_accessor :ram, :frequency, :percent_usage, :rolling_restart_frequency, :reaper_status_logs, :pre_term
+  attr_accessor :ram, :frequency, :percent_usage, :rolling_restart_frequency, :reaper_status_logs, :pre_term, :on_calculation
   self.ram           = 512  # mb
   self.frequency     = 10   # seconds
   self.percent_usage = 0.99 # percent of RAM to use
   self.rolling_restart_frequency = 6 * 3600 # 6 hours in seconds
   self.reaper_status_logs = true
-  self.pre_term = lambda { |_| } # nop
+  self.pre_term = nil
+  self.on_calculation = nil
 
   def config
     yield self
   end
 
-  def reaper(ram = self.ram, percent = self.percent_usage, reaper_status_logs = self.reaper_status_logs, pre_term = self.pre_term)
-    Reaper.new(ram * percent_usage, nil, reaper_status_logs, pre_term)
+  def reaper(ram = self.ram, percent = self.percent_usage, reaper_status_logs = self.reaper_status_logs, pre_term = self.pre_term, on_calculation = self.on_calculation)
+    Reaper.new(ram * percent_usage, nil, reaper_status_logs, pre_term, on_calculation)
   end
 
   def start(frequency = self.frequency, reaper = self.reaper)

lib/puma_worker_killer/reaper.rb

@@ -1,10 +1,11 @@
 module PumaWorkerKiller
   class Reaper
-    def initialize(max_ram, master = nil, reaper_status_logs = true, pre_term)
+    def initialize(max_ram, master = nil, reaper_status_logs = true, pre_term = nil, on_calculation = nil)
       @cluster = PumaWorkerKiller::PumaMemory.new(master)
       @max_ram = max_ram
       @reaper_status_logs = reaper_status_logs
       @pre_term = pre_term
+      @on_calculation = on_calculation
     end
 
     # used for tes
@@ -14,7 +15,10 @@ def get_total_memory
 
     def reap
       return false if @cluster.workers_stopped?
-      if (total = get_total_memory) > @max_ram
+      total = get_total_memory
+      @on_calculation.call(total) unless @on_calculation.nil?
+
+      if total > @max_ram
         @cluster.master.log "PumaWorkerKiller: Out of memory. #{@cluster.workers.count} workers consuming total: #{total} mb out of max: #{@max_ram} mb. Sending TERM to pid #{@cluster.largest_worker.pid} consuming #{@cluster.largest_worker_memory} mb."
 
         # Fetch the largest_worker so that both `@pre_term` and `term_worker` are called with the same worker
@@ -25,7 +29,7 @@ def reap
         #   A new request comes in, Worker B takes it, and consumes 101 mb memory
         #   term_largest_worker (previously here) gets called and terms Worker B (thus not passing the about-to-be-terminated worker to `@pre_term`)
         largest_worker = @cluster.largest_worker
-        @pre_term.call(largest_worker)
+        @pre_term.call(largest_worker) unless @pre_term.nil?
         @cluster.term_worker(largest_worker)
 
       elsif @reaper_status_logs

test/fixtures/on_calculation.ru

@@ -0,0 +1,8 @@
+load File.expand_path("../fixture_helper.rb", __FILE__)
+
+PumaWorkerKiller.config do |config|
+  config.on_calculation = lambda { |usage| puts("Current memory footprint: #{usage} mb") }
+end
+PumaWorkerKiller.start
+
+run HelloWorldApp

test/puma_worker_killer_test.rb

@@ -45,6 +45,18 @@ def test_pre_term
     end
   end
 
+  def test_on_calculation
+    file     = fixture_path.join("on_calculation.ru")
+    port     = 0
+    command  = "bundle exec puma #{ file } -t 1:1 -w 2 --preload --debug -p #{ port }"
+    options  = { wait_for: "booted", timeout: 5, env: { "PUMA_FREQUENCY" => 1, 'PUMA_RAM' => 1} }
+
+    WaitForIt.new(command, options) do |spawn|
+      assert_contains(spawn, "Out of memory")
+      assert_contains(spawn, "Current memory footprint:") # defined in on_calculate.ru
+    end
+  end
+
   def assert_contains(spawn, string)
     assert spawn.wait(string), "Expected logs to contain '#{string}' but it did not, contents: #{ spawn.log.read }"
   end