solarwinds_apm 7.0.2 → 7.1.1

data/lib/solarwinds_apm/sampling/oboe_sampler.rb

@@ -18,7 +18,6 @@ module SolarWindsAPM
     TRIGGERED_TRACE_ATTRIBUTE = 'TriggeredTrace'
 
     TRACESTATE_REGEXP = /^[0-9a-f]{16}-[0-9a-f]{2}$/
-    BUCKET_INTERVAL = 1000
     DICE_SCALE = 1_000_000
 
     OTEL_SAMPLING_DECISION = ::OpenTelemetry::SDK::Trace::Samplers::Decision
@@ -30,34 +29,36 @@ module SolarWindsAPM
       @counters = SolarWindsAPM::Metrics::Counter.new
       @buckets = {
         SolarWindsAPM::BucketType::DEFAULT =>
-          SolarWindsAPM::TokenBucket.new(SolarWindsAPM::TokenBucketSettings.new(nil, nil, BUCKET_INTERVAL)),
+          SolarWindsAPM::TokenBucket.new(SolarWindsAPM::TokenBucketSettings.new(nil, nil, 'DEFUALT')),
         SolarWindsAPM::BucketType::TRIGGER_RELAXED =>
-          SolarWindsAPM::TokenBucket.new(SolarWindsAPM::TokenBucketSettings.new(nil, nil, BUCKET_INTERVAL)),
+          SolarWindsAPM::TokenBucket.new(SolarWindsAPM::TokenBucketSettings.new(nil, nil, 'TRIGGER_RELAXED')),
         SolarWindsAPM::BucketType::TRIGGER_STRICT =>
-          SolarWindsAPM::TokenBucket.new(SolarWindsAPM::TokenBucketSettings.new(nil, nil, BUCKET_INTERVAL))
+          SolarWindsAPM::TokenBucket.new(SolarWindsAPM::TokenBucketSettings.new(nil, nil, 'TRIGGER_STRICT'))
       }
       @settings = {} # parsed setting from swo backend
-
-      @buckets.each_value(&:start)
+      @settings_mutex = ::Mutex.new
     end
 
     # return sampling result
     # params: {:trace_id=>, :parent_context=>, :links=>, :name=>, :kind=>, :attributes=>}
     # propagator -> processor -> sampler
     def should_sample?(params)
-      @logger.debug { "should_sample? params: #{params.inspect}" }
       _, parent_context, _, _, _, attributes = params.values
 
       parent_span = ::OpenTelemetry::Trace.current_span(parent_context)
       type = SolarWindsAPM::SpanType.span_type(parent_span)
 
-      @logger.debug { "[#{self.class}/#{__method__}] span type is #{type}" }
+      @logger.debug { "[#{self.class}/#{__method__}] should_sample? params: #{params.inspect}; span type is #{type}" }
 
       # For local spans, we always trust the parent
       if type == SolarWindsAPM::SpanType::LOCAL
