aws-xray-sdk 0.10.2 → 0.11.0

data/lib/aws-xray-sdk/sampling/local/sampling_rule.rb

@@ -0,0 +1,63 @@
+require 'aws-xray-sdk/exceptions'
+require 'aws-xray-sdk/sampling/local/reservoir'
+require 'aws-xray-sdk/search_pattern'
+
+module XRay
+  # One SamplingRule object represents one rule defined from the rules hash definition.
+  # It can be either a custom rule or the default rule.
+  class LocalSamplingRule
+    attr_reader :fixed_target, :rate, :host,
+                :method, :path, :reservoir, :default
+
+    # @param [Hash] rule_definition Hash that defines a single rule.
+    # @param default A boolean flag indicates if this rule is the default rule.
+    def initialize(rule_definition:, default: false)
+      @fixed_target = rule_definition[:fixed_target]
+      @rate = rule_definition[:rate]
+
+      @host = rule_definition[:host]
+      @method = rule_definition[:http_method]
+      @path = rule_definition[:url_path]
+
+      @default = default
+      validate
+      @reservoir = LocalReservoir.new traces_per_sec: @fixed_target
+    end
+
+    # Determines whether or not this sampling rule applies to
+    # the incoming request based on some of the request's parameters.
+    # Any None parameters provided will be considered an implicit match.
+    def applies?(sampling_req)
+      return false if sampling_req.nil? || sampling_req.empty?
+
+      host = sampling_req[:host]
+      url_path = sampling_req[:url_path]
+      http_method = sampling_req[:http_method]
+
+      host_match = !host || SearchPattern.wildcard_match?(pattern: @host, text: host)
+      path_match = !url_path || SearchPattern.wildcard_match?(pattern: @path, text: url_path)
+      method_match = !http_method || SearchPattern.wildcard_match?(pattern: @method, text: http_method)
+      host_match && path_match && method_match
+    end
+
+    private
+
+    def validate
+      if @fixed_target < 0 || @rate < 0
+        raise InvalidSamplingConfigError, 'All rules must have non-negative values for fixed_target and rate.'
+      end
+
+      if @default
+        # validate default rule
+        if @host || @method || @path
+          raise InvalidSamplingConfigError, 'The default rule must not specify values for url_path, service_name, or http_method.'
+        end
+      else
+        # validate custom rule
+        unless @host && @method && @path
+          raise InvalidSamplingConfigError, 'All non-default rules must have values for url_path, service_name, and http_method.'
+        end
+      end
+    end
+  end
+end

data/lib/aws-xray-sdk/sampling/reservoir.rb

@@ -1,35 +1,79 @@
+require 'date'
+require 'aws-xray-sdk/sampling/sampling_decision'
+
 module XRay
-  # Keeps track of the number of sampled segments within
-  # a single second. This class is implemented to be
-  # thread-safe to achieve accurate sampling.
+  # Centralized thread-safe reservoir which holds fixed sampling
+  # quota for the current instance, borrowed count and TTL.
   class Reservoir
-    # @param [Integer] traces_per_sec The number of guranteed sampled
-    #   segments per second.
-    def initialize(traces_per_sec: 0)
-      @traces_per_sec = traces_per_sec
-      @used_this_sec = 0
-      @this_sec = Time.now.to_i
+    attr_reader :quota, :ttl
+
+    def initialize
+      @quota = nil
+      @ttl = nil
+
+      @this_sec = 0
+      @taken_this_sec = 0
+      @borrowed_this_sec = 0
+
+      @report_interval = 1
+      @report_elapsed = 0
+
       @lock = Mutex.new
     end
 
-    # Returns `true` if there are segments left within the
-    # current second, otherwise returns `false`.
-    def take
-      # nothing to provide if reserved is set to 0
-      return false if @traces_per_sec.zero?
+    # Decide whether to borrow or take one quota from
+    # the reservoir. Return `false` if it can neither
+    # borrow nor take. This method is thread-safe.
+    def borrow_or_take(now, borrowable)
       @lock.synchronize do
