diagnostics sampling

Author: kenny-statsig

.github/workflows/build-and-test.yml

@@ -16,7 +16,7 @@ jobs:
     strategy:
       max-parallel: 1
       matrix:
-        version: [2.5.0, 3.2.0]
+        version: [3.2.0, 2.5.0]
     steps:
       - uses: actions/checkout@v2
       - name: Set up Ruby

lib/diagnostics.rb

@@ -7,14 +7,18 @@ class Diagnostics
     extend T::Sig
 
     sig { returns(String) }
-    attr_reader :context
+    attr_accessor :context
 
     sig { returns(T::Array[T::Hash[Symbol, T.untyped]]) }
     attr_reader :markers
 
+    sig { returns(T::Hash[String, Numeric]) }
+    attr_accessor :sample_rates
+
     def initialize(context)
       @context = context
       @markers = []
+      @sample_rates = {}
     end
 
     sig do
@@ -65,12 +69,27 @@ def serialize
       }
     end
 
+    def serialize_with_sampling
+      marker_keys = @markers.map { |e| e[:key] }
+      unique_marker_keys = marker_keys.uniq { |e| e }
+      sampled_marker_keys = unique_marker_keys.select do |key|
+        @sample_rates.key?(key) && !self.class.sample(@sample_rates[key])
+      end
+      final_markers = @markers.select do |marker|
+        !sampled_marker_keys.include?(marker[:key])
+      end
+      {
+        context: @context.clone,
+        markers: final_markers
+      }
+    end
+
     def clear_markers
       @markers.clear
     end
 
-    def self.sample(rate)
-      rand(rate).zero?
+    def self.sample(rate_over_ten_thousand)
+      rand * 10_000 < rate_over_ten_thousand
     end
 
     class Context

lib/spec_store.rb

@@ -149,7 +149,7 @@ def maybe_restart_background_threads
     end
 
     def sync_config_specs
-      @diagnostics = Diagnostics.new('config_sync')
+      @diagnostics.context = 'config_sync'
       if @options.data_store&.should_be_used_for_querying_updates(Interfaces::IDataStore::CONFIG_SPECS_KEY)
         load_config_specs_from_storage_adapter
       else
@@ -159,7 +159,7 @@ def sync_config_specs
     end
 
     def sync_id_lists
-      @diagnostics = Diagnostics.new('config_sync')
+      @diagnostics.context = 'config_sync'
       if @options.data_store&.should_be_used_for_querying_updates(Interfaces::IDataStore::ID_LISTS_KEY)
         get_id_lists_from_adapter
       else
@@ -280,6 +280,7 @@ def process_specs(specs_string, from_adapter: false)
       specs_json['feature_gates'].each { |gate| new_gates[gate['name']] = gate }
       specs_json['dynamic_configs'].each { |config| new_configs[config['name']] = config }
       specs_json['layer_configs'].each { |layer| new_layers[layer['name']] = layer }
+      specs_json['diagnostics']&.each { |key, value| @diagnostics.sample_rates[key] = value }
 
       if specs_json['layers'].is_a?(Hash)
         specs_json['layers'].each { |layer_name, experiments|
@@ -293,8 +294,6 @@ def process_specs(specs_string, from_adapter: false)
       @specs[:experiment_to_layer] = new_exp_to_layer
       @specs[:sdk_keys_to_app_ids] = specs_json['sdk_keys_to_app_ids'] || {}
 
-      specs_json['diagnostics']
-
       unless from_adapter
         save_config_specs_to_storage_adapter(specs_string)
       end

lib/statsig_driver.rb

@@ -243,7 +243,7 @@ def maybe_restart_background_threads
 
   def run_with_diagnostics(task:, caller:)
     diagnostics = nil
-    if Statsig::Diagnostics::API_CALL_KEYS.include?(caller) && Statsig::Diagnostics.sample(10_000)
+    if Statsig::Diagnostics::API_CALL_KEYS.include?(caller) && Statsig::Diagnostics.sample(1)
       diagnostics = Statsig::Diagnostics.new('api_call')
       tracker = diagnostics.track(caller)
     end

lib/statsig_logger.rb

@@ -100,11 +100,14 @@ def log_layer_exposure(user, layer, parameter_name, config_evaluation, context =
 
     def log_diagnostics_event(diagnostics, user = nil)
       return if @options.disable_diagnostics_logging
-      return if diagnostics.nil? || diagnostics.markers.empty?
+      return if diagnostics.nil?
 
       event = StatsigEvent.new($diagnostics_event)
       event.user = user
-      event.metadata = diagnostics.serialize
+      serialized = diagnostics.serialize_with_sampling
+      return if serialized[:markers].empty?
+
+      event.metadata = serialized
       log_event(event)
       diagnostics.clear_markers
     end

test/test_diagnostics.rb

@@ -32,14 +32,10 @@ def setup
       @events.push(*JSON.parse(request.body)['events'])
       return ''
     })
