Merge pull request #38142 from code-dot-org/cdo-throttle

Add Cdo::Throttle module

lib/cdo/throttle.rb

@@ -0,0 +1,45 @@
+require 'cdo/shared_cache'
+require 'dynamic_config/dcdo'
+
+module Cdo
+  module Throttle
+    CACHE_PREFIX = "cdo_throttle/".freeze
+
+    # @param [String] id - Unique identifier to throttle on.
+    # @param [Integer] limit - Number of requests allowed over period.
+    # @param [Integer] period - Period of time in seconds.
+    # @param [Integer] throttle_for - How long id should stay throttled in seconds. Optional.
+    # Defaults to Cdo::Throttle.throttle_time.
+    # @returns [Boolean] Whether or not the request should be throttled.
+    def self.throttle(id, limit, period, throttle_for = throttle_time)
+      full_key = CACHE_PREFIX + id.to_s
+      value = CDO.shared_cache.read(full_key) || empty_value
+      now = Time.now.utc
+      value[:request_timestamps] << now
+
+      if value[:throttled_until]&.future?
+        should_throttle = true
+      else
+        value[:throttled_until] = nil
+        earliest = now - period
+        value[:request_timestamps].select! {|timestamp| timestamp >= earliest}
+        should_throttle = value[:request_timestamps].size > limit
+        value[:throttled_until] = now + throttle_for if should_throttle
+      end
+
+      CDO.shared_cache.write(full_key, value)
+      should_throttle
+    end
+
+    def self.empty_value
+      {
+        throttled_until: nil,
+        request_timestamps: []
+      }
+    end
+
+    def self.throttle_time
+      DCDO.get('throttle_time_default', 60)
+    end
+  end
+end

lib/test/cdo/test_throttle.rb

@@ -0,0 +1,37 @@
+require_relative '../test_helper'
+require 'cdo/throttle'
+require 'timecop'
+
+class ThrottleTest < Minitest::Test
+  def teardown
+    CDO.shared_cache.clear
+  end
+
+  def test_throttle_with_limit_1
+    Timecop.freeze
+    refute Cdo::Throttle.throttle("my_key", 1, 2) # 1/1 reqs per 2s - not throttled
+    Timecop.travel(Time.now.utc + 1)
+    assert Cdo::Throttle.throttle("my_key", 1, 2) # 2/1 reqs per 2s - throttled
+    Timecop.travel(Time.now.utc + Cdo::Throttle.throttle_time - 1)
+    assert Cdo::Throttle.throttle("my_key", 1, 2) # still throttled
+    Timecop.travel(Time.now.utc + Cdo::Throttle.throttle_time)
+    refute Cdo::Throttle.throttle("my_key", 1, 2) # 1/1 reqs per 2s after waiting - not throttled anymore
+    Timecop.travel(Time.now.utc + 1)
+    assert Cdo::Throttle.throttle("my_key", 1, 2) # 2/1 reqs per 2s - throttled again
+  end
+
+  def test_throttle_with_limit_greater_than_1
+    Timecop.freeze
+    refute Cdo::Throttle.throttle("my_key", 2, 2) # 1/2 reqs per 2s - not throttled
+    Timecop.travel(Time.now.utc + 1)
+    refute Cdo::Throttle.throttle("my_key", 2, 2) # 2/2 reqs per 2s - not throttled
+    Timecop.travel(Time.now.utc + 0.5)
+    assert Cdo::Throttle.throttle("my_key", 2, 2) # 3/2 reqs per 2s - throttled
+    Timecop.travel(Time.now.utc + Cdo::Throttle.throttle_time)
+    refute Cdo::Throttle.throttle("my_key", 2, 2) # 1/2 reqs per 2s after waiting - not throttled anymore
+    Timecop.travel(Time.now.utc + 1)
+    refute Cdo::Throttle.throttle("my_key", 2, 2) # 2/2 reqs per 2s - not throttled
+    Timecop.travel(Time.now.utc + 0.5)
+    assert Cdo::Throttle.throttle("my_key", 2, 2) # 3/2 reqs per 2s - throttled again
+  end
+end