-        return OTEL_SAMPLING_RESULT.new(decision: OTEL_SAMPLING_DECISION::RECORD_AND_SAMPLE, tracestate: DEFAULT_TRACESTATE) if parent_span.context.trace_flags.sampled?
-
-        return OTEL_SAMPLING_RESULT.new(decision: OTEL_SAMPLING_DECISION::DROP, tracestate: DEFAULT_TRACESTATE)
+        if parent_span.context.trace_flags.sampled?
+          return OTEL_SAMPLING_RESULT.new(decision: OTEL_SAMPLING_DECISION::RECORD_AND_SAMPLE,
+                                          tracestate: DEFAULT_TRACESTATE)
+        else
+          return OTEL_SAMPLING_RESULT.new(decision: OTEL_SAMPLING_DECISION::DROP,
+                                          tracestate: DEFAULT_TRACESTATE)
+        end
       end
 
       sample_state = SampleState.new(OTEL_SAMPLING_DECISION::DROP,
@@ -78,10 +79,9 @@ module SolarWindsAPM
         # TraceOptions.parse_trace_options return TriggerTraceOptions
         sample_state.trace_options = ::SolarWindsAPM::TraceOptions.parse_trace_options(sample_state.headers['X-Trace-Options'], @logger)
 
-        @logger.debug { "X-Trace-Options present: #{sample_state.trace_options}" }
+        @logger.debug { "[#{self.class}/#{__method__}] sample_state.trace_options: #{sample_state.trace_options.inspect}" }
 
         if sample_state.headers['X-Trace-Options-Signature']
-          @logger.debug { 'X-Trace-Options-Signature present; validating' }
 
           # this validate_signature is the function from trace_options file
           sample_state.trace_options.response.auth = TraceOptions.validate_signature(
@@ -92,36 +92,27 @@ module SolarWindsAPM
           )
 
           # If the request has an invalid signature, drop the trace
-          if sample_state.trace_options.response.auth != Auth::OK # Auth::OK is a string from trace_options.rb: 'ok'
-            @logger.debug { 'X-Trace-Options-Signature invalid; tracing disabled' }
+          if sample_state.trace_options.response.auth != Auth::OK
+            @logger.debug { "[#{self.class}/#{__method__}] signature invalid; tracing disabled (auth=#{sample_state.trace_options.response.auth})" }
 
             xtracestate = generate_new_tracestate(parent_span, sample_state)
-            return OTEL_SAMPLING_RESULT.new(decision: OTEL_SAMPLING_DECISION::DROP, tracestate: xtracestate, attributes: sample_state.attributes)
+            return OTEL_SAMPLING_RESULT.new(decision: OTEL_SAMPLING_DECISION::DROP,
+                                            tracestate: xtracestate,
+                                            attributes: sample_state.attributes)
           end
         end
 
-        unless sample_state.trace_options.trigger_trace
-          sample_state.trace_options.response.trigger_trace = TriggerTrace::NOT_REQUESTED # 'not-requested'
-        end
-
-        # Apply trace options to span attributes
+        # Apply trace options to span attributes and list ignored keys in response
+        sample_state.trace_options.response.trigger_trace = TriggerTrace::NOT_REQUESTED unless sample_state.trace_options.trigger_trace
         sample_state.attributes[SW_KEYS_ATTRIBUTE] = sample_state.trace_options[:sw_keys] if sample_state.trace_options[:sw_keys]
-
-        sample_state.trace_options.custom.each do |k, v|
-          sample_state.attributes[k] = v
-        end
-
-        # List ignored keys in response
+        sample_state.trace_options.custom.each { |k, v| sample_state.attributes[k] = v }
         sample_state.trace_options.response.ignored = sample_state.trace_options[:ignored].map { |k, _| k } if sample_state.trace_options[:ignored].any?
       end
 
       unless sample_state.settings
-        @logger.debug { 'settings unavailable; sampling disabled' }
+        @logger.debug { "[#{self.class}/#{__method__}] settings unavailable; sampling disabled" }
 
-        if sample_state.trace_options&.trigger_trace
-          @logger.debug { 'trigger trace requested but unavailable' }
-          sample_state.trace_options.response.trigger_trace = TriggerTrace::SETTINGS_NOT_AVAILABLE # 'settings-not-available'
-        end
+        sample_state.trace_options.response.trigger_trace = TriggerTrace::SETTINGS_NOT_AVAILABLE if sample_state.trace_options&.trigger_trace
 
         xtracestate = generate_new_tracestate(parent_span, sample_state)
 
@@ -131,31 +122,23 @@ module SolarWindsAPM
       end
 
       @logger.debug { "[#{self.class}/#{__method__}] sample_state before deciding sampling algo: #{sample_state.inspect}" }
+
       # Decide which sampling algo to use and add sampling attribute to decision attributes
       # https://swicloud.atlassian.net/wiki/spaces/NIT/pages/3815473156/Tracing+Decision+Tree
       if sample_state.trace_state && TRACESTATE_REGEXP.match?(sample_state.trace_state)
-        @logger.debug { 'context is valid for parent-based sampling' }
         parent_based_algo(sample_state)
-
       elsif sample_state.settings[:flags].anybits?(Flags::SAMPLE_START)
         if sample_state.trace_options&.trigger_trace
-          @logger.debug { 'trigger trace requested' }
           trigger_trace_algo(sample_state)
         else
-          @logger.debug { 'defaulting to dice roll' }
           dice_roll_algo(sample_state)
         end
       else
-        @logger.debug { 'SAMPLE_START is unset; sampling disabled' }
         disabled_algo(sample_state)
       end
 
-      @logger.debug { "final sampling state: #{sample_state.inspect}" }
-
       xtracestate = generate_new_tracestate(parent_span, sample_state)
-
-      # if need to set 'sw.w3c.tracestate' to attributes
-      # sample_state.attributes['sw.w3c.tracestate'] = ::SolarWindsAPM::Utils.trace_state_header(xtracestate)
+      @logger.debug { "[#{self.class}/#{__method__}] final sampling state: #{sample_state.inspect}" }
 
       OTEL_SAMPLING_RESULT.new(decision: sample_state.decision,
                                tracestate: xtracestate,
@@ -163,13 +146,10 @@ module SolarWindsAPM
     end
 
     def parent_based_algo(sample_state)
-      # original js code: const [context] = s.params
-      # the context is used for metrics e.g. this.#counters.throughTraceCount.add(1, {}, context)
+      @logger.debug { "[#{self.class}/#{__method__}] parent_based_algo start" }
 
-      # compare the parent_id
       sample_state.attributes[PARENT_ID_ATTRIBUTE] = sample_state.trace_state[0, 16]
-
-      if sample_state.trace_options&.trigger_trace # need to implement trace_options
+      if sample_state.trace_options&.trigger_trace
         @logger.debug { 'trigger trace requested but ignored' }
         sample_state.trace_options.response.trigger_trace = TriggerTrace::IGNORED # 'ignored'
       end
@@ -198,18 +178,22 @@ module SolarWindsAPM
           @logger.debug { 'parent is sampled; record and sample' }
 
           @counters[:trace_count].add(1)
-          @counters[:through_trace_count].add(1) # ruby metrics only add incremented value and attributes
+          @counters[:through_trace_count].add(1)
 
           sample_state.decision = OTEL_SAMPLING_DECISION::RECORD_AND_SAMPLE
         end
       end
+
+      @logger.debug { "[#{self.class}/#{__method__}] parent_based_algo end" }
     end
 
     def trigger_trace_algo(sample_state)
+      @logger.debug { "[#{self.class}/#{__method__}] trigger_trace_algo start" }
+
       if sample_state.settings[:flags].nobits?(Flags::TRIGGERED_TRACE)
         @logger.debug { 'TRIGGERED_TRACE unset; record only' }
 
-        sample_state.trace_options.response.trigger_trace = TriggerTrace::TRIGGER_TRACING_DISABLED # 'trigger-tracing-disabled'
+        sample_state.trace_options.response.trigger_trace = TriggerTrace::TRIGGER_TRACING_DISABLED
         sample_state.decision = OTEL_SAMPLING_DECISION::RECORD_ONLY
       else
         @logger.debug { 'TRIGGERED_TRACE set; trigger tracing' }
@@ -244,9 +228,12 @@ module SolarWindsAPM
           sample_state.decision = OTEL_SAMPLING_DECISION::RECORD_ONLY
         end
       end
+      @logger.debug { "[#{self.class}/#{__method__}] trigger_trace_algo end" }
     end
 
     def dice_roll_algo(sample_state)
+      @logger.debug { "[#{self.class}/#{__method__}] dice_roll_algo start" }
+
       dice = SolarWindsAPM::Dice.new(rate: sample_state.settings[:sample_rate], scale: DICE_SCALE)
       sample_state.attributes[SAMPLE_RATE_ATTRIBUTE] = dice.rate
       sample_state.attributes[SAMPLE_SOURCE_ATTRIBUTE] = sample_state.settings[:sample_source]
@@ -278,9 +265,11 @@ module SolarWindsAPM
         @logger.debug { 'dice roll failure; record only' }
         sample_state.decision = OTEL_SAMPLING_DECISION::RECORD_ONLY
       end
+      @logger.debug { "[#{self.class}/#{__method__}] dice_roll_algo end" }
     end
 
     def disabled_algo(sample_state)
+      @logger.debug { "[#{self.class}/#{__method__}] disabled_algo start" }
       if sample_state.trace_options&.trigger_trace
         @logger.debug { 'trigger trace requested but tracing disabled' }
         sample_state.trace_options.response.trigger_trace = TriggerTrace::TRACING_DISABLED
@@ -293,14 +282,17 @@ module SolarWindsAPM
         @logger.debug { 'SAMPLE_THROUGH_ALWAYS is set; record' }
         sample_state.decision = OTEL_SAMPLING_DECISION::RECORD_ONLY
       end
+      @logger.debug { "[#{self.class}/#{__method__}] disabled_algo end" }
     end
 
     def update_settings(settings)
-      return unless settings[:timestamp] > (@settings[:timestamp] || 0)
+      @settings_mutex.synchronize do
+        return unless settings[:timestamp] > (@settings[:timestamp] || 0)
 
-      @settings = settings
-      @buckets.each do |type, bucket|
-        bucket.update(@settings[:buckets][type]) if @settings[:buckets][type]
+        @settings = settings
+        @buckets.each do |type, bucket|
+          bucket.update(@settings[:buckets][type]) if @settings[:buckets][type]
+        end
       end
     end
 
@@ -308,41 +300,41 @@ module SolarWindsAPM
     # handle_response_headers functionality is replace by generate_new_tracestate
     def generate_new_tracestate(parent_span, sample_state)
       if !parent_span.context.valid? || parent_span.context.tracestate.nil?
-        @logger.debug { 'create new tracestate' }
+        action = 'create'
         decision = sw_from_span_and_decision(parent_span, sample_state.decision)
         trace_state = ::OpenTelemetry::Trace::Tracestate.from_hash({ 'sw' => decision })
       else
-        @logger.debug { 'update tracestate' }
+        action = 'update'
         decision = sw_from_span_and_decision(parent_span, sample_state.decision)
         trace_state = parent_span.context.tracestate.set_value('sw', decision)
       end
 
       stringified_trace_options = SolarWindsAPM::TraceOptions.stringify_trace_options_response(sample_state.trace_options&.response)
-      @logger.debug { "[#{self.class}/#{__method__}] stringified_trace_options: #{stringified_trace_options}" }
-
       trace_state = trace_state.set_value('xtrace_options_response', stringified_trace_options)
-      @logger.debug { "[#{self.class}/#{__method__}] new trace_state: #{trace_state.inspect}" }
+      @logger.debug { "[#{self.class}/#{__method__}] Tracestate #{action}: decision=#{decision[-2, 2]}, xtrace_resp=#{stringified_trace_options}, trace_state=#{trace_state.inspect}" }
       trace_state
     end
 
     def sw_from_span_and_decision(parent_span, otel_decision)
       trace_flag = otel_decision == OTEL_SAMPLING_DECISION::RECORD_AND_SAMPLE ? '01' : '00'
-      [parent_span.context.hex_span_id, trace_flag].join('-')
+      "#{parent_span.context.hex_span_id}-#{trace_flag}"
     end
 
     def get_settings(params)
-      return if @settings.empty?
-
-      expiry = (@settings[:timestamp] + @settings[:ttl]) * 1000
-      time_now = Time.now.to_i * 1000
-      if time_now > expiry
-        @logger.debug { 'settings expired, removing' }
-        @settings = {}
-        return
+      @settings_mutex.synchronize do
+        return if @settings.empty?
+
+        expiry = (@settings[:timestamp] + @settings[:ttl]) * 1000
+        time_now = Time.now.to_i * 1000
+        if time_now > expiry
+          @logger.debug { "[#{self.class}/#{__method__}] settings expired, removing" }
+          @settings = {}
+          return
+        end
+        sampling_setting = SolarWindsAPM::SamplingSettings.merge(@settings, local_settings(params))
+        @logger.debug { "[#{self.class}/#{__method__}] sampling_setting: #{sampling_setting.inspect}" }
+        sampling_setting
       end
-      sampling_setting = SolarWindsAPM::SamplingSettings.merge(@settings, local_settings(params))
-      @logger.debug { "sampling_setting: #{sampling_setting.inspect}" }
-      sampling_setting
     end
   end
 end

data/lib/solarwinds_apm/sampling/sampler.rb

@@ -21,6 +21,9 @@ module SolarWindsAPM
     ATTR_URL_PATH = 'url.path'
     ATTR_HTTP_TARGET = RUBY_SEM_CON::HTTP_TARGET
 
+    SW_XTRACEOPTIONS_KEY = 'sw_xtraceoptions'
+    SW_SIGNATURE_KEY = 'sw_signature'
+
     # tracing_mode is getting from SolarWindsAPM::Config
     def initialize(config, logger)
       super(logger)
@@ -28,25 +31,22 @@ module SolarWindsAPM
       @trigger_mode = config[:trigger_trace_enabled]
       @transaction_settings = config[:transaction_settings]
       @ready = false
+      @logger.debug { "[#{self.class}/#{__method__}] Sampler initialized: tracing_mode=#{@tracing_mode}, trigger_mode=#{@trigger_mode}, transaction_settings_count=#{@transaction_settings.inspect}" }
     end
 
-    # wait for getting the first settings
     def wait_until_ready(timeout = 10)
-      thread = Thread.new { settings_ready }
-      thread.join(timeout) || (thread.kill
-                               false)
-      @ready
-    end
-
-    def settings_ready(timeout = 10)
-      deadline = Time.now
-      loop do
-        break unless @settings.empty?
-
+      deadline = Time.now + timeout
+      while Time.now < deadline
+        # The @settings hash is populated by another thread (e.g., HttpSampler)
+        unless @settings.empty?
+          @ready = !@settings[:signature_key].nil?
+          return @ready
+        end
         sleep 0.1
-        break if (Time.now - deadline).round(0) >= timeout
       end
-      @ready = true unless @settings[:signature_key].nil?
+
+      @logger.warn { "[#{self.class}/#{__method__}] Timed out waiting for settings after #{timeout} seconds." }
+      @ready # Will be false if timeout is reached
     end
 
     def resolve_tracing_mode(config)
@@ -58,59 +58,45 @@ module SolarWindsAPM
     def local_settings(params)
       _trace_id, _parent_context, _links, span_name, span_kind, attributes = params.values
       settings = { tracing_mode: @tracing_mode, trigger_mode: @trigger_mode }
-      return settings if @transaction_settings.nil? || @transaction_settings.empty?
 
-      @logger.debug { "Current @transaction_settings: #{@transaction_settings.inspect}" }
-      http_metadata = http_span_metadata(span_kind, attributes)
-      @logger.debug { "http_metadata: #{http_metadata.inspect}" }
+      if @transaction_settings.nil? || @transaction_settings.empty?
+        @logger.debug { "[#{self.class}/#{__method__}] No transaction settings, using defaults settings: #{settings.inspect}" }
+      else
+        http_metadata = http_span_metadata(span_kind, attributes)
+        # below is for filter out unwanted transaction
+        trans_settings = ::SolarWindsAPM::TransactionSettings.new(url_path: http_metadata[:url], name: span_name, kind: span_kind)
+        tracing_mode   = trans_settings.calculate_trace_mode == 1 ? TracingMode::ALWAYS : TracingMode::NEVER
 
-      # below is for filter out unwanted transaction
-      trans_settings = ::SolarWindsAPM::TransactionSettings.new(url_path: http_metadata[:url], name: span_name, kind: span_kind)
-      tracing_mode   = trans_settings.calculate_trace_mode == 1 ? TracingMode::ALWAYS : TracingMode::NEVER
+        settings[:tracing_mode] = tracing_mode
+      end
 
-      settings[:tracing_mode] = tracing_mode
+      @logger.debug { "[#{self.class}/#{__method__}] Transaction settings after calculation #{settings.inspect}" }
       settings
     end
 
     # if context have sw-related value, it should be stored in context
-    # named sw_xtraceoptions in header propagator
-    # original x_trace_options will parse headers in the class, apm-js separate the task
-    # apm-js will make headers as hash
+    # named sw_xtraceoptions and sw_signature in header from propagator
     def request_headers(params)
-      parent_context = params[:parent_context]
-      header = obtain_sw_value(parent_context, 'sw_xtraceoptions')
-      signature = obtain_sw_value(parent_context, 'sw_signature')
-      @logger.debug { "[#{self.class}/#{__method__}] trace_options option_header: #{header}; trace_options sw_signature: #{signature}" }
-
-      {
-        'X-Trace-Options' => header,
-        'X-Trace-Options-Signature' => signature
-      }
+      header, signature = obtain_traceoptions_headers_signature(params[:parent_context])
+      @logger.debug { "[#{self.class}/#{__method__}] trace_options header: #{header.inspect}, signature: #{signature.inspect} from parent_context: #{params[:parent_context].inspect}" }
+      { 'X-Trace-Options' => header, 'X-Trace-Options-Signature' => signature }
     end
 
-    def obtain_sw_value(context, type)
-      sw_value = nil
-      instance_variable = context&.instance_variable_get('@entries')
-      instance_variable&.each do |key, value|
-        next unless key.instance_of?(::String)
-
-        sw_value = value if key == type
-      end
-      sw_value
+    def obtain_traceoptions_headers_signature(context)
+      header = context.value(SW_XTRACEOPTIONS_KEY)
+      signature = context.value(SW_SIGNATURE_KEY)
+      [header, signature]
     end
 
     def update_settings(settings)
       parsed = parse_settings(settings)
       if parsed
-        @logger.debug { "valid settings #{parsed.inspect} from setting #{settings.inspect}" }
-
-        super(parsed) # call oboe_sampler update_settings function to update the buckets
-
+        @logger.debug { "[#{self.class}/#{__method__}] Valid settings #{parsed.inspect} from setting #{settings.inspect}" }
         @logger.warn { "Warning from parsed settings: #{parsed[:warning]}" } if parsed[:warning]
-
+        super(parsed) # call oboe_sampler update_settings function to update the buckets
         parsed
       else
-        @logger.debug { "invalid settings: #{settings.inspect}" }
+        @logger.debug { "[#{self.class}/#{__method__}] Invalid settings: #{settings.inspect}" }
         nil
       end
     end
@@ -136,24 +122,25 @@ module SolarWindsAPM
         url: url
       }
 
-      @logger.debug { "Retrieved http metadata: #{http_metadata.inspect}" }
+      @logger.debug { "[#{self.class}/#{__method__}] Retrieved http metadata: #{http_metadata.inspect}" }
       http_metadata
     end
 
     def parse_settings(unparsed)
       return unless unparsed.is_a?(Hash)
 
-      return unless unparsed['value'].is_a?(Numeric) &&
-                    unparsed['timestamp'].is_a?(Numeric) &&
-                    unparsed['ttl'].is_a?(Numeric)
-
       sample_rate = unparsed['value']
-      timestamp = unparsed['timestamp']
-      ttl = unparsed['ttl']
+      timestamp   = unparsed['timestamp']
+      ttl         = unparsed['ttl']
+      flags       = unparsed['flags']
+
+      return unless sample_rate.is_a?(Numeric) &&
+                    timestamp.is_a?(Numeric) &&
+                    ttl.is_a?(Numeric)
 
-      return unless unparsed['flags'].is_a?(String)
+      return unless flags.is_a?(String)
 
-      flags = unparsed['flags'].split(',').reduce(Flags::OK) do |final_flag, f|
+      flags = flags.split(',').reduce(Flags::OK) do |final_flag, f|
         flag = {
           'OVERRIDE' => Flags::OVERRIDE,
           'SAMPLE_START' => Flags::SAMPLE_START,
@@ -167,6 +154,7 @@ module SolarWindsAPM
 
       buckets = {}
       signature_key = nil
+      warning = nil
 
       if unparsed['arguments'].is_a?(Hash)
         args = unparsed['arguments']

data/lib/solarwinds_apm/sampling/sampling_constants.rb

@@ -55,7 +55,7 @@ module SolarWindsAPM
 
   TokenBucketSettings = Struct.new(:capacity,    # Number
                                    :rate,        # Number
-                                   :interval) # Number
+                                   :type)        # String
 
   module SampleSource
     LOCAL_DEFAULT = 2
@@ -105,11 +105,11 @@ module SolarWindsAPM
     end
 
     def self.valid_trace_id?(trace_id)
-      !!(trace_id =~ VALID_TRACEID_REGEX) && trace_id != INVALID_TRACEID
+      VALID_TRACEID_REGEX.match?(trace_id) && trace_id != INVALID_TRACEID
     end
 
     def self.valid_span_id?(span_id)
-      !!(span_id =~ VALID_SPANID_REGEX) && span_id != INVALID_SPANID
+      VALID_SPANID_REGEX.match?(span_id) && span_id != INVALID_SPANID
     end
 
     def self.span_context_valid?(span_context)

data/lib/solarwinds_apm/sampling/settings.rb

@@ -9,6 +9,7 @@
 module SolarWindsAPM
   module SamplingSettings
     def self.merge(remote, local)
+      SolarWindsAPM.logger.debug { "[#{self.class}/#{__method__}] SamplingSettings merge with remote: #{remote.inspect}; local: #{local.inspect}" }
       flags = local[:tracing_mode] || remote[:flags]
 
       if local[:trigger_mode] == :enabled
@@ -22,6 +23,7 @@ module SolarWindsAPM
         flags |= SolarWindsAPM::Flags::OVERRIDE
       end
 
+      SolarWindsAPM.logger.debug { "[#{self.class}/#{__method__}] final flags: #{flags}" }
       remote.merge(flags: flags)
     end
   end

data/lib/solarwinds_apm/sampling/token_bucket.rb

@@ -10,106 +10,88 @@
 # capacity is updated through update_settings
 module SolarWindsAPM
   class TokenBucket
-    # Maximum value of a signed 32-bit integer
-    MAX_INTERVAL = (2**31) - 1
-
-    attr_reader :capacity, :rate, :interval, :tokens
+    attr_reader :type
 
     def initialize(token_bucket_settings)
-      self.capacity = token_bucket_settings.capacity || 0
-      self.rate = token_bucket_settings.rate || 0
-      self.interval = token_bucket_settings.interval || MAX_INTERVAL
-      self.tokens = @capacity
-      @timer = nil
+      @capacity = token_bucket_settings.capacity || 0
+      @rate = token_bucket_settings.rate || 0
+      @tokens = @capacity
+      @last_update_time = Time.now.to_f
+      @type = token_bucket_settings.type
+      @lock = ::Mutex.new
     end
 
-    # oboe sampler update_settings will update the token
-    # (thread safe as update_settings is guarded by mutex from oboe sampler)
-    def update(settings)
-      settings.instance_of?(Hash) ? update_from_hash(settings) : update_from_hash(tb_to_hash(settings))
+    def capacity
+      @lock.synchronize { @capacity }
     end
 
-    def update_from_hash(settings)
-      if settings[:capacity]
-        difference = settings[:capacity] - @capacity
-        self.capacity = settings[:capacity]
-        self.tokens = @tokens + difference
-      end
-
-      self.rate = settings[:rate] if settings[:rate]
-      start
+    def rate
+      @lock.synchronize { @rate }
     end
 
-    def tb_to_hash(settings)
-      { capacity: settings.capacity,
-        rate: settings.rate,
-        interval: settings.interval }
-    end
-
-    def capacity=(capacity)
-      @capacity = [0, capacity].max
-    end
-
-    def rate=(rate)
-      @rate = [0, rate].max
-    end
-
-    # self.interval= sets the @interval and @sleep_interval
-    # @sleep_interval is used in the timer thread to sleep between replenishing the bucket
-    def interval=(interval)
-      @interval = interval.clamp(0, MAX_INTERVAL)
-      @sleep_interval = @interval / 1000.0
+    def tokens
+      @lock.synchronize do
+        calculate_tokens
+        @tokens
+      end
     end
 
-    def tokens=(tokens)
-      @tokens = tokens.clamp(0, @capacity)
+    # oboe sampler update_settings will update the token
+    def update(settings)
+      settings.instance_of?(Hash) ? update_from_hash(settings) : update_from_hash(tb_to_hash(settings))
     end
 
     # Attempts to consume tokens from the bucket
-    # @param n [Integer] Number of tokens to consume
+    # @param token [Integer] Number of tokens to consume
     # @return [Boolean] Whether there were enough tokens
-    # TODO: we need to include thread-safety here since sampler is shared across threads
-    # and we may have multiple threads trying to consume tokens at the same time
     def consume(token = 1)
-      if @tokens >= token
-        self.tokens = @tokens - token
-        SolarWindsAPM.logger.debug { "[#{self.class}/#{__method__}] Consumed #{token} from total #{@tokens} (#{(@tokens.to_f / @capacity * 100).round(1)}% remaining)" }
-        true
-      else
-        SolarWindsAPM.logger.debug { "[#{self.class}/#{__method__}] Token consumption failed: requested=#{token}, available=#{@tokens}, capacity=#{@capacity}" }
-        false
-      end
-    end
-
-    # Starts replenishing the bucket
-    def start
-      return if running
-
-      @timer = Thread.new do
-        loop do
-          task
-          sleep(@sleep_interval)
+      @lock.synchronize do
+        calculate_tokens
+        if @tokens >= token
+          @tokens -= token
+          SolarWindsAPM.logger.debug { "[#{self.class}/#{__method__}] #{@type} Consumed #{token} (#{(@tokens.to_f / @capacity * 100).round(1)}% remaining)" }
+          true
+        else
+          SolarWindsAPM.logger.debug { "[#{self.class}/#{__method__}] #{@type} Token consumption failed: requested=#{token}, available=#{@tokens}, capacity=#{@capacity}" }
+          false
         end
       end
     end
 
-    # Stops replenishing the bucket
-    def stop
-      return unless running
+    private
 
-      @timer.kill
-      @timer = nil
+    def calculate_tokens
+      now = Time.now.to_f
+      elapsed = now - @last_update_time
+      @last_update_time = now
+      @tokens += elapsed * @rate
+      @tokens = [@tokens, @capacity].min
     end
 
-    # Whether the bucket is actively being replenished
-    def running
-      !@timer.nil? && @timer.alive?
-    end
+    # settings is from json sampler
+    def update_from_hash(settings)
+      @lock.synchronize do
+        calculate_tokens
+
+        if settings[:capacity]
+          new_capacity = [0, settings[:capacity]].max
+          difference = new_capacity - @capacity
+          @capacity = new_capacity
+          @tokens += difference
+          @tokens = [0, @tokens].max
+        end
 
-    private
+        @rate = [0, settings[:rate]].max if settings[:rate]
+      end
+    end
 
-    def task
-      self.tokens = tokens + @rate
+    # settings is from http sampler
+    def tb_to_hash(settings)
+      tb_hash = {}
+      tb_hash[:capacity] = settings.capacity if settings.respond_to?(:capacity)
+      tb_hash[:rate] = settings.rate if settings.respond_to?(:rate)
+      tb_hash[:type] = settings.type if settings.respond_to?(:type)
+      tb_hash
     end
   end
 end