-    Spy.on(Statsig::Diagnostics, 'sample').and_return do
-      true
-    end
   end
 
   def teardown
     super
-    Spy.off(Statsig::Diagnostics, 'sample')
     WebMock.reset!
     WebMock.disable!
   end
@@ -190,12 +186,16 @@ def test_data_adapter_init
 
   def test_api_call_diagnostics
     Statsig.initialize('secret-key')
+    Spy.on(Statsig::Diagnostics, 'sample').and_return do
+      true
+    end
     user = StatsigUser.new(user_id: 'test-user')
     Statsig.check_gate_with_exposure_logging_disabled(user, 'non-existent-gate')
     Statsig.get_config_with_exposure_logging_disabled(user, 'non-existent-config')
     Statsig.get_experiment_with_exposure_logging_disabled(user, 'non-existent-experiment')
     Statsig.get_layer_with_exposure_logging_disabled(user, 'non-existent-layer')
     Statsig.shutdown
+    Spy.off(Statsig::Diagnostics, 'sample')
 
     keys = Statsig::Diagnostics::API_CALL_KEYS
 
@@ -223,6 +223,72 @@ def test_disable_diagnostics_logging
     assert_equal(0, @events.length)
   end
 
+  def test_diagnostics_sampling
+    json_file = File.read("#{__dir__}/data/download_config_specs.json")
+    @mock_response = JSON.parse(json_file)
+    @mock_response['diagnostics'] = {
+      "download_config_specs": 5000,
+      "get_id_list": 5000,
+      "get_id_list_sources": 5000
+    }
+    stub_request(:post, 'https://statsigapi.net/v1/download_config_specs').to_return(
+      status: 200,
+      body: @mock_response.to_json,
+      headers: { 'x-statsig-region' => 'az-westus-2' }
+    )
+    driver = StatsigDriver.new(
+      'secret-key',
+      StatsigOptions.new(
+        disable_rulesets_sync: true,
+        disable_idlists_sync: true,
+        logging_interval_seconds: 9999
+      )
+    )
+    logger = driver.instance_variable_get('@logger')
+    logger.flush
+
+    assert_equal(1, @events.length)
+    event = @events[0]
+    assert_equal('statsig::diagnostics', event['eventName'])
+
+    metadata = event['metadata']
+    assert_equal('initialize', metadata['context'])
+    @events = []
+
+    10.times do
+      driver.manually_sync_rulesets
+    end
+    logger.flush
+
+    assert(
+      @events.length < 10 && @events.length.positive?,
+      "Expected between 0 and 10 events, received #{@events.length}"
+    )
+    event = @events[0]
+    assert_equal('statsig::diagnostics', event['eventName'])
+
+    metadata = event['metadata']
+    assert_equal('config_sync', metadata['context'])
+    @events = []
+
+    10.times do
+      driver.manually_sync_idlists
+    end
+    logger.flush
+
+    assert(
+      @events.length < 10 && @events.length.positive?,
+      "Expected between 0 and 10 events, received #{@events.length}"
+    )
+    event = @events[0]
+    assert_equal('statsig::diagnostics', event['eventName'])
+
+    metadata = event['metadata']
+    assert_equal('config_sync', metadata['context'])
+
+    driver.shutdown
+  end
+
   private
 
   def assert_marker_equal(marker, key, action, step = nil, tags = {})