Merge branch 'rate-limit/rule-extras' into 'master' (b94e1223) · Commits · GitLab.org / ruby / gems / labkit-ruby

lib/labkit/rate_limit.rb

+5 −2

Original line number	Diff line number	Diff line
		@@ -44,9 +44,12 @@ module Labkit
		# @param redis [Object, nil] Redis client; falls back to config.redis
		# @param logger [Logger, nil] logger; falls back to config.logger
		# @param cost [Numeric] amount to add to the counter; see Limiter#check
		# @param rule_context [Hash, nil] per-request context for one-arity
		# callables on rule limit/period; see Limiter#check
		# @return [Result]
		def check(name:, identifier:, rules:, redis: nil, logger: nil, cost: 1)
		Limiter.new(name: name, rules: rules, redis: redis, logger: logger).check(identifier, cost: cost)
		def check(name:, identifier:, rules:, redis: nil, logger: nil, cost: 1, rule_context: nil)
		Limiter.new(name: name, rules: rules, redis: redis, logger: logger)
		.check(identifier, cost: cost, rule_context: rule_context)
		end
		end
		end

lib/labkit/rate_limit/evaluator.rb

+37 −16

Original line number	Diff line number	Diff line
		@@ -65,8 +65,8 @@ module Labkit
		@logger = logger
		end

		def check(identifier, cost: 1)
		check_rules(identifier, cost)
		def check(identifier, cost: 1, rule_context: nil)
		check_rules(identifier, cost, rule_context)
		rescue StandardError => e
		# Intentionally broad: fail-open applies to any unexpected error (network,
		# timeout, OOM) not only Redis protocol errors.
		@@ -78,8 +78,8 @@ module Labkit
		# Read-without-increment counterpart to {#check}. Same matching and Result
		# shape; the underlying Redis counter is not mutated and the TTL is not
		# extended. A missing Redis key is treated as count=0 (matched, not exceeded).
		def peek(identifier)
		peek_rules(identifier)
		def peek(identifier, rule_context: nil)
		peek_rules(identifier, rule_context)
		rescue StandardError => e
		report_error_metrics
		log_error(e, identifier)
		@@ -95,7 +95,7 @@ module Labkit
		# is missing the count_distinct key fail open + log + bump errors_total, and
		# the loop continues to the next rule (the rule is treated as not applicable
		# rather than aborting the whole evaluation).
		def check_rules(identifier, cost)
		def check_rules(identifier, cost, rule_context)
		@rules.each do \|rule\|
		next unless rule_matches?(rule, identifier)

		@@ -105,7 +105,7 @@ module Labkit
		next
		end

		result = evaluate_rule(rule, identifier, cost)
		result = evaluate_rule(rule, identifier, cost, rule_context)
		report_matched_metrics(result)
		return result unless rule.action == :log
		end
		@@ -120,12 +120,12 @@ module Labkit
		# peek does not need the count_distinct identifier key - it reads SCARD on
		# the rule-keyed compound key, which contains the cardinality across all
		# members. So missing-key fail-open does not apply here.
		def peek_rules(identifier)
		def peek_rules(identifier, rule_context)
		@rules.each do \|rule\|
		next if rule.action == :log
		next unless rule_matches?(rule, identifier)

		return peek_rule(rule, identifier)
		return peek_rule(rule, identifier, rule_context)
		end

		Result.new(matched: false, action: :allow)
		@@ -140,10 +140,10 @@ module Labkit
		value.nil? \|\| value.to_s.empty?
		end

		def evaluate_rule(rule, identifier, cost)
		def evaluate_rule(rule, identifier, cost, rule_context)
		redis_key = build_redis_key(rule, identifier)
		resolved_limit = Integer(resolve_value(rule.limit))
		resolved_period = Integer(resolve_value(rule.period))
		resolved_limit = Integer(resolve_value(rule.limit, rule_context))
		resolved_period = Integer(resolve_value(rule.period, rule_context))

		# cost is ignored for count_distinct rules: SADD is binary (a member is
		# either added or not), and the post-add count is SCARD regardless.
		@@ -157,10 +157,10 @@ module Labkit
		build_result(rule, resolved_limit, resolved_period, count, ttl)
		end

		def peek_rule(rule, identifier)
		def peek_rule(rule, identifier, rule_context)
		redis_key = build_redis_key(rule, identifier)
		resolved_limit = Integer(resolve_value(rule.limit))
		resolved_period = Integer(resolve_value(rule.period))
		resolved_limit = Integer(resolve_value(rule.limit, rule_context))
		resolved_period = Integer(resolve_value(rule.period, rule_context))

		count, ttl = rule.count_distinct ? scard_with_ttl(redis_key) : read_with_ttl(redis_key)
		build_result(rule, resolved_limit, resolved_period, count, ttl)
		@@ -195,8 +195,29 @@ module Labkit
		value.to_s
		end

		def resolve_value(val)
		val.respond_to?(:call) ? val.call : val
		# Resolve a limit/period value. Plain values pass through; callables are
		# invoked according to their arity:
		#
		# - Zero-arity callables call with no args (e.g. -> { ApplicationSetting.current.foo }).
		# - Callables with arity >= 1 receive +rule_context+, which may be nil
		# if the caller didn't pass it. They must therefore handle nil -
		# typically with `ctx&.[](:key) \|\| default`.
		# - Variadic callables (negative arity, e.g. ->(*args) { ... }) take
		# the zero-arg path. Opt into rule_context by writing the lambda
		# with exactly one required parameter: ->(ctx) { ... }. This avoids
		# the footgun where ->(*args) silently receives [rule_context] and
		# the caller's overrides never take effect.
		# - Callables that respond to +call+ but not +arity+ (e.g. a class with
		# `def call` and no explicit arity) take the zero-arg path, preserving
		# the pre-rule_context behaviour for custom callable objects.
		#
		# This lets rules carry callables that depend on per-request context (e.g.
		# per-namespace settings) without rebuilding the Rule on every call or
		# smuggling state through globals.
		def resolve_value(val, rule_context = nil)
		return val unless val.respond_to?(:call)

		val.respond_to?(:arity) && val.arity >= 1 ? val.call(rule_context) : val.call
		end

		def encode_char_value(value)

