From 2d109122b7979739393488033f7b6db9a052587f Mon Sep 17 00:00:00 2001
From: Oleg Pudeyev <code@olegp.name>
Date: Tue, 3 Sep 2024 08:47:48 -0400
Subject: [PATCH] DEBUG-2334 default token bucket allow? size to 1

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.
---
 lib/datadog/core/rate_limiter.rb              |  8 +++----
 lib/datadog/tracing/sampling/rule_sampler.rb  |  2 +-
 lib/datadog/tracing/sampling/span/rule.rb     |  2 +-
 spec/datadog/core/rate_limiter_spec.rb        | 22 +++++++++++++++++++
 .../tracing/sampling/rule_sampler_spec.rb     |  2 +-
 5 files changed, 29 insertions(+), 7 deletions(-)

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 }|