Add Rate Limiter and Sampling Targets Poller Logic

Author: jj22ee

sampler/xray/lib/opentelemetry/sampler/xray/aws_xray_remote_sampler.rb

@@ -4,14 +4,14 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-require 'net/http'
 require 'json'
+require 'net/http'
 require 'opentelemetry/sdk'
-require_relative 'sampling_rule'
+require_relative 'aws_xray_sampling_client'
 require_relative 'fallback_sampler'
-require_relative 'sampling_rule_applier'
 require_relative 'rule_cache'
-require_relative 'aws_xray_sampling_client'
+require_relative 'sampling_rule'
+require_relative 'sampling_rule_applier'
 
 module OpenTelemetry
   module Sampler
@@ -68,7 +68,8 @@ def initialize(endpoint: '127.0.0.1:2000', polling_interval: DEFAULT_RULES_POLLI
           # Start the Sampling Rules poller
           start_sampling_rules_poller
 
-          # TODO: Start the Sampling Targets poller
+          # Start the Sampling Targets poller
+          start_sampling_targets_poller
         end
 
         def should_sample?(trace_id:, parent_context:, links:, name:, kind:, attributes:)
@@ -113,6 +114,15 @@ def start_sampling_rules_poller
           end
         end
 
+        def start_sampling_targets_poller
+          @target_poller = Thread.new do
+            loop do
+              sleep(((@target_polling_interval * 1000) + @target_polling_jitter_millis) / 1000.0)
+              retrieve_and_update_sampling_targets
+            end
+          end
+        end
+
         def retrieve_and_update_sampling_rules
           sampling_rules_response = @sampling_client.fetch_sampling_rules
           if sampling_rules_response&.body && sampling_rules_response.body != ''
@@ -125,6 +135,19 @@ def retrieve_and_update_sampling_rules
           OpenTelemetry.handle_error(exception: e, message: 'Error occurred when retrieving or updating Sampling Rules')
         end
 
+        def retrieve_and_update_sampling_targets
+          request_body = {
+            SamplingStatisticsDocuments: @rule_cache.create_sampling_statistics_documents(@client_id)
+          }
+          sampling_targets_response = @sampling_client.fetch_sampling_targets(request_body)
+          if sampling_targets_response&.body && sampling_targets_response.body != ''
+            response_body = JSON.parse(sampling_targets_response.body)
+            update_sampling_targets(response_body)
+          else
+            OpenTelemetry.logger.debug('SamplingTargets Response is falsy')
+          end
+        end
+
         def update_sampling_rules(response_object)
           sampling_rules = []
           if response_object && response_object['SamplingRuleRecords']
@@ -140,6 +163,33 @@ def update_sampling_rules(response_object)
           end
         end
 
+        def update_sampling_targets(response_object)
+          if response_object && response_object['SamplingTargetDocuments']
+            target_documents = {}
+
+            response_object['SamplingTargetDocuments'].each do |new_target|
+              target_documents[new_target['RuleName']] = new_target
+            end
+
+            refresh_sampling_rules, next_polling_interval = @rule_cache.update_targets(
+              target_documents,
+              response_object['LastRuleModification']
+            )
+
+            @target_polling_interval = next_polling_interval
+
+            if refresh_sampling_rules
+              OpenTelemetry.logger.debug('Performing out-of-band sampling rule polling to fetch updated rules.')
+              @rule_poller&.kill
+              start_sampling_rules_poller
+            end
+          else
+            OpenTelemetry.logger.debug('SamplingTargetDocuments from SamplingTargets request is not defined')
+          end
+        rescue StandardError => e
+          OpenTelemetry.logger.debug("Error occurred when updating Sampling Targets: #{e}")
+        end
+
         class << self
           def generate_client_id
             hex_chars = ('0'..'9').to_a + ('a'..'f').to_a

sampler/xray/lib/opentelemetry/sampler/xray/fallback_sampler.rb

@@ -4,17 +4,24 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+require_relative 'rate_limiting_sampler'
+
 module OpenTelemetry
   module Sampler
     module XRay
       # FallbackSampler samples 1 req/sec and additional 5% of requests using TraceIdRatioBasedSampler.
       class FallbackSampler
         def initialize
           @fixed_rate_sampler = OpenTelemetry::SDK::Trace::Samplers::TraceIdRatioBased.new(0.05)
+          @rate_limiting_sampler = RateLimitingSampler.new(1)
         end
 
         def should_sample?(trace_id:, parent_context:, links:, name:, kind:, attributes:)
-          # TODO: implement and use Rate Limiting Sampler
+          sampling_result = @rate_limiting_sampler.should_sample?(
+            trace_id: trace_id, parent_context: parent_context, links: links, name: name, kind: kind, attributes: attributes
+          )
+
+          return sampling_result if sampling_result.instance_variable_get(:@decision) != OpenTelemetry::SDK::Trace::Samplers::Decision::DROP
 
           @fixed_rate_sampler.should_sample?(trace_id: trace_id, parent_context: parent_context, links: links, name: name, kind: kind, attributes: attributes)
         end

sampler/xray/lib/opentelemetry/sampler/xray/rate_limiter.rb

@@ -0,0 +1,48 @@
+# frozen_string_literal: true
+
+# Copyright OpenTelemetry Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+module OpenTelemetry
+  module Sampler
+    module XRay
+      # RateLimiter keeps track of the current reservoir quota balance available (measured via available time)
+      # If enough time has elapsed, the RateLimiter will allow quota balance to be consumed/taken (decrease available time)
+      # A RateLimitingSampler uses this RateLimiter to determine if it should sample or not based on the quota balance available.
+      class RateLimiter
+        def initialize(quota, max_balance_in_seconds = 1)
+          @max_balance_millis = max_balance_in_seconds * 1000.0
+          @quota = quota
+          @wallet_floor_millis = Time.now.to_f * 1000
+          # current "balance" would be `ceiling - floor`
+          @lock = Mutex.new
+        end
+
+        def take(cost = 1)
+          return false if @quota.zero?
+
+          quota_per_millis = @quota / 1000.0
+
+          # assume divide by zero not possible
+          cost_in_millis = cost / quota_per_millis
+
+          @lock.synchronize do
+            wallet_ceiling_millis = Time.now.to_f * 1000
+            current_balance_millis = wallet_ceiling_millis - @wallet_floor_millis
+            current_balance_millis = [current_balance_millis, @max_balance_millis].min
+            pending_remaining_balance_millis = current_balance_millis - cost_in_millis
+
+            if pending_remaining_balance_millis >= 0
+              @wallet_floor_millis = wallet_ceiling_millis - pending_remaining_balance_millis
+              return true
+            end
+
+            # No changes to the wallet state
+            false
+          end
+        end
+      end
+    end
+  end
+end

sampler/xray/lib/opentelemetry/sampler/xray/rate_limiting_sampler.rb

@@ -0,0 +1,43 @@
+# frozen_string_literal: true
+
+# Copyright OpenTelemetry Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+require_relative 'rate_limiter'
+
+module OpenTelemetry
+  module Sampler
+    module XRay
+      # RateLimitingSampler is a Sampler that uses a RateLimiter to determine
+      # if it should sample or not based on the quota balance available.
+      class RateLimitingSampler
+        def initialize(quota)
+          @quota = quota
+          @reservoir = RateLimiter.new(quota)
+        end
+
+        def should_sample?(trace_id:, parent_context:, links:, name:, kind:, attributes:)
+          tracestate = OpenTelemetry::Trace.current_span(parent_context).context.tracestate
+          if @reservoir.take(1)
+            OpenTelemetry::SDK::Trace::Samplers::Result.new(
+              decision: OpenTelemetry::SDK::Trace::Samplers::Decision::RECORD_AND_SAMPLE,
+              tracestate: tracestate,
+              attributes: attributes
+            )
+          else
+            OpenTelemetry::SDK::Trace::Samplers::Result.new(
+              decision: OpenTelemetry::SDK::Trace::Samplers::Decision::DROP,
+              tracestate: tracestate,
+              attributes: attributes
+            )
+          end
+        end
+
+        def to_s
+          "RateLimitingSampler{rate limiting sampling with sampling config of #{@quota} req/sec and 0% of additional requests}"
+        end
+      end
+    end
+  end
+end

sampler/xray/lib/opentelemetry/sampler/xray/rule_cache.rb

@@ -53,6 +53,52 @@ def update_rules(new_rule_appliers)
           end
         end
 
+        def create_sampling_statistics_documents(client_id)
+          statistics_documents = []
+
+          @cache_lock.synchronize do
+            @rule_appliers.each do |rule|
+              statistics = rule.snapshot_statistics
+              now_in_seconds = Time.now.to_i
+
+              sampling_statistics_doc = {
+                ClientID: client_id,
+                RuleName: rule.sampling_rule.rule_name,
+                Timestamp: now_in_seconds,
+                RequestCount: statistics.request_count,
+                BorrowCount: statistics.borrow_count,
+                SampledCount: statistics.sample_count
+              }
+
+              statistics_documents << sampling_statistics_doc
+            end
+          end
+
+          statistics_documents
+        end
+
+        def update_targets(target_documents, last_rule_modification)
+          min_polling_interval = nil
+          next_polling_interval = DEFAULT_TARGET_POLLING_INTERVAL_SECONDS
+
+          @cache_lock.synchronize do
+            @rule_appliers.each_with_index do |rule, index|
+              target = target_documents[rule.sampling_rule.rule_name]
+              if target
+                @rule_appliers[index] = rule.with_target(target)
+                min_polling_interval = target['Interval'] if target['Interval'] && (min_polling_interval.nil? || min_polling_interval > target['Interval'])
+              else
+                OpenTelemetry.logger.debug('Invalid sampling target: missing rule name')
+              end
+            end
+
+            next_polling_interval = min_polling_interval if min_polling_interval
+
+            refresh_sampling_rules = last_rule_modification * 1000 > @last_updated_epoch_millis
+            return [refresh_sampling_rules, next_polling_interval]
+          end
+        end
+
         private
 
         def sort_rules_by_priority

sampler/xray/lib/opentelemetry/sampler/xray/sampling_rule_applier.rb

@@ -9,6 +9,7 @@
 require 'date'
 require_relative 'sampling_rule'
 require_relative 'statistics'
+require_relative 'rate_limiting_sampler'
 require_relative 'utils'
 
 module OpenTelemetry
@@ -26,10 +27,24 @@ def initialize(sampling_rule, statistics = OpenTelemetry::Sampler::XRay::Statist
           @sampling_rule = sampling_rule
           @fixed_rate_sampler = OpenTelemetry::SDK::Trace::Samplers::TraceIdRatioBased.new(@sampling_rule.fixed_rate)
 
-          # TODO: Add Reservoir Sampler (Rate Limiting Sampler)
+          @reservoir_sampler = if @sampling_rule.reservoir_size.positive?
+                                 OpenTelemetry::Sampler::XRay::RateLimitingSampler.new(1)
+                               else
+                                 OpenTelemetry::Sampler::XRay::RateLimitingSampler.new(0)
+                               end
 
           @reservoir_expiry_time = MAX_DATE_TIME_SECONDS
           @statistics = statistics
+          @statistics_lock = Mutex.new
+
+          @statistics.reset_statistics
+          @borrowing_enabled = true
+
+          apply_target(target) if target
+        end
+
+        def with_target(target)
+          self.class.new(@sampling_rule, @statistics, target)
         end
 
         def matches?(attributes, resource)
@@ -78,26 +93,63 @@ def matches?(attributes, resource)
         end
 
         def should_sample?(trace_id:, parent_context:, links:, name:, kind:, attributes:)
-          # TODO: Record Sampling Statistics
-
+          has_borrowed = false
           result = OpenTelemetry::SDK::Trace::Samplers::Result.new(
             decision: OpenTelemetry::SDK::Trace::Samplers::Decision::DROP,
             tracestate: OpenTelemetry::Trace::Tracestate::DEFAULT
           )
 
-          # TODO: Apply Reservoir Sampling
+          now = Time.now
+          reservoir_expired = now >= @reservoir_expiry_time
+
+          unless reservoir_expired
+            result = @reservoir_sampler.should_sample?(
+              trace_id: trace_id, parent_context: parent_context, links: links, name: name, kind: kind, attributes: attributes
+            )
+            has_borrowed = @borrowing_enabled && result.instance_variable_get(:@decision) != OpenTelemetry::SDK::Trace::Samplers::Decision::DROP
+          end
 
           if result.instance_variable_get(:@decision) == OpenTelemetry::SDK::Trace::Samplers::Decision::DROP
             result = @fixed_rate_sampler.should_sample?(
               trace_id: trace_id, parent_context: parent_context, links: links, name: name, kind: kind, attributes: attributes
             )
           end
 
+          @statistics_lock.synchronize do
+            @statistics.sample_count += result.instance_variable_get(:@decision) == OpenTelemetry::SDK::Trace::Samplers::Decision::DROP ? 0 : 1
+            @statistics.borrow_count += has_borrowed ? 1 : 0
+            @statistics.request_count += 1
+          end
+
           result
         end
 
+        def snapshot_statistics
+          @statistics_lock.synchronize do
+            statistics_copy = @statistics.dup
+            @statistics.reset_statistics
+            return statistics_copy
+          end
+        end
+
         private
 
+        def apply_target(target)
+          @borrowing_enabled = false
+
+          @reservoir_sampler = OpenTelemetry::Sampler::XRay::RateLimitingSampler.new(target['ReservoirQuota']) if target['ReservoirQuota']
+
+          @reservoir_expiry_time = if target['ReservoirQuotaTTL']
+                                     Time.at(target['ReservoirQuotaTTL'])
+                                   else
+                                     Time.now
+                                   end
+
+          return unless target['FixedRate']
+
+          @fixed_rate_sampler = OpenTelemetry::SDK::Trace::Samplers::TraceIdRatioBased.new(target['FixedRate'])
+        end
+
         def get_arn(resource, attributes)
           resource_hash = resource.attribute_enumerator.to_h
           arn = resource_hash[SEMCONV::Resource::AWS_ECS_CONTAINER_ARN] ||