Discard bucket when rate limit settings change

Author: marksmith

lib/aikido/zen/rate_limiter.rb

@@ -16,12 +16,7 @@ def initialize(
     )
       @config = config
       @settings = settings
-      @buckets = Hash.new { |store, route|
-        synchronize {
-          settings = settings_for(route)
-          store[route] = Bucket.new(ttl: settings.period, max_size: settings.max_requests)
-        }
-      }
+      @buckets = {}
     end
 
     # Calculate based on the configuration whether a request will be
@@ -30,24 +25,30 @@ def initialize(
     # @param request [Aikido::Zen::Request]
     # @return [Aikido::Zen::RateLimiter::Result, nil]
     def calculate_rate_limits(request)
-      route, enabled = resolve_route_enabled(request)
-      return nil unless enabled
+      route, settings = @settings.endpoints.match(request.route)
 
-      bucket = @buckets[route]
-      key = @config.rate_limiting_discriminator.call(request)
-      bucket.increment(key)
-    end
+      rate_limit_settings = settings&.rate_limiting
+
+      return nil unless rate_limit_settings&.enabled?
+
+      bucket = nil
 
-    private
+      synchronize do
+        bucket = @buckets[route]
 
-    def resolve_route_enabled(request)
-      @settings.endpoints.match(request.route) do |route, settings|
-        [route, settings.rate_limiting.enabled?]
+        if bucket.nil? || bucket.settings_changed?(rate_limit_settings)
+          bucket = Bucket.new(
+            ttl: rate_limit_settings.period,
+            max_size: rate_limit_settings.max_requests,
+            settings: rate_limit_settings
+          )
+
+          @buckets[route] = bucket
+        end
       end
-    end
 
-    def settings_for(route)
-      @settings.endpoints[route].rate_limiting
+      key = @config.rate_limiting_discriminator.call(request)
+      bucket.increment(key)
     end
   end
 end

lib/aikido/zen/rate_limiter/bucket.rb

@@ -29,14 +29,21 @@ class RateLimiter::Bucket
     # @!visibility private
     #
     # Use the monotonic clock to ensure time differences are consistent
-    # and not affected by timezones.or daylight savings changes.
+    # and not affected by timezones or daylight savings changes.
     DEFAULT_CLOCK = -> { Process.clock_gettime(Process::CLOCK_MONOTONIC).round }
 
-    def initialize(ttl:, max_size:, clock: DEFAULT_CLOCK)
+    # @param ttl [Integer] the time to live in seconds
+    # @param max_size [Integer] the number of requests before throttling
+    # @param clock [#call() => Integer] a callable that returns the current
+    #   monotonic time in seconds
+    # @param settings [Aikido::Zen::RuntimeSettings::RateLimitSettings, nil]
+    #   the settings that might change
+    def initialize(ttl:, max_size:, clock: DEFAULT_CLOCK, settings: nil)
       @ttl = ttl
       @max_size = max_size
       @data = Hash.new { |h, k| h[k] = [] }
       @clock = clock
+      @settings = settings
     end
 
     # Increments the key if the number of entries within the current TTL window
@@ -67,10 +74,21 @@ def increment(key)
       end
     end
 
+    # Checks whether the settings provided at initialization have changed.
+    #
+    # @param settings [Aikido::Zen::RuntimeSettings::RateLimitSettings]
+    # @return [Boolean] true if settings were provided at initialization and
+    #   the settings have changed
+    def settings_changed?(settings)
+      !@settings.nil? && @settings != settings
+    end
+
     private
 
     def evict(key, at: @clock.call)
-      synchronize { @data[key].delete_if { |time| time < (at - @ttl) } }
+      synchronize do
+        @data[key].delete_if { |time| time < (at - @ttl) }
+      end
     end
   end
 end

lib/aikido/zen/runtime_settings/rate_limit_settings.rb

@@ -28,6 +28,9 @@ def self.disabled
       new(enabled: false)
     end
 
+    # @return [Boolean]
+    attr_reader :enabled
+
     # @return [Integer] the fixed window to bucket requests in, in seconds.
     attr_reader :period
 
@@ -43,5 +46,13 @@ def initialize(enabled: false, max_requests: 1000, period: 60)
     def enabled?
       @enabled
     end
+
+    def ==(other)
+      other.is_a?(self.class) &&
+        other.enabled == enabled &&
+        other.period == period &&
+        other.max_requests == max_requests
+    end
+    alias_method :eql?, :==
   end
 end