Skip to content

Commit

Permalink
DEBUG-2334 default token bucket allow? size to 1
Browse files Browse the repository at this point in the history
The only usage of token bucket rate limiter currently in
our code uses a size of 1 (to perform one action).
Default the size parameter to 1 to simplify the API for
using token bucket on the caller side.

Dynamic instrumentation will also use token bucket and
will also call allow? with the size of 1.
  • Loading branch information
p committed Sep 3, 2024
1 parent ae50107 commit 2d10912
Show file tree
Hide file tree
Showing 5 changed files with 29 additions and 7 deletions.
8 changes: 4 additions & 4 deletions lib/datadog/core/rate_limiter.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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?`.
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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?

Expand Down Expand Up @@ -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

Expand Down
2 changes: 1 addition & 1 deletion lib/datadog/tracing/sampling/rule_sampler.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down
2 changes: 1 addition & 1 deletion lib/datadog/tracing/sampling/span/rule.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
22 changes: 22 additions & 0 deletions spec/datadog/core/rate_limiter_spec.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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 }

Expand Down
2 changes: 1 addition & 1 deletion spec/datadog/tracing/sampling/rule_sampler_spec.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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 }|
Expand Down

0 comments on commit 2d10912

Please sign in to comment.