Merge pull request #57 from figma/batch_update_test_duration

[CI] batch update test job duration

Author: lihao-figma

ruby/lib/ci/queue/configuration.rb

@@ -10,6 +10,7 @@ class Configuration
       attr_accessor :max_test_failed, :redis_ttl
       attr_accessor :strategy, :timing_file, :timing_fallback_duration, :export_timing_file
       attr_accessor :suite_max_duration, :suite_buffer_percent
+      attr_accessor :branch
       attr_reader :circuit_breakers
       attr_writer :seed, :build_id
       attr_writer :queue_init_timeout, :report_timeout, :inactive_workers_timeout
@@ -24,6 +25,7 @@ def from_env(env)
             statsd_endpoint: env['CI_QUEUE_STATSD_ADDR'],
             redis_ttl: env['CI_QUEUE_REDIS_TTL']&.to_i ||  8 * 60 * 60,
             known_flaky_tests: load_known_flaky_tests(env['CI_QUEUE_KNOWN_FLAKY_TESTS']),
+            branch: env['BUILDKITE_BRANCH'],
           )
         end
 
@@ -55,7 +57,8 @@ def initialize(
         queue_init_timeout: nil, redis_ttl: 8 * 60 * 60, report_timeout: nil, inactive_workers_timeout: nil,
         export_flaky_tests_file: nil, known_flaky_tests: [],
         strategy: :random, timing_file: nil, timing_fallback_duration: 100.0, export_timing_file: nil,
-        suite_max_duration: 120_000, suite_buffer_percent: 10
+        suite_max_duration: 120_000, suite_buffer_percent: 10,
+        branch: nil
       )
         @build_id = build_id
         @circuit_breakers = [CircuitBreaker::Disabled]
