Merge pull request #3 from iloveitaly/heroku-will-terminate

Addding heroku_will_terminate?

Gemfile

@@ -1,2 +1,5 @@
 source 'https://rubygems.org'
 gemspec
+
+gem 'pry'
+gem 'pry-byebug'

README.md

@@ -15,7 +15,15 @@ Add this line to your application's Gemfile:
 gem 'resque-heroku-signals'
 ```
 
-Since this gem monkeypatches the Heroku worker the `gemspec` is locked to a `x.x.x` version of Resque. Issue a PR if this is not compatible with the version of resque you are using. 
+Since this gem monkeypatches the Heroku worker the `gemspec` is locked to a `x.x.x` version of Resque. Issue a PR if this is not compatible with the version of resque you are using.
+
+## Determining When a Process Will Shutdown
+
+Heroku sends a `TERM` signal to a process before hard killing it. If your job communicates with slow external APIs, you may want to make sure you have enough time to receive and handle the response from the external system before executing the API requests.
+
+Ideally, using an idempotency key with each external API request is the best way to ensure that a given API request only runs. However, depending on your application logic this may not be practical and knowing if a process will be terminated in less than 30s by Heroku is a useful tool.
+
+Use `Resque.heroku_will_terminate?` to determine if Heroku will terminate your process within 30s. 
 
 ## Example Procfile
 

lib/resque-heroku-signals.rb

@@ -1,9 +1,19 @@
 require 'resque'
 
+$HEROKU_WILL_TERMINATE_RESQUE = false
+
+Resque.class_eval do
+  def self.heroku_will_terminate?
+    !!$HEROKU_WILL_TERMINATE_RESQUE
+  end
+end
+
 # https://github.com/resque/resque/issues/1559#issuecomment-310908574
 Resque::Worker.class_eval do
   def unregister_signal_handlers
     trap('TERM') do
+      $HEROKU_WILL_TERMINATE_RESQUE = true
+
       trap('TERM') do
         log_with_severity :info, "[resque-heroku] received second term signal, throwing term exception"
 
@@ -15,7 +25,6 @@ def unregister_signal_handlers
       end
 
       log_with_severity :info, "[resque-heroku] received first term signal from heroku, ignoring"
-
     end
 
     trap('INT', 'DEFAULT')

spec/resque/heroku_spec.rb

@@ -8,14 +8,48 @@ def self.perform(uuid)
   end
 end
 
+def LongCleanupJob
+  def self.perform(uuid)
+    sleep 2
+
+    if Resque.heroku_will_terminate?
+      FileUtils.touch(uuid)
+    end
+  end
+end
+
+def get_job_pid(worker)
+  pid = nil
+  wait condition: -> () {
+    pid = worker.instance_variable_get("@child")
+    pid.nil?
+  }
+  pid
+end
+
+def wait(timeout: 5000, condition:)
+  t0 = Time.now
+  while condition.call && ((Time.now - t0) * 1000 < timeout)
+    sleep 0.1
+  end
+end
+
 RSpec.describe 'resque-heroku-signals' do
 
   context "Signal handling" do
     before do
       @uuid = SecureRandom.uuid
       ENV["TERM_CHILD"] = "1"
 
+      # match production timeouts, this ensures that there are not any edge
+      # cases in the underlying resque logic that may conflict with our patches
+      ENV["RESQUE_PRE_SHUTDOWN_TIMEOUT"] = "20"
+      ENV["RESQUE_TERM_TIMEOUT"] = "8"
+
       @worker = Resque::Worker.new(:jobs)
+
+      # by default, resque doesn't log anything
+      @worker.very_verbose = true
     end
 
     after do
@@ -38,35 +72,46 @@ def self.perform(uuid)
     end
 
     it "ignores the first TERM signal but raises an exception on the second signal" do
-      thread = Thread.new do
-        @worker.work(1)
-      end
+      thread = Thread.new { @worker.work(0.1) { puts "hello" } }
+
       Resque::Job.create(:jobs, DummyJob, @uuid)
       pid = get_job_pid(@worker)
-      Process.kill("TERM", pid)
-      sleep 0 # It seems like the second signal isn't received without this
-      Process.kill("TERM", pid)
+
+      expect(pid).to_not be_nil
+
+      Process.kill(:TERM, pid)
+
+
+      # It seems like the second signal isn't received without this
+      sleep 0.1
+      Process.kill(:TERM, pid)
 
       @worker.shutdown
       thread.join
 
+      # Implied assertion here is if the file does not exist, then an exception was raised
       expect(File.exist?(@uuid)).to eq(false)
     end
-  end
-end
 
-def get_job_pid(worker)
-  pid = nil
-  wait condition: -> () {
-    pid = worker.instance_variable_get("@child")
-    pid.nil?
-  }
-  pid
-end
+    it 'provides a flag indicating if heroku will soon terminate the worker' do
+      expect(Resque.heroku_will_terminate?).to be false
 
-def wait(timeout: 5000, condition:)
-  t0 = Time.now
-  while condition.call && ((Time.now - t0) * 1000 < timeout)
-    sleep 0.1
+      # `work` must be run in a separate thread, a new process is only created
+      # after a job is picked up. The current thread is used to run the job loop
+      thread = Thread.new { @worker.work(0.1) { puts "hello" } }
+
+      Resque::Job.create(:jobs, LongCleanupJob, @uuid)
+      pid = get_job_pid(@worker)
+
+      expect(pid).to_not be_nil
+
+      Process.kill(:TERM, pid)
+
+      @worker.shutdown
+      thread.join
+
+      # if worker will terminate, then the file is written
+      expect(File.exist?(@uuid)).to eq(true)
+    end
   end
 end