Add support for Rails 5, more tests, proc key arguments.

nickelser · nickelser · commit 1fcc8340155a · 2016-06-30T18:00:31.000-07:00
diff --git a/lib/active_job/traffic_control.rb b/lib/active_job/traffic_control.rb
@@ -36,7 +36,7 @@ def logger
       def client
         @client ||= begin
           logger.error "defaulting to Redis as the lock client; please set "\
-                       " `ActiveJob::TrafficControl.client` to a Redis or Memcached client,"
+                       " `ActiveJob::TrafficControl.client` to a Redis or Memcached client."
           @client_klass = Suo::Client::Redis
           Redis.new(url: ENV["REDIS_URL"])
         end
diff --git a/lib/active_job/traffic_control/base.rb b/lib/active_job/traffic_control/base.rb
@@ -18,6 +18,22 @@ def cleaned_name
         def cache_client
           ActiveJob::TrafficControl.cache_client
         end
+
+        def lock_key(prefix, job, config_attr)
+          if config_attr
+            if config_attr[:key].respond_to?(:call)
+              "traffic_control:#{prefix}:#{config_attr[:key].call(job)}"
+            else
+              @static_job_key ||= begin
+                if config_attr[:key].present?
+                  "traffic_control:#{prefix}:#{config_attr[:key]}"
+                else
+                  "traffic_control:#{prefix}:#{cleaned_name}"
+                end
+              end
+            end
+          end
+        end
       end
 
       # convenience methods
@@ -28,12 +44,12 @@ def cleaned_name
       def reenqueue(range, reason)
         later_delay = rand(range).seconds
         retry_job(wait: later_delay)
-        logger.error "Re-enqueing #{self.class.name} to run in #{later_delay}s due to #{reason}"
+        logger.info "Re-enqueing #{self.class.name} to run in #{later_delay}s due to #{reason}"
         ActiveSupport::Notifications.instrument "re_enqueue.active_job", job: self, reason: reason
       end
 
       def drop(reason)
-        logger.error "Dropping #{self.class.name} due to #{reason}"
+        logger.info "Dropping #{self.class.name} due to #{reason}"
         ActiveSupport::Notifications.instrument "dropped.active_job", job: self, reason: reason
       end
 
diff --git a/lib/active_job/traffic_control/concurrency.rb b/lib/active_job/traffic_control/concurrency.rb
@@ -5,11 +5,14 @@ module TrafficControl
     module Concurrency
       extend ::ActiveSupport::Concern
 
+      CONCURRENCY_REENQUEUE_DELAY = ENV["RACK_ENV"] == "test" ? 1...5 : 30...(60 * 5)
+
       class_methods do
         attr_accessor :job_concurrency
 
         def concurrency(threshold, drop: true, key: nil, wait_timeout: 0.1, stale_timeout: 60 * 10)
           raise ArgumentError, "Concurrent jobs needs to be an integer > 0" if threshold.to_i < 1
