aws-xray-sdk 0.10.2 → 0.11.0

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

@@ -0,0 +1,72 @@
+require 'securerandom'
+require 'aws-xray-sdk/sampling/sampling_rule'
+require 'aws-xray-sdk/logger'
+
+module XRay
+  # Connector class that translates Sampling poller functions to
+  # actual X-Ray back-end APIs and communicates with X-Ray daemon
+  # as the signing proxy.
+  class ServiceConnector
+    include Logging
+    attr_accessor :xray_client
+    @@client_id = SecureRandom.hex(12)
+
+    def initialize
+      update_xray_client
+    end
+
+    def fetch_sampling_rules
+      rules = []
+      records = @xray_client.get_sampling_rules.sampling_rule_records
+      records.each { |record| rules << SamplingRule.new(record.sampling_rule) if rule_valid?(record.sampling_rule) }
+      rules
+    end
+
+    def fetch_sampling_targets(rules)
+      now = Time.now.to_i
+      reports = generate_reports(rules, now)
+      resp = @xray_client.get_sampling_targets({ sampling_statistics_documents: reports })
+      {
+        last_modified: resp.last_rule_modification,
+        documents: resp.sampling_target_documents
+      }
+    end
+
+    def update_xray_client(ip: '127.0.0.1', port: 2000)
+      require 'aws-sdk-xray'
+      @xray_client = Aws::XRay::Client.new(
+        endpoint: %(http://#{ip}:#{port}),
+        access_key_id: 'dummy', # AWS Ruby SDK doesn't support unsigned request
+        secret_access_key: 'dummy',
+        region: 'us-west-2' # not used
+      )
+    end
+
+    def daemon_config=(v)
+      update_xray_client ip: v.tcp_ip, port: v.tcp_port
+    end
+
+    private
+
+    def generate_reports(rules, now)
+      reports = []
+      rules.each do |rule|
+        report = rule.snapshot_statistics
+        report[:rule_name] = rule.name
+        report[:timestamp] = now
+        report[:client_id] = @@client_id
+        reports << report
+      end
+      reports
+    end
+
+    def rule_valid?(rule)
+      return false if rule.version != 1
+      # rules has resource ARN and attributes configured
+      # doesn't apply to this SDK
+      return false unless rule.resource_arn == '*'
+      return false unless rule.attributes.nil? || rule.attributes.empty?
+      true
+    end
+  end
+end

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

@@ -1,105 +1,99 @@
+require 'aws-xray-sdk/logger'
+require 'aws-xray-sdk/sampling/local/sampler'
+require 'aws-xray-sdk/sampling/lead_poller'
+require 'aws-xray-sdk/sampling/rule_cache'
 require 'aws-xray-sdk/sampling/sampler'
 require 'aws-xray-sdk/sampling/sampling_rule'
-require 'aws-xray-sdk/exceptions'
+require 'aws-xray-sdk/sampling/sampling_decision'
 
 module XRay
-  # The default sampler that uses internally defined
-  # sampling rule and reservoir models to decide sampling decision.
-  # It also uses the default sampling rule.
-  # An example definition:
-  #   {
-  #     version: 1,
-  #     rules: [
-  #       {
-  #         description: 'Player moves.',
-  #         service_name: '*',
-  #         http_method: '*',
-  #         url_path: '/api/move/*',
-  #         fixed_target: 0,
-  #         rate: 0.05
-  #       }
-  #     ],
-  #     default: {
-  #       fixed_target: 1,
-  #       rate: 0.1
-  #     }
-  #   }
-  # This example defines one custom rule and a default rule.
-  # The custom rule applies a five-percent sampling rate with no minimum
-  # number of requests to trace for paths under /api/move/. The default
-  # rule traces the first request each second and 10 percent of additional requests.
-  # The SDK applies custom rules in the order in which they are defined.
-  # If a request matches multiple custom rules, the SDK applies only the first rule.
+  # Making sampling decisions based on service sampling rules defined
+  # by X-Ray control plane APIs. It will fall back to local sampling rules
+  # if service sampling rules are not available or expired.
   class DefaultSampler
     include Sampler
-    DEFAULT_RULES = {
-      version: 1,
-      default: {
-        fixed_target: 1,
-        rate: 0.05
-      },
-      rules: []
-    }.freeze
+    include Logging
+    attr_reader :cache, :local_sampler, :poller
+    attr_accessor :origin
 
     def initialize
-      load_sampling_rules(DEFAULT_RULES)
+      @local_sampler = LocalSampler.new
+      @cache = RuleCache.new
+      @poller = LeadPoller.new(@cache)
+
+      @started = false
+      @origin = nil
+      @lock = Mutex.new
     end
 
-    # Return True if the sampler decide to sample based on input
-    # information and sampling rules. It will first check if any
-    # custom rule should be applied, if not it falls back to the
-    # default sampling rule.
-    # All arugments are extracted from incoming requests by
-    # X-Ray middleware to perform path based sampling.
-    def sample_request?(service_name:, url_path:, http_method:)
-      # directly fallback to non-path-based if all arguments are nil
-      return sample? unless service_name || url_path || http_method
-      @custom_rules ||= []
-      @custom_rules.each do |c|
-        return should_sample?(c) if c.applies?(target_name: service_name, target_path: url_path, target_method: http_method)
+    # Start background threads to poll sampling rules
+    def start
+      @lock.synchronize do
+        unless @started
+          @poller.start
+          @started = true
+        end
+      end
+    end
+
+    # Return the rule name if it decides to sample based on
+    # a service sampling rule matching. If there is no match
+    # it will fallback to local defined sampling rules.
+    def sample_request?(sampling_req)
+      start unless @started
+      now = Time.now.to_i
+      if sampling_req.nil?
+        sampling_req = { service_type: @origin } if @origin
+      elsif !sampling_req.key?(:service_type)
+        sampling_req[:service_type] = @origin if @origin
+      end
+
+      matched_rule = @cache.get_matched_rule(sampling_req, now: now)
+      if !matched_rule.nil?
+        logger.debug %(Rule #{matched_rule.name} is selected to make a sampling decision.')
+        process_matched_rule(matched_rule, now)
+      else
+        logger.warn %(No effective centralized sampling rule match. Fallback to local rules.)
+        @local_sampler.sample_request?(sampling_req)
       end
-      sample?
     end
 
-    # Decides if should sample based on non-path-based rule.
-    # Currently only the default rule is not path-based.
     def sample?
-      should_sample?(@default_rule)
+      sample_request? nil
     end
 
-    # @param [Hash] v The sampling rules definition.
+    # @param [Hash] v Local sampling rules definition.
+    # This configuration has lower priority than service
+    # sampling rules and only has effect when those rules
+    # are not available or expired.
     def sampling_rules=(v)
-      load_sampling_rules(v)
+      @local_sampler.sampling_rules = v
     end
 
-    # @return [Array] An array of [SamplingRule]
-    def sampling_rules
-      all_rules = []
-      all_rules << @default_rule
-      all_rules << @custom_rules unless @custom_rules.empty?
-      all_rules
+    def daemon_config=(v)
+      @poller.connector.daemon_config = v
     end
 
     private
 
-    def should_sample?(rule)
-      return true if rule.reservoir.take
-      Random.rand <= rule.rate
-    end
-
-    def load_sampling_rules(v)
-      version = v[:version]
-      if version != 1
-        raise InvalidSamplingConfigError, %('Sampling rule version #{version} is not supported.')
-      end
-      unless v[:default]
-        raise InvalidSamplingConfigError, 'A default rule must be provided.'
-      end
-      @default_rule = SamplingRule.new rule_definition: v[:default], default: true
-      @custom_rules = []
-      v[:rules].each do |d|
-        @custom_rules << SamplingRule.new(rule_definition: d)
+    def process_matched_rule(rule, now)
+      # As long as a rule is matched we increment request counter.
+      rule.increment_request_count
+      reservoir = rule.reservoir
+      sample = true
+      # We check if we can borrow or take from reservoir first.
+      decision = reservoir.borrow_or_take(now, rule.borrowable?)
+      if decision == SamplingDecision::BORROW
+        rule.increment_borrow_count
+      elsif decision == SamplingDecision::TAKE
+        rule.increment_sampled_count
+        # Otherwise we compute based on fixed rate of this sampling rule.
+      elsif rand <= rule.rate
+        rule.increment_sampled_count
+      else
+        sample = false
       end
+      sample ? rule.name : false
     end
   end
 end

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

@@ -0,0 +1,72 @@
+require 'aws-xray-sdk/logger'
+require 'aws-xray-sdk/sampling/connector'
+require 'aws-xray-sdk/sampling/rule_poller'
+
+module XRay
+  # The poller to report the current statistics of all
+  # sampling rules and retrieve the new allocated
+  # sampling quota and TTL from X-Ray service. It also
+  # controls the rule poller.
+  class LeadPoller
+    include Logging
+    attr_reader :connector
+    @@interval = 10 # working frequency of the lead poller
+    @@rule_interval = 5 * 60 # 5 minutes on polling rules
+
+    def initialize(cache)
+      @cache = cache
+      @connector = ServiceConnector.new
+      @rule_poller = RulePoller.new cache: @cache, connector: @connector
+      @rule_poller_elapsed = 0
+    end
+
+    def start
+      @rule_poller.run
+      Thread.new { worker }
+    end
+
+    def worker
+      loop do
+        sleep_time = @@interval + rand
+        sleep sleep_time
+        @rule_poller_elapsed += sleep_time
+        refresh_cache
+        if @rule_poller_elapsed >= @@rule_interval
+          @rule_poller.run
+          @rule_poller_elapsed = 0
+        end
+      end
+    end
+
+    private
+
+    def refresh_cache
+      candidates = get_candidates(@cache.rules)
+      if candidates.empty?
+        logger.debug %(No X-Ray sampling rules to report statistics. Skipping.)
+        return
+      end
+
+      result = @connector.fetch_sampling_targets(candidates)
+      targets = {}
+      result[:documents].each { |doc| targets[doc.rule_name] = doc }
+      @cache.load_targets(targets)
+
+      return unless @cache.last_updated && result[:last_modified].to_i > @cache.last_updated
+      logger.info 'Performing out-of-band sampling rule polling to fetch updated rules.'
+      @rule_poller.run
+      @rule_poller_elapsed = 0
+    rescue StandardError => e
+      logger.warn %(failed to fetch X-Ray sampling targets due to #{e.message})
+    end
+
+    # Don't report a rule statistics if any of the conditions is met:
+    # 1. The report time hasn't come(some rules might have larger report intervals).
+    # 2. The rule is never matched.
+    def get_candidates(rules)
+      candidates = []
+      rules.each { |rule| candidates << rule if rule.ever_matched? && rule.time_to_report? }
+      candidates
+    end
+  end
+end

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

@@ -0,0 +1,35 @@
+module XRay
+  # Keeps track of the number of sampled segments within
+  # a single second in the local process. This class is
+  # implemented to be thread-safe to achieve accurate sampling.
+  class LocalReservoir
+    # @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
+      @lock = Mutex.new
+    end
+
+    # Returns `true` if there are quota 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?
+      @lock.synchronize do
+        now = Time.now.to_i
+        # refresh time frame
+        if now != @this_sec
+          @used_this_sec = 0
+          @this_sec = now
+        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
+      end
+    end
+  end
+end

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

@@ -0,0 +1,110 @@
+require 'aws-xray-sdk/sampling/sampler'
+require 'aws-xray-sdk/sampling/local/sampling_rule'
+require 'aws-xray-sdk/exceptions'
+
+module XRay
+  # The local sampler that uses locally defined
+  # sampling rule and reservoir models to decide sampling decision.
+  # It also uses the default sampling rule.
+  # An example definition:
+  #   {
+  #     version: 2,
+  #     rules: [
+  #       {
+  #         description: 'Player moves.',
+  #         host: '*',
+  #         http_method: '*',
+  #         url_path: '/api/move/*',
+  #         fixed_target: 0,
+  #         rate: 0.05
+  #       }
+  #     ],
+  #     default: {
+  #       fixed_target: 1,
+  #       rate: 0.1
+  #     }
+  #   }
+  # This example defines one custom rule and a default rule.
+  # The custom rule applies a five-percent sampling rate with no minimum
+  # number of requests to trace for paths under /api/move/. The default
+  # rule traces the first request each second and 10 percent of additional requests.
+  # The SDK applies custom rules in the order in which they are defined.
+  # If a request matches multiple custom rules, the SDK applies only the first rule.
+  class LocalSampler
+    include Sampler
+    DEFAULT_RULES = {
+      version: 2,
+      default: {
+        fixed_target: 1,
+        rate: 0.05
+      },
+      rules: []
+    }.freeze
+
+    SUPPORTED_VERSION = [1, 2].freeze
+
+    def initialize
+      load_sampling_rules(DEFAULT_RULES)
+    end
+
+    # Return True if the sampler decide to sample based on input
+    # information and sampling rules. It will first check if any
+    # custom rule should be applied, if not it falls back to the
+    # default sampling rule.
+    # All arugments are extracted from incoming requests by
+    # X-Ray middleware to perform path based sampling.
+    def sample_request?(sampling_req)
+      sample = sample?
+      return sample if sampling_req.nil? || sampling_req.empty?
+      @custom_rules ||= []
+      @custom_rules.each do |c|
+        return should_sample?(c) if c.applies?(sampling_req: sampling_req)
+      end
+      # use previously made decision based on default rule
+      # if no path-based rule has been matched
+      sample
+    end
+
+    # Decides if should sample based on non-path-based rule.
+    # Currently only the default rule is non-path-based.
+    def sample?
+      should_sample?(@default_rule)
+    end
+
+    # @param [Hash] v The sampling rules definition.
+    def sampling_rules=(v)
+      load_sampling_rules(v)
+    end
+
+    # @return [Array] An array of [SamplingRule]
+    def sampling_rules
+      all_rules = []
+      all_rules << @default_rule
+      all_rules << @custom_rules unless @custom_rules.empty?
+      all_rules
+    end
+
+    private
+
+    def should_sample?(rule)
+      return true if rule.reservoir.take
+      Random.rand <= rule.rate
+    end
+
+    def load_sampling_rules(v)
+      version = v[:version]
+      unless SUPPORTED_VERSION.include?(version)
+        raise InvalidSamplingConfigError, %('Sampling rule version #{version} is not supported.')
+      end
+      unless v[:default]
+        raise InvalidSamplingConfigError, 'A default rule must be provided.'
+      end
+      @default_rule = LocalSamplingRule.new rule_definition: v[:default], default: true
+      @custom_rules = []
+      v[:rules].each do |d|
+        d[:host] = d[:service_name] if version == 1
+        @custom_rules << LocalSamplingRule.new(rule_definition: d)
+      end
+    end
+  end
+end