lib/labkit/rate_limit/limiter.rb

+13 −4

Original line number	Diff line number	Diff line
		@@ -35,10 +35,18 @@ module Labkit
		# @param cost [Numeric] amount to add to the counter. Defaults to 1
		# (count-mode). Pass a non-1 Numeric for cost-mode counters such as
		# resource-usage limits; passing 0 reads the counter without writing.
		# @param rule_context [Hash, nil] optional per-request context passed to
		# one-arity callables on +limit+/+period+. Lets rules resolve dynamic
		# configuration (e.g. per-namespace settings) without rebuilding the
		# Rule or doing out-of-band DB queries. Zero-arity callables ignore it.
		# The key contract is owned by the rule's callable, not validated
		# here: if the rule reads ctx[:limit] and the caller passes
		# ctx[:lmit], the callable's fallback branch fires silently. Keep
		# the rule definition and the call site colocated.
		# @return [Result]
		def check(identifier, cost: 1)
		def check(identifier, cost: 1, rule_context: nil)
		id = identifier.is_a?(Identifier) ? identifier : Identifier.new(identifier)
		@evaluator.check(id, cost: cost)
		@evaluator.check(id, cost: cost, rule_context: rule_context)
		end

		# Read the current rate-limit state without incrementing the counter.
		@@ -54,10 +62,11 @@ module Labkit
		# open identically to {#check}.
		#
		# @param identifier [Identifier, Hash] caller attributes for this request
		# @param rule_context [Hash, nil] see {#check}
		# @return [Result]
		def peek(identifier)
		def peek(identifier, rule_context: nil)
		id = identifier.is_a?(Identifier) ? identifier : Identifier.new(identifier)
		@evaluator.peek(id)
		@evaluator.peek(id, rule_context: rule_context)
		end

		private

spec/labkit/rate_limit/evaluator_spec.rb

+128 −0