-        now = Time.now.to_i
-        # refresh time frame
-        if now != @this_sec
-          @used_this_sec = 0
-          @this_sec = now
+        reset_new_sec(now)
+        # Don't borrow if the quota is available and fresh.
+        if quota_fresh?(now)
+          return SamplingDecision::NOT_SAMPLE if @taken_this_sec >= @quota
+          @taken_this_sec += 1
+          return SamplingDecision::TAKE
         end
-        # return false if reserved item ran out
-        return false unless @used_this_sec < @traces_per_sec
-        # otherwise increment used counter and return true
-        @used_this_sec += 1
-        return true
+
+        # Otherwise try to borrow if the quota is not present or expired.
+        if borrowable
+          return SamplingDecision::NOT_SAMPLE if @borrowed_this_sec >= 1
+          @borrowed_this_sec += 1
+          return SamplingDecision::BORROW
+        end
+
+        # Cannot sample if quota expires and cannot borrow
+        SamplingDecision::NOT_SAMPLE
+      end
+    end
+
+    def load_target_info(quota:, ttl:, interval:)
+      @quota = quota unless quota.nil?
+      @ttl = ttl.to_i unless ttl.nil?
+      @interval = interval / 10 unless interval.nil?
+    end
+
+    def time_to_report?
+      if @report_elapsed + 1 >= @report_interval
+        @report_elapsed = 0
+        true
+      else
+        @report_elapsed += 1
+        false
       end
     end
+
+    private
+
+    # Reset the counter if now enters a new one-second window
+    def reset_new_sec(now)
+      return if now == @this_sec
+      @taken_this_sec = 0
+      @borrowed_this_sec = 0
+      @this_sec = now
+    end
+
+    def quota_fresh?(now)
+      @quota && @quota >= 0 && @ttl && @ttl >= now
+    end
   end
 end

data/lib/aws-xray-sdk/sampling/rule_cache.rb

@@ -0,0 +1,86 @@
+require 'aws-xray-sdk/logger'
+
+module XRay
+  # Cache sampling rules and quota retrieved by `TargetPoller`
+  # and `RulePoller`. It will not return anything if it expires.
+  class RuleCache
+    include Logging
+    attr_accessor :last_updated
+    @@TTL = 60 * 60 # 1 hour
+
+    def initialize
+      @rules = []
+      @last_updated = nil
+      @lock = Mutex.new
+    end
+
+    def get_matched_rule(sampling_req, now: Time.now.to_i)
+      return nil if expired?(now)
+      matched = nil
+      rules.each do |rule|
+        matched = rule if matched.nil? && rule.applies?(sampling_req)
+        matched = rule if matched.nil? && rule.default?
+      end
+      matched
+    end
+
+    def load_rules(new_rules)
+      @lock.synchronize do
+        # Simply assign rules and sort if cache is empty
+        if @rules.empty?
+          @rules = new_rules
+          return sort_rules
+        end
+
+        # otherwise we need to merge new rules and current rules
+        curr_rules = {}
+        @rules.each do |rule|
+          curr_rules[rule.name] = rule
+        end
+        # Update the rules in the cache
+        @rules = new_rules
+        # Transfer state information
+        @rules.each do |rule|
+          curr_rule = curr_rules[rule.name]
+          rule.merge(curr_rule) unless curr_rule.nil?
+        end
+        sort_rules
+      end
+    end
+
+    def load_targets(targets_h)
+      @lock.synchronize do
+        @rules.each do |rule|
+          target = targets_h[rule.name]
+          next if target.nil?
+          rule.rate = target.fixed_rate
+          rule.reservoir.load_target_info(
+            quota: target.reservoir_quota,
+            ttl: target.reservoir_quota_ttl,
+            interval: target.interval
+          )
+        end
+      end
+    end
+
+    def rules
+      @lock.synchronize do
+        @rules
+      end
+    end
+
+    private
+
+    # The cache should maintain the order of the rules based on
+    # priority. If priority is the same we sort name by alphabet
+    # as rule name is unique.
+    def sort_rules
+      @rules.sort_by! { |rule| [rule.priority, rule.name] }      
+    end
+
+    def expired?(now)
+      # The cache is treated as expired if it is never loaded.
+      @last_updated.nil? || now > @last_updated + @@TTL
+    end
+  end
+end