@@ -87,6 +90,7 @@ def initialize(
         @export_timing_file = export_timing_file
         @suite_max_duration = suite_max_duration
         @suite_buffer_percent = suite_buffer_percent
+        @branch = branch
       end
 
       def queue_init_timeout

ruby/lib/ci/queue/redis/test_time_record.rb

@@ -3,6 +3,12 @@ module CI
   module Queue
     module Redis
       class TestTimeRecord < Worker
+        def initialize(redis, config)
+          super
+          @ema_buffer = []
+          @ema_batch_size = 100
+        end
+
         def record(test_name, duration)
           record_test_time(test_name, duration)
           record_test_duration_moving_average(test_name, duration)
@@ -18,7 +24,7 @@ def fetch
 
         private
 
-        attr_reader :redis
+        attr_reader :redis, :ema_batch_size
 
         def record_test_time(test_name, duration)
           redis.pipelined do |pipeline|
@@ -32,7 +38,21 @@ def record_test_time(test_name, duration)
         end
 
         def record_test_duration_moving_average(test_name, duration)
-          UpdateTestDurationMovingAverage.new(redis).update(test_name, duration)
+          return if !should_update_moving_average?
+          @ema_buffer << [test_name, duration]
+          flush_ema_buffer if @ema_buffer.size >= ema_batch_size
+        end
+
+        def flush_ema_buffer
+          return nil if @ema_buffer.empty?
+          UpdateTestDurationMovingAverage.new(redis).update_batch(@ema_buffer)
+          @ema_buffer.clear
+          nil
+        end
+
+        def should_update_moving_average?
+          current_branch = config.respond_to?(:branch) ? config.branch : nil
+          current_branch.to_s == 'master'
         end
 
         def record_test_name(test_name)

ruby/lib/ci/queue/redis/update_test_duration_moving_average.rb

@@ -4,21 +4,25 @@ module CI
   module Queue
     module Redis
       class UpdateTestDurationMovingAverage
-        LUA_SCRIPT= <<~LUA
+        LUA_SCRIPT_BATCH = <<~LUA
           local hash_key = KEYS[1]
-          local test_id = ARGV[1]
-          local new_duration = tonumber(ARGV[2])
-          local smoothing = tonumber(ARGV[3])
-          local current_avg = redis.call('HGET', hash_key, test_id)
-          if current_avg then
-            current_avg = tonumber(current_avg)
-            local new_avg = smoothing * new_duration + (1 - smoothing) * current_avg
-            redis.call('HSET', hash_key, test_id, new_avg)
-            return tostring(new_avg)
-          else
-            redis.call('HSET', hash_key, test_id, new_duration)
-            return tostring(new_duration)
+          local smoothing = tonumber(ARGV[1])
+
+          for i = 2, #ARGV, 2 do
+            local test_id = ARGV[i]
+            local new_duration = tonumber(ARGV[i+1])
+            local current_avg = redis.call('HGET', hash_key, test_id)
+
+            if current_avg then
+              current_avg = tonumber(current_avg)
+              local new_avg = smoothing * new_duration + (1 - smoothing) * current_avg
+              redis.call('HSET', hash_key, test_id, new_avg)
+            else
+              redis.call('HSET', hash_key, test_id, new_duration)
+            end
           end
+
+          return 'OK'
         LUA
 
         def initialize(redis, key: "test_duration_moving_averages", smoothing_factor: 0.2)
@@ -27,8 +31,10 @@ def initialize(redis, key: "test_duration_moving_averages", smoothing_factor: 0.
           @smoothing_factor = smoothing_factor
         end
 
-        def update(test_id, duration)
-          @redis.eval(LUA_SCRIPT, keys: [@key], argv: [test_id, duration, @smoothing_factor]).to_f
+        def update_batch(pairs)
+          return 0 if pairs.nil? || pairs.empty?
+          argv = [@smoothing_factor] + pairs.flat_map { |test_id, duration| [test_id, duration] }
+          @redis.eval(LUA_SCRIPT_BATCH, keys: [@key], argv: argv)
         end
       end
     end

ruby/lib/minitest/queue/test_time_recorder.rb

@@ -13,6 +13,10 @@ def record(test)
         test_id = "#{test.klass}##{test.name}"
         @build.record(test_id, test_duration_in_milliseconds)
       end
+
+      def report
+        @build.flush_ema_buffer if @build.respond_to?(:flush_ema_buffer)
+      end
     end
   end
 end

ruby/test/ci/queue/redis/moving_average_test.rb

@@ -23,11 +23,11 @@ def test_initialize_with_custom_smoothing_factor
     updater = CI::Queue::Redis::UpdateTestDurationMovingAverage.new(@redis, key: @key, smoothing_factor: 0.5)
     reader = CI::Queue::Redis::TestDurationMovingAverages.new(@redis, key: @key)
 
-    updater.update('test1', 100.0)
+    updater.update_batch([["test1", 100.0]])
     first_avg = reader['test1']
     assert_equal 100.0, first_avg
 
-    updater.update('test1', 200.0)
+    updater.update_batch([["test1", 200.0]])
     second_avg = reader['test1']
     assert_in_delta 150.0, second_avg, 0.001
   end
@@ -36,41 +36,39 @@ def test_update_creates_new_entry
     updater = CI::Queue::Redis::UpdateTestDurationMovingAverage.new(@redis, key: @key)
     reader = CI::Queue::Redis::TestDurationMovingAverages.new(@redis, key: @key)
 
-    result = updater.update('test1', 10.5)
-
+    updater.update_batch([["test1", 10.5]])
+    result = reader['test1']
     assert_equal 10.5, result
     assert_equal 1, reader.size
   end
 
   def test_update_calculates_exponential_moving_average
     updater = CI::Queue::Redis::UpdateTestDurationMovingAverage.new(@redis, key: @key, smoothing_factor: 0.2)
+    reader = CI::Queue::Redis::TestDurationMovingAverages.new(@redis, key: @key)
 
-    first_avg = updater.update('test1', 100.0)
-    assert_equal 100.0, first_avg
+    updater.update_batch([["test1", 100.0]])
+    assert_equal 100.0, reader['test1']
 
     # EMA = 0.2 * 150 + 0.8 * 100 = 30 + 80 = 110
-    second_avg = updater.update('test1', 150.0)
-    assert_in_delta 110.0, second_avg, 0.001
+    updater.update_batch([["test1", 150.0]])
+    assert_in_delta 110.0, reader['test1'], 0.001
 
-    third_avg = updater.update('test1', 200.0)
-    assert_in_delta 128.0, third_avg, 0.001
+    updater.update_batch([["test1", 200.0]])
+    assert_in_delta 128.0, reader['test1'], 0.001
   end
 
   def test_update_multiple_tests
     updater = CI::Queue::Redis::UpdateTestDurationMovingAverage.new(@redis, key: @key)
     reader = CI::Queue::Redis::TestDurationMovingAverages.new(@redis, key: @key)
 
-    updater.update('test1', 10.0)
-    updater.update('test2', 20.0)
-    updater.update('test3', 30.0)
+    updater.update_batch([["test1", 10.0], ["test2", 20.0], ["test3", 30.0]])
 
     assert_equal 3, reader.size
   end
 
   def test_bracket_operator_loads_and_returns_value
     updater = CI::Queue::Redis::UpdateTestDurationMovingAverage.new(@redis, key: @key)
-    updater.update('test1', 10.5)
-    updater.update('test2', 20.5)
+    updater.update_batch([["test1", 10.5], ["test2", 20.5]])
 
     reader = CI::Queue::Redis::TestDurationMovingAverages.new(@redis, key: @key)
     assert_equal 10.5, reader['test1']
@@ -81,16 +79,14 @@ def test_bracket_operator_returns_nil_for_missing_key
     updater = CI::Queue::Redis::UpdateTestDurationMovingAverage.new(@redis, key: @key)
     reader = CI::Queue::Redis::TestDurationMovingAverages.new(@redis, key: @key)
 
-    updater.update('test1', 10.5)
+    updater.update_batch([["test1", 10.5]])
 
     assert_nil reader['nonexistent']
   end
 
   def test_load_all_loads_all_values_from_redis
     updater = CI::Queue::Redis::UpdateTestDurationMovingAverage.new(@redis, key: @key)
-    updater.update('test1', 10.0)
-    updater.update('test2', 20.0)
-    updater.update('test3', 30.0)
+    updater.update_batch([["test1", 10.0], ["test2", 20.0], ["test3", 30.0]])
 
     reader = CI::Queue::Redis::TestDurationMovingAverages.new(@redis, key: @key)
     reader.load_all
@@ -104,7 +100,7 @@ def test_load_all_handles_batches
     updater = CI::Queue::Redis::UpdateTestDurationMovingAverage.new(@redis, key: @key)
 
     1500.times do |i|
-      updater.update("test_#{i}", i.to_f)
+      updater.update_batch([["test_#{i}", i.to_f]])
     end
 
     # Create new instance and load all
@@ -123,20 +119,19 @@ def test_size_returns_correct_count
 
     assert_equal 0, reader.size
 
-    updater.update('test1', 10.0)
+    updater.update_batch([["test1", 10.0]])
     assert_equal 1, reader.size
 
-    updater.update('test2', 20.0)
+    updater.update_batch([["test2", 20.0]])
     assert_equal 2, reader.size
 
-    updater.update('test1', 15.0)
+    updater.update_batch([["test1", 15.0]])
     assert_equal 2, reader.size
   end
 
   def test_updates_persist_to_redis
     updater = CI::Queue::Redis::UpdateTestDurationMovingAverage.new(@redis, key: @key)
-    updater.update('test1', 10.5)
-    updater.update('test2', 20.5)
+    updater.update_batch([["test1", 10.5], ["test2", 20.5]])
 
     # Verify data is actually in Redis
     values = @redis.hgetall(@key)
@@ -150,8 +145,8 @@ def test_concurrent_updates_from_different_instances
     updater2 = CI::Queue::Redis::UpdateTestDurationMovingAverage.new(@redis, key: @key, smoothing_factor: 0.2)
     reader = CI::Queue::Redis::TestDurationMovingAverages.new(@redis, key: @key)
 
-    updater1.update('test1', 100.0)
-    updater2.update('test1', 200.0)
+    updater1.update_batch([["test1", 100.0]])
+    updater2.update_batch([["test1", 200.0]])
 
     # Expected: 0.2 * 200 + 0.8 * 100 = 120
     assert_in_delta 120.0, reader['test1'], 0.001
@@ -161,9 +156,7 @@ def test_handles_floating_point_precision
     updater = CI::Queue::Redis::UpdateTestDurationMovingAverage.new(@redis, key: @key)
 
     # Test with various floating point values
-    updater.update('test1', 0.123456789)
-    updater.update('test2', 999.999999)
-    updater.update('test3', 0.000001)
+    updater.update_batch([["test1", 0.123456789], ["test2", 999.999999], ["test3", 0.000001]])
 
     # Load in new instance
     reader = CI::Queue::Redis::TestDurationMovingAverages.new(@redis, key: @key)

ruby/test/ci/queue/redis_test.rb

@@ -436,8 +436,7 @@ def test_suite_bin_packing_uses_moving_average_for_duration
     @redis.flushdb
 
     updater = CI::Queue::Redis::UpdateTestDurationMovingAverage.new(@redis)
-    updater.update('TestSuite#test_1', 5000.0)
-    updater.update('TestSuite#test_2', 3000.0)
+    updater.update_batch([["TestSuite#test_1", 5000.0], ["TestSuite#test_2", 3000.0]])
 
     tests = [
       MockTest.new('TestSuite#test_1'),
@@ -467,7 +466,7 @@ def test_moving_average_takes_precedence_over_timing_file
     timing_data = { 'TestSuite#test_1' => 10_000.0 }
 
     updater = CI::Queue::Redis::UpdateTestDurationMovingAverage.new(@redis)
-    updater.update('TestSuite#test_1', 2000.0)
+    updater.update_batch([["TestSuite#test_1", 2000.0]])
 
     tests = [MockTest.new('TestSuite#test_1')]
 
@@ -539,8 +538,7 @@ def test_mixed_duration_sources_in_suite_splitting
     require 'tempfile'
 
     updater = CI::Queue::Redis::UpdateTestDurationMovingAverage.new(@redis)
-    updater.update('MixedTest#test_1', 60_000.0)
-    updater.update('MixedTest#test_2', 50_000.0)
+    updater.update_batch([["MixedTest#test_1", 60_000.0], ["MixedTest#test_2", 50_000.0]])
 
     timing_data = {
       'MixedTest#test_3' => 40_000.0,
@@ -582,9 +580,7 @@ def test_moving_average_ordering_by_duration
     @redis.flushdb
 
     updater = CI::Queue::Redis::UpdateTestDurationMovingAverage.new(@redis)
-    updater.update('FastTest#test_1', 1000.0)
-    updater.update('SlowTest#test_1', 10_000.0)
-    updater.update('MediumTest#test_1', 5000.0)
+    updater.update_batch([["FastTest#test_1", 1000.0], ["SlowTest#test_1", 10_000.0], ["MediumTest#test_1", 5000.0]])
 
     tests = [
       MockTest.new('FastTest#test_1'),
@@ -616,7 +612,7 @@ def test_moving_average_with_partial_coverage
 
     # Only one test has moving average data
     updater = CI::Queue::Redis::UpdateTestDurationMovingAverage.new(@redis)
-    updater.update('PartialTest#test_1', 3000.0)
+    updater.update_batch([["PartialTest#test_1", 3000.0]])
 
     tests = [
       MockTest.new('PartialTest#test_1'),
@@ -642,7 +638,7 @@ def test_moving_average_updates_persist_across_workers
 
     # Manually update moving average as if a previous worker completed the test
     updater = CI::Queue::Redis::UpdateTestDurationMovingAverage.new(@redis)
-    updater.update('PersistTest#test_1', 5500.0)
+    updater.update_batch([["PersistTest#test_1", 5500.0]])
 
     # New worker should see the persisted moving average
     tests = [MockTest.new('PersistTest#test_1')]