Original line number	Diff line number	Diff line
		@@ -199,6 +199,134 @@ RSpec.describe Labkit::RateLimit::Evaluator do
		end
		end

		describe "rule_context: per-request callable context" do
		it "passes rule_context to a one-arity limit callable" do
		received = nil
		rule = make_rule(
		name: "ctx_limit",
		limit: lambda { \|ctx\|
		received = ctx
		10
		},
		period: 60
		)

		evaluator(rules: [rule]).check(identifier, rule_context: { limit: 99, period: 100 })

		expect(received).to eq({ limit: 99, period: 100 })
		end

		it "passes rule_context to a one-arity period callable" do
		received = nil
		rule = make_rule(
		name: "ctx_period",
		limit: 10,
		period: lambda { \|ctx\|
		received = ctx
		60
		}
		)

		evaluator(rules: [rule]).check(identifier, rule_context: { period: 100 })

		expect(received).to eq({ period: 100 })
		end

		it "passes nil to a one-arity callable when rule_context is omitted" do
		received = :unset
		rule = make_rule(name: "ctx_nil", limit: lambda { \|ctx\|
		received = ctx
		10
		})

		evaluator(rules: [rule]).check(identifier)

		expect(received).to be_nil
		end

		it "still calls a zero-arity callable with no args even when rule_context is provided" do
		call_count = 0
		rule = make_rule(name: "zero_arity", limit: lambda {
		call_count += 1
		10
		})

		evaluator(rules: [rule]).check(identifier, rule_context: { limit: 99 })

		expect(call_count).to eq(1)
		end

		it "resolves limit from rule_context for an exceeded check" do
		rule = make_rule(
		name: "ctx_exceeded", action: :block,
		limit: ->(ctx) { ctx&.dig(:limit) \|\| 0 },
		period: 60
		)
		# Pre-populate the counter so a low override limit is exceeded immediately.
		raw_redis.set("labkit:rl:rack_request:ctx_exceeded:user:42", "5")
		raw_redis.expire("labkit:rl:rack_request:ctx_exceeded:user:42", 60)

		result = evaluator(rules: [rule]).check(identifier, rule_context: { limit: 3 })

		expect(result.exceeded?).to be(true)
		expect(result.action).to eq(:block)
		expect(result.info.resolved_limit).to eq(3)
		end

		it "passes rule_context through peek as well" do
		received = nil
		rule = make_rule(
		name: "ctx_peek",
		limit: lambda { \|ctx\|
		received = ctx
		10
		}
		)

		evaluator(rules: [rule]).peek(identifier, rule_context: { source: "peek" })

		expect(received).to eq({ source: "peek" })
		end

		it "treats a variadic callable as zero-arity (rule_context not passed)" do
		# Variadic lambdas take the zero-arg path so ->(*args) doesn't silently
		# receive [rule_context] and ignore the caller's overrides. Opt in to
		# rule_context with exactly one required parameter: ->(ctx) { ... }.
		received = :unset
		rule = make_rule(name: "variadic", limit: lambda { \|*args\|
		received = args
		10
		})

		evaluator(rules: [rule]).check(identifier, rule_context: { limit: 7 })

		expect(received).to eq([])
		end

		it "invokes a custom callable without #arity with no args (preserves legacy behaviour)" do
		callable_class = Class.new do
		attr_reader :call_count

		def initialize
		@call_count = 0
		end

		def call
		@call_count += 1
		10
		end
		end
		callable = callable_class.new
		rule = make_rule(name: "no_arity", limit: callable)

		result = evaluator(rules: [rule]).check(identifier, rule_context: { limit: 99 })

		expect(callable.call_count).to eq(1)
		expect(result.matched?).to be(true)
		expect(result.info.resolved_limit).to eq(10)
		end
		end

		describe "EVALSHA / NOSCRIPT fallback (wiring smoke)" do
		# Dispatch correctness (EVALSHA vs EVAL, NOSCRIPT recovery, non-NOSCRIPT
		# propagation) lives in spec/labkit/redis/script_spec.rb. This one test

spec/labkit/rate_limit/limiter_spec.rb

+20 −1

Original line number	Diff line number	Diff line
		@@ -117,7 +117,8 @@ RSpec.describe Labkit::RateLimit::Limiter do

		it "delegates to the evaluator's peek path" do
		lim = limiter
		expect(lim.instance_variable_get(:@evaluator)).to receive(:peek).with(instance_of(Labkit::RateLimit::Identifier))
		expect(lim.instance_variable_get(:@evaluator))
		.to receive(:peek).with(instance_of(Labkit::RateLimit::Identifier), rule_context: nil)
		lim.peek({ user: 42 })
		end

		@@ -236,6 +237,24 @@ RSpec.describe Labkit::RateLimit::Limiter do
		end
		end

		describe "rule_context: kwarg passthrough" do
		it "threads rule_context: through to the evaluator's check" do
		lim = limiter
		expect(lim.instance_variable_get(:@evaluator))
		.to receive(:check)
		.with(instance_of(Labkit::RateLimit::Identifier), cost: 1, rule_context: { limit: 7 })
		lim.check({ user: 42 }, rule_context: { limit: 7 })
		end

		it "threads rule_context: through to the evaluator's peek" do
		lim = limiter
		expect(lim.instance_variable_get(:@evaluator))
		.to receive(:peek)
		.with(instance_of(Labkit::RateLimit::Identifier), rule_context: { source: "peek" })
		lim.peek({ user: 42 }, rule_context: { source: "peek" })
		end
		end

		describe "#check with a count_distinct rule" do
		before do
		# SADD_SCRIPT.eval returns [scard, ttl]