diff --git a/lib/datadog/core/rate_limiter.rb b/lib/datadog/core/rate_limiter.rb index 0b749fcee2b..61a55aa2d22 100644 --- a/lib/datadog/core/rate_limiter.rb +++ b/lib/datadog/core/rate_limiter.rb @@ -13,7 +13,7 @@ class RateLimiter # to be side-effect free. # # @return [Boolean] whether a resource conforms with the current limit - def allow?(size); end + def allow?(size = 1); end # The effective rate limiting ratio based on # recent calls to `allow?`. @@ -61,7 +61,7 @@ def initialize(rate, max_tokens = rate) # the tokens from the bucket. # # @return [Boolean] +true+ if message conforms with current bucket limit - def allow?(size) + def allow?(size = 1) allowed = should_allow?(size) update_rate_counts(allowed) allowed @@ -125,7 +125,7 @@ def increment_conforming_count @conforming_messages += 1 end - def should_allow?(size) + def should_allow?(size = 1) # rate limit of 0 blocks everything return false if @rate.zero? @@ -170,7 +170,7 @@ def update_rate_counts(allowed) # with no limits. class UnlimitedLimiter < RateLimiter # @return [Boolean] always +true+ - def allow?(_) + def allow?(_ = 1) true end diff --git a/lib/datadog/tracing/sampling/rule_sampler.rb b/lib/datadog/tracing/sampling/rule_sampler.rb index daf625c87d3..4b7e812e75f 100644 --- a/lib/datadog/tracing/sampling/rule_sampler.rb +++ b/lib/datadog/tracing/sampling/rule_sampler.rb @@ -121,7 +121,7 @@ def sample_trace(trace) return false unless sampled - rate_limiter.allow?(1).tap do |allowed| + rate_limiter.allow?.tap do |allowed| set_priority(trace, allowed) set_limiter_metrics(trace, rate_limiter.effective_rate) diff --git a/lib/datadog/tracing/sampling/span/rule.rb b/lib/datadog/tracing/sampling/span/rule.rb index 827e17a6821..45fff8664a2 100644 --- a/lib/datadog/tracing/sampling/span/rule.rb +++ b/lib/datadog/tracing/sampling/span/rule.rb @@ -54,7 +54,7 @@ def initialize( def sample!(trace_op, span_op) return :not_matched unless @matcher.match?(span_op) - if @rate_limiter.allow?(1) && @sampler.sample!(trace_op) + if @rate_limiter.allow? && @sampler.sample!(trace_op) span_op.set_metric(Span::Ext::TAG_MECHANISM, Sampling::Ext::Mechanism::SPAN_SAMPLING_RATE) span_op.set_metric(Span::Ext::TAG_RULE_RATE, @sample_rate) span_op.set_metric(Span::Ext::TAG_MAX_PER_SECOND, @rate_limit) diff --git a/spec/datadog/core/rate_limiter_spec.rb b/spec/datadog/core/rate_limiter_spec.rb index 11bbd1c146d..6567e1760de 100644 --- a/spec/datadog/core/rate_limiter_spec.rb +++ b/spec/datadog/core/rate_limiter_spec.rb @@ -102,6 +102,28 @@ end end + context 'when size is not given' do + subject(:allow?) { bucket.allow? } + + context 'when tokens are available' do + it 'returns true' do + is_expected.to be true + end + end + + context 'when tokens are not available' do + let(:max_tokens) { 1 } + + before do + bucket.allow? + end + + it 'returns false' do + is_expected.to be false + end + end + end + context 'with negative rate' do let(:rate) { -1 } diff --git a/spec/datadog/tracing/sampling/rule_sampler_spec.rb b/spec/datadog/tracing/sampling/rule_sampler_spec.rb index 15d1e0808ad..43c2ea25d22 100644 --- a/spec/datadog/tracing/sampling/rule_sampler_spec.rb +++ b/spec/datadog/tracing/sampling/rule_sampler_spec.rb @@ -17,7 +17,7 @@ before do allow(rate_limiter).to receive(:effective_rate).and_return(effective_rate) - allow(rate_limiter).to receive(:allow?).with(1).and_return(allow?) + allow(rate_limiter).to receive(:allow?).and_return(allow?) end shared_examples 'a simple rule that matches all span operations' do |options = { sample_rate: 1.0 }|