+
           @job_concurrency = {
             threshold: threshold.to_i,
             drop: drop,
@@ -19,31 +22,23 @@ def concurrency(threshold, drop: true, key: nil, wait_timeout: 0.1, stale_timeou
           }
         end
 
-        def concurrency_key
-          if job_concurrency
-            @concurrency_key ||= begin
-              if job_concurrency[:key].present?
-                job_concurrency[:key]
-              else
-                "traffic_control:concurrency:#{cleaned_name}"
-              end
-            end
-          end
+        def concurrency_lock_key(job)
+          lock_key("concurrency", job, job_concurrency)
         end
       end
 
       included do
         include ActiveJob::TrafficControl::Base
 
-        around_perform do |_, block|
+        around_perform do |job, block|
           if self.class.job_concurrency.present?
             lock_options = {
               resources: self.class.job_concurrency[:threshold],
               acquisition_lock: self.class.job_concurrency[:wait_timeout],
               stale_lock_expiration: self.class.job_concurrency[:stale_timeout]
             }
 
-            with_lock_client(self.class.concurrency_key, lock_options) do |client|
+            with_lock_client(self.class.concurrency_lock_key(job), lock_options) do |client|
               locked = client.lock do
                 block.call
                 true
diff --git a/lib/active_job/traffic_control/throttle.rb b/lib/active_job/traffic_control/throttle.rb
@@ -10,33 +10,31 @@ module Throttle
 
         def throttle(threshold:, period:, drop: false, key: nil)
           raise ArgumentError, "Threshold needs to be an integer > 0" if threshold.to_i < 1
-          @job_throttling = {threshold: threshold, period: period, drop: drop, key: key}
+
+          @job_throttling = {
+            threshold: threshold,
+            period: period,
+            drop: drop,
+            key: key
+          }
         end
 
-        def throttling_key
-          if job_throttling
-            @throttling_key ||= begin
-              if job_throttling[:key].present?
-                job_throttling[:key]
-              else
-                "traffic_control:throttling:#{cleaned_name}"
-              end
-            end
-          end
+        def throttling_lock_key(job)
+          lock_key("throttle", job, job_throttling)
         end
       end
 
       included do
         include ActiveJob::TrafficControl::Base
 
-        around_perform do |_, block|
+        around_perform do |job, block|
           if self.class.job_throttling.present?
             lock_options = {
               resources: self.class.job_throttling[:threshold],
               stale_lock_expiration: self.class.job_throttling[:period]
             }
 
-            with_lock_client(self.class.throttling_key, lock_options) do |client|
+            with_lock_client(self.class.throttling_lock_key(job), lock_options) do |client|
               token = client.lock
 
               if token
diff --git a/test/active_job/traffic_control_test.rb b/test/active_job/traffic_control_test.rb
@@ -19,6 +19,40 @@ class ThrottleTestJob < ActiveJob::Base
     throttle threshold: 2, period: 1.second, drop: true
 
     def perform
+      sleep 0.5
+      $count += 1
+    end
+  end
+
+  class ThrottleWithKeyTestJob < ActiveJob::Base
+    include ActiveJob::TrafficControl::Throttle
+
+    throttle threshold: 2, period: 1.second, drop: true, key: "throttle_test_key"
+
+    def perform
+      sleep 0.5
+      $count += 1
+    end
+  end
+
+  class ThrottleWithProcKeyTestJob < ActiveJob::Base
+    include ActiveJob::TrafficControl::Throttle
+
+    throttle threshold: 2, period: 1.second, drop: true, key: -> (_) { "throttle_proc_job_name" }
+
+    def perform
+      sleep 0.5
+      $count += 1
+    end
+  end
+
+  class ThrottleNotDroppedTestJob < ActiveJob::Base
+    include ActiveJob::TrafficControl::Throttle
+
+    throttle threshold: 2, period: 1.second, drop: false
+
+    def perform
+      sleep 0.5
       $count += 1
     end
   end
@@ -34,6 +68,39 @@ def perform
     end
   end
 
+  class ConcurrencyNotDroppedTestJob < ActiveJob::Base
+    include ActiveJob::TrafficControl::Concurrency
+
+    concurrency 1, drop: false
+
+    def perform
+      sleep 0.5
+      $count += 1
+    end
+  end
+
+  class ConcurrencyWithKeyTestJob < ActiveJob::Base
+    include ActiveJob::TrafficControl::Concurrency
+
+    concurrency 1, drop: true, key: "concurrency_test_key"
+
+    def perform
+      sleep 0.5
+      $count += 1
+    end
+  end
+
+  class ConcurrencyWithProcKeyTestJob < ActiveJob::Base
+    include ActiveJob::TrafficControl::Concurrency
+
+    concurrency 1, drop: true, key: -> (_) { "concurrency_proc_job_name" }
+
+    def perform
+      sleep 0.5
+      $count += 1
+    end
+  end
+
   class InheritedConcurrencyJob < ConcurrencyTestJob
     def perform
       $count += 1
@@ -70,33 +137,80 @@ def test_disable
     assert_equal 2, $count
   end
 
-  def test_throttle
-    t1 = Thread.new { ThrottleTestJob.perform_now }
-    t2 = Thread.new { ThrottleTestJob.perform_now }
-    t3 = Thread.new { ThrottleTestJob.perform_now }
+  def throttle_helper(klass)
+    t1 = Thread.new { klass.perform_now }
+    t2 = Thread.new { klass.perform_now }
+    t3 = Thread.new { klass.perform_now }
     [t1, t2, t3].map(&:join)
+    sleep 0.5
     assert_equal 2, $count
     sleep 1
-    ThrottleTestJob.perform_now
+    klass.perform_now
+    assert_equal 3, $count
+  end
+
+  def test_throttle
+    throttle_helper(ThrottleTestJob)
+  end
+
+  def test_throttle_with_key
+    throttle_helper(ThrottleWithKeyTestJob)
+  end
+
+  def test_throttle_with_proc_key
+    throttle_helper(ThrottleWithProcKeyTestJob)
+  end
+
+  def test_throttle_not_dropped
+    return unless ActiveJob::Base.queue_adapter == :async
+
+    t1 = Thread.new { ThrottleNotDroppedTestJob.perform_now }
+    t2 = Thread.new { ThrottleNotDroppedTestJob.perform_now }
+    t3 = Thread.new { ThrottleNotDroppedTestJob.perform_now }
+    [t1, t2, t3].map(&:join)
+    sleep 0.5
+    assert_equal 2, $count
+    sleep 6
     assert_equal 3, $count
   end
 
+  def concurrency_helper(klass)
+    t1 = Thread.new { klass.perform_now }
+    t2 = Thread.new { klass.perform_now }
+    [t1, t2].map(&:join)
+    sleep 0.5
+    assert_equal 1, $count
+    klass.perform_now
+    assert_equal 2, $count
+  end
+
   def test_concurrency
-    t1 = Thread.new { ConcurrencyTestJob.perform_now }
-    t2 = Thread.new { ConcurrencyTestJob.perform_now }
+    concurrency_helper(ConcurrencyTestJob)
+  end
+
+  def test_concurrency_with_key
+    concurrency_helper(ConcurrencyWithKeyTestJob)
+  end
+
+  def test_concurrency_with_proc_key
+    concurrency_helper(ConcurrencyWithProcKeyTestJob)
+  end
+
+  def test_concurrent_not_dropped
+    return unless ActiveJob::Base.queue_adapter == :async
+
+    t1 = Thread.new { ConcurrencyNotDroppedTestJob.perform_now }
+    t2 = Thread.new { ConcurrencyNotDroppedTestJob.perform_now }
     [t1, t2].map(&:join)
-    sleep 1
     assert_equal 1, $count
-    ConcurrencyTestJob.perform_later
-    sleep 1
+    sleep 6
     assert_equal 2, $count
   end
 
   def test_concurrency_is_not_inherited
-    t1 = Thread.new { InheritedConcurrencyJob.perform_later }
-    t2 = Thread.new { InheritedConcurrencyJob.perform_later }
+    t1 = Thread.new { InheritedConcurrencyJob.perform_now }
+    t2 = Thread.new { InheritedConcurrencyJob.perform_now }
     [t1, t2].map(&:join)
-    sleep 1
     assert_equal 2, $count
   end
 
diff --git a/test/test_helper.rb b/test/test_helper.rb
@@ -1,5 +1,7 @@
 $LOAD_PATH.unshift File.expand_path("../../lib", __FILE__)
 
+ENV["RACK_ENV"] = ENV["RAILS_ENV"] = "test"
+
 if ENV["CODECLIMATE_REPO_TOKEN"]
   require "codeclimate-test-reporter"
   ::SimpleCov.add_filter "helper"
@@ -18,5 +20,10 @@
   l
 end
 
-ActiveJob::Base.queue_adapter = :inline
+begin
+  ActiveJob::Base.queue_adapter = :async # ActiveJob 5
+rescue => _
+  ActiveJob::Base.queue_adapter = :inline
+end
+
 ActiveJob::Base.logger = test_logger