data/lib/aws-xray-sdk/sampling/rule_poller.rb

@@ -0,0 +1,39 @@
+require 'aws-xray-sdk/logger'
+
+module XRay
+  # Polls sampling rules from X-Ray service
+  class RulePoller
+    include Logging
+    attr_reader :cache, :connector
+
+    def initialize(cache:, connector:)
+      @cache = cache
+      @connector = connector
+      @worker = Thread.new { poll }
+    end
+
+    def run
+      @worker.run
+    end
+
+    private
+
+    def poll
+      loop do
+        Thread.stop
+        refresh_cache
+      end
+    end
+
+    def refresh_cache
+      now = Time.now.to_i
+      rules = @connector.fetch_sampling_rules
+      unless rules.nil? || rules.empty?
+        @cache.load_rules(rules)
+        @cache.last_updated = now
+      end
+    rescue StandardError => e
+      logger.warn %(failed to fetch X-Ray sampling rules due to #{e.backtrace})
+    end
+  end
+end

data/lib/aws-xray-sdk/sampling/sampler.rb

@@ -6,12 +6,12 @@ module XRay
   module Sampler
     # Decides if a segment should be sampled for an incoming request.
     # Used in case of middleware.
-    def sample_request?(service_name:, url_path:, http_method:)
+    def sample_request?(sampling_req:)
       raise 'Not implemented'
     end
 
-    # Decides if a segment should be sampled merely based on internal
-    # sampling rules.
+    # Sample purely based on cached sampling rules without
+    # any incoming rules matching information.
     def sample?
       raise 'Not implemented'
     end

data/lib/aws-xray-sdk/sampling/sampling_decision.rb

@@ -0,0 +1,8 @@
+module XRay
+  # Stores the enum style sampling decisions for default sampler
+  module SamplingDecision
+    TAKE = 'take'.freeze
+    BORROW = 'borrow'.freeze
+    NOT_SAMPLE = 'no'.freeze
+  end
+end

data/lib/aws-xray-sdk/sampling/sampling_rule.rb

@@ -3,55 +3,122 @@ require 'aws-xray-sdk/sampling/reservoir'
 require 'aws-xray-sdk/search_pattern'
 
 module XRay
-  # One SamplingRule object represents one rule defined from the rules hash definition.
-  # It can be either a custom rule or the default rule.
+  # Service sampling rule data model
   class SamplingRule
-    attr_reader :fixed_target, :rate, :service_name,
-                :method, :path, :reservoir, :default
+    attr_reader :name, :priority,
+                :request_count, :borrow_count, :sampled_count
+    attr_accessor :reservoir, :rate
 
-    # @param [Hash] rule_definition Hash that defines a single rule.
-    # @param default A boolean flag indicates if this rule is the default rule.
-    def initialize(rule_definition:, default: false)
-      @fixed_target = rule_definition[:fixed_target]
-      @rate = rule_definition[:rate]
+    # @param Struct defined here https://docs.aws.amazon.com/sdk-for-ruby/v3/api/Aws/XRay/Types/SamplingRule.html.
+    def initialize(record)
+      @name = record.rule_name
+      @priority = record.priority
+      @rate = record.fixed_rate
 
-      @service_name = rule_definition[:service_name]
-      @method = rule_definition[:http_method]
-      @path = rule_definition[:url_path]
+      @host = record.host
+      @method = record.http_method
+      @path = record.url_path
+      @service = record.service_name
+      @service_type = record.service_type
 
-      @default = default
-      validate
-      @reservoir = Reservoir.new traces_per_sec: @fixed_target
+      @reservoir_size = record.reservoir_size
+      @reservoir = Reservoir.new
+      reset_statistics
+
+      @lock = Mutex.new
     end
 
     # Determines whether or not this sampling rule applies to
     # the incoming request based on some of the request's parameters.
-    # Any None parameters provided will be considered an implicit match.
-    def applies?(target_name:, target_path:, target_method:)
-      name_match   = !target_name   || SearchPattern.wildcard_match?(pattern: @service_name, text: target_name)
-      path_match   = !target_path   || SearchPattern.wildcard_match?(pattern: @path, text: target_path)
-      method_match = !target_method || SearchPattern.wildcard_match?(pattern: @method, text: target_method)
-      name_match && path_match && method_match
+    # Any Nil parameters provided will be considered as implicit matches
+    # as the rule matching is a best effort.
+    def applies?(sampling_req)
+      return false if sampling_req.nil? || sampling_req.empty?
+
+      host = sampling_req[:host]
+      http_method = sampling_req[:http_method]
+      url_path = sampling_req[:url_path]
+      service = sampling_req[:service]
+
+      host_match = !host || SearchPattern.wildcard_match?(pattern: @host, text: host)
+      path_match = !url_path || SearchPattern.wildcard_match?(pattern: @path, text: url_path)
+      method_match = !http_method || SearchPattern.wildcard_match?(pattern: @method, text: http_method)
+      service_match = !service || SearchPattern.wildcard_match?(pattern: @service, text: service)
+
+      # if sampling request contains service type we assmue
+      # the origin (a.k.a AWS plugins are set and effective)
+      if sampling_req.key?(:service_type)
+        service_type = sampling_req[:service_type]
+        service_type_match = SearchPattern.wildcard_match?(pattern: @service_type, text: service_type)
+      else
+        service_type_match = @service_type == '*'
+      end
+      host_match && path_match && method_match && service_match && service_type_match
     end
 
-    private
+    def snapshot_statistics
+      @lock.synchronize do
+        report = {
+          request_count: @request_count,
+          borrow_count: @borrow_count,
+          sampled_count: @sampled_count
+        }
+        reset_statistics
+        report
+      end
+    end
 
-    def validate
-      if @fixed_target < 0 || @rate < 0
-        raise InvalidSamplingConfigError, 'All rules must have non-negative values for fixed_target and rate.'
+    def merge(rule)
+      @lock.synchronize do
+        @request_count = rule.request_count
+        @borrow_count = rule.borrow_count
+        @sampled_count = rule.sampled_count
+        @reservoir = rule.reservoir
+        rule.reservoir = nil
       end
+    end
 
-      if @default
-        # validate default rule
-        if @service_name || @method || @path
-          raise InvalidSamplingConfigError, 'The default rule must not specify values for url_path, service_name, or http_method.'
-        end
-      else
-        # validate custom rule
-        unless @service_name && @method && @path
-          raise InvalidSamplingConfigError, 'All non-default rules must have values for url_path, service_name, and http_method.'
-        end
+    def borrowable?
+      @reservoir_size != 0
+    end
+
+    # Return `true` if this rule is the default rule.
+    def default?
+      @name == 'Default'
+    end
+
+    def ever_matched?
+      @request_count > 0
+    end
+
+    def time_to_report?
+      @reservoir.time_to_report?
+    end
+
+    def increment_request_count
+      @lock.synchronize do
+        @request_count += 1
+      end
+    end
+
+    def increment_borrow_count
+      @lock.synchronize do
+        @borrow_count += 1
       end
     end
+
+    def increment_sampled_count
+      @lock.synchronize do
+        @sampled_count += 1
+      end
+    end
+
+    private
+
+    def reset_statistics
+      @request_count = 0
+      @borrow_count = 0
+      @sampled_count = 0
+    end
   end
 end