Initial test version of lazy loading

Author: danielahrnsbrak

ruby/README.md

@@ -38,6 +38,31 @@ minitest-queue --queue redis://example.com run -Itest test/**/*_test.rb
 
 Additionally you can configure the requeue settings (see main README) with `--max-requeues` and `--requeue-tolerance`.
 
+#### Lazy Loading (Minitest only)
+
+For large test suites, you can enable lazy loading to reduce memory usage on worker nodes. With lazy loading enabled:
+
+- The **leader** worker loads all test files to build a manifest mapping test classes to their source files
+- **Consumer** workers only load test files on-demand as they pick up tests from the queue
+
+This can provide significant memory savings (40-80%+) for workers that don't need to run the entire test suite.
+
+```bash
+minitest-queue --queue redis://example.com run \
+  --lazy-load \
+  --test-helpers test/test_helper.rb \
+  -Itest test/**/*_test.rb
+```
+
+Options:
+- `--lazy-load`: Enable lazy loading mode
+- `--test-helpers`: Comma-separated list of helper files to load before tests (e.g., `test/test_helper.rb`)
+
+You can also enable lazy loading via environment variables:
+- `CI_QUEUE_LAZY_LOAD=true`
+- `CI_QUEUE_TEST_HELPERS=test/test_helper.rb`
+
+**Note:** Lazy loading is currently only available for Minitest.
 
 If you'd like to centralize the error reporting you can do so with:
 

ruby/lib/ci/queue.rb

@@ -14,6 +14,7 @@
 require 'ci/queue/file'
 require 'ci/queue/grind'
 require 'ci/queue/bisect'
+require 'ci/queue/lazy_loader'
 
 module CI
   module Queue

ruby/lib/ci/queue/configuration.rb

@@ -6,6 +6,7 @@ class Configuration
       attr_accessor :requeue_tolerance, :namespace, :failing_test, :statsd_endpoint
       attr_accessor :max_test_duration, :max_test_duration_percentile, :track_test_duration
       attr_accessor :max_test_failed, :redis_ttl, :warnings_file, :debug_log, :max_missed_heartbeat_seconds
+      attr_accessor :lazy_load, :test_helpers
       attr_reader :circuit_breakers
       attr_writer :seed, :build_id
       attr_writer :queue_init_timeout, :report_timeout, :inactive_workers_timeout
@@ -22,6 +23,8 @@ def from_env(env)
             debug_log: env['CI_QUEUE_DEBUG_LOG'],
             max_requeues: env['CI_QUEUE_MAX_REQUEUES']&.to_i || 0,
             requeue_tolerance: env['CI_QUEUE_REQUEUE_TOLERANCE']&.to_f || 0,
+            lazy_load: env['CI_QUEUE_LAZY_LOAD'] == 'true',
+            test_helpers: env['CI_QUEUE_TEST_HELPERS'],
           )
         end
 
@@ -46,7 +49,8 @@ def initialize(
         grind_count: nil, max_duration: nil, failure_file: nil, max_test_duration: nil,
         max_test_duration_percentile: 0.5, track_test_duration: false, max_test_failed: nil,
         queue_init_timeout: nil, redis_ttl: 8 * 60 * 60, report_timeout: nil, inactive_workers_timeout: nil,
-        export_flaky_tests_file: nil, warnings_file: nil, debug_log: nil, max_missed_heartbeat_seconds: nil)
+        export_flaky_tests_file: nil, warnings_file: nil, debug_log: nil, max_missed_heartbeat_seconds: nil,
+        lazy_load: false, test_helpers: nil)
         @build_id = build_id
         @circuit_breakers = [CircuitBreaker::Disabled]
         @failure_file = failure_file
@@ -73,6 +77,8 @@ def initialize(
         @warnings_file = warnings_file
         @debug_log = debug_log
         @max_missed_heartbeat_seconds = max_missed_heartbeat_seconds
+        @lazy_load = lazy_load
+        @test_helpers = test_helpers
       end
 
       def queue_init_timeout
@@ -118,6 +124,16 @@ def build_id
       def global_max_requeues(tests_count)
         (tests_count * Float(requeue_tolerance)).ceil
       end
+
+      def lazy_load?
+        @lazy_load
+      end
+
+      def test_helper_paths
+        return [] unless @test_helpers
+
+        @test_helpers.split(',').map(&:strip)
+      end
     end
   end
 end

ruby/lib/ci/queue/lazy_loader.rb

@@ -0,0 +1,152 @@
+# frozen_string_literal: true
+
+require 'json'
+require 'set'
+
+module CI
+  module Queue
+    class LazyLoadError < StandardError; end
+
+    # LazyLoader handles on-demand loading of test files based on a manifest
+    # that maps class names to their source file paths.
+    #
+    # The manifest is a Hash mapping class names (e.g., "MyModule::MyTest") to
+    # absolute file paths where those classes are defined.
+    #
+    # Test IDs follow the format "ClassName#method_name" (e.g., "MyTest#test_foo").
+    class LazyLoader
+      attr_reader :loaded_files
+
+      def initialize
+        @loaded_files = Set.new
+        @manifest = {}
+      end
+
+      # Build manifest from loaded tests
+      # Returns a hash mapping class_name -> file_path
+      # Handles both Minitest::Queue::SingleExample and regular test objects
+      def self.build_manifest(tests)
+        manifest = {}
+        tests.each do |test|
+          # SingleExample has `runnable` (the class) and `method_name`
+          # Regular test objects have `class` and `name`
+          if test.respond_to?(:runnable)
+            # Minitest::Queue::SingleExample
+            class_name = test.runnable.name
+            source_location = test.source_location&.first
+          else
+            # Regular test object
+            class_name = test.class.name
+            method_name = test.respond_to?(:method_name) ? test.method_name : test.name
+            source_location = test.method(method_name).source_location&.first rescue nil
+          end
+
+          next if class_name.nil? || source_location.nil?
+
+          # Warn about duplicate class names - only one file path will be used
+          if manifest.key?(class_name) && manifest[class_name] != source_location
+            warn "[ci-queue] WARNING: Duplicate class name '#{class_name}' found in multiple files:\n" \
+                 "  - #{manifest[class_name]}\n" \
+                 "  - #{source_location}\n" \
+                 "Only one will be used for lazy loading. Rename one class to avoid test failures."
+          end
+
+          manifest[class_name] = source_location
+        end
+        manifest
+      end
+
+      # Store manifest in Redis
+      # key: Redis key to store the manifest
+      def store_manifest(redis, key, manifest, ttl:)
+        return if manifest.empty?
+
+        redis.hset(key, manifest)
+        redis.expire(key, ttl)
+      end
+
+      # Fetch manifest from Redis with retry logic
+      # key: Redis key where manifest is stored
+      # retries: number of retries if manifest is empty
+      def fetch_manifest(redis, key, retries: 3, retry_delay: 0.5)
+        return @manifest unless @manifest.empty?
+
+        retries.times do |attempt|
+          @manifest = redis.hgetall(key)
+          break unless @manifest.empty?
+
+          sleep(retry_delay * (attempt + 1)) if attempt < retries - 1
+        end
+
+        if @manifest.empty?
+          raise LazyLoadError, "Failed to fetch manifest from Redis after #{retries} attempts. " \
+                               "The leader may not have finished populating the queue."
+        end
+
+        @manifest
+      end
+
+      # Set manifest directly (e.g., from leader)
+      def set_manifest(manifest)
+        @manifest = manifest
+      end
+
+      # Load test class if not already loaded
+      # Returns the test runnable instance
+      def load_test(class_name, method_name)
+        load_class(class_name)
+        instantiate_test(class_name, method_name)
+      end
+
+      # Load a class from the manifest
+      def load_class(class_name)
+        file_path = @manifest[class_name]
+        raise LazyLoadError, "No manifest entry for #{class_name}" unless file_path
+
+        # Use `load` instead of `require` - classes may be partially defined
+        # (constant exists but methods missing due to autoloader/parent class).
+        # Ruby's require checks $LOADED_FEATURES and may skip re-execution,
+        # whereas load always executes the file.
+        # We track loaded files ourselves to prevent duplicate loading.
+        unless @loaded_files.include?(file_path)
+          load(file_path)
+          @loaded_files.add(file_path)
+        end
+      end
+
+      # Find a class by name using const_get
+      # Raises LazyLoadError with helpful message if class not found
+      def find_class(class_name)
+        class_name.split('::').reduce(Object) do |ns, const|
+          ns.const_get(const, false)
+        end
+      rescue NameError => e
+        file_path = @manifest[class_name]
+        raise LazyLoadError, "Class #{class_name} not found after loading #{file_path}. " \
+                             "The file may not define the expected class. Original error: #{e.message}"
+      end
+
+      # Instantiate a test runnable
+      def instantiate_test(class_name, method_name)
+        klass = find_class(class_name)
+        klass.new(method_name)
+      end
+
+      # Parse a test identifier into class_name and method_name
+      def self.parse_test_id(test_id)
+        # Test IDs are in format "ClassName#method_name" or "Module::ClassName#method_name"
+        class_name, method_name = test_id.split('#', 2)
+        [class_name, method_name]
+      end
+
+      # Build a test identifier from class_name and method_name
+      def self.build_test_id(class_name, method_name)
+        "#{class_name}##{method_name}"
+      end
+
+      def files_loaded_count
+        @loaded_files.size
+      end
+    end
+  end
+end

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

@@ -1,5 +1,6 @@
 # frozen_string_literal: true
 require 'ci/queue/static'
+require 'ci/queue/lazy_loader'
 require 'concurrent/set'
 
 module CI
@@ -18,6 +19,7 @@ class Worker < Base
         def initialize(redis, config)
           @reserved_tests = Concurrent::Set.new
           @shutdown_required = false
+          @lazy_loader = LazyLoader.new if config.lazy_load?
           super(redis, config)
         end
 
@@ -32,8 +34,67 @@ def populate(tests, random: Random.new)
           self
         end
 
+        # Populate queue with lazy loading support
+        # Only the leader loads all test files and builds the manifest
+        # Workers load files on-demand when they claim tests
+        def populate_lazy(test_files:, random:, config:)
+          @lazy_load_mode = true
+          @lazy_loader ||= LazyLoader.new
+
+          push do
+            begin
+              # This block only runs on master - load files and build manifest
+              test_files.each do |f|
+                require(f)
+              rescue LoadError => e
+                raise LazyLoadError, "Failed to load test file #{f}: #{e.message}"
+              end
+
+              tests = Minitest.loaded_tests
+              if tests.empty?
+                raise LazyLoadError, "No tests found after loading #{test_files.size} test files. " \
+                                     "Ensure test files define Minitest::Test subclasses."
+              end
+
+              # Build and store manifest
+              manifest = LazyLoader.build_manifest(tests)
+              @lazy_loader.set_manifest(manifest)
+              @lazy_loader.store_manifest(redis, key('manifest'), manifest, ttl: config.redis_ttl)
+
+              # Count files loaded for metrics
+              source_files = Set.new
+              tests.each do |test|
+                source_location = test.source_location&.first
+                source_files.add(source_location) if source_location
+              end
+              @files_loaded_by_leader = source_files.size
+
+              puts "Leader loaded #{@files_loaded_by_leader} test files, found #{tests.size} tests."
+
+              # Return shuffled test IDs
+              Queue.shuffle(tests, random).map(&:id)
+            rescue LazyLoadError
+              raise
+            rescue => error
+              build.report_worker_error(error)
+              raise LazyLoadError, "Failed to build manifest: #{error.class}: #{error.message}"
+            end
+          end
+          self
+        end
+
         def populated?
-          !!defined?(@index)
+          !!defined?(@index) || @lazy_load_mode
+        end
+
+        def lazy_load?
+          @lazy_load_mode || config.lazy_load?
+        end
+
+        def files_loaded_count
+          return 0 unless @lazy_loader
+
+          @lazy_loader.files_loaded_count
         end
 
         def shutdown!
@@ -52,11 +113,13 @@ def master?
 
         def poll
           wait_for_master
+          fetch_manifest_for_lazy_load if lazy_load?
           attempt = 0
           until shutdown_required? || config.circuit_breakers.any?(&:open?) || exhausted? || max_test_failed?
-            if test = reserve
+            if test_id = reserve
               attempt = 0
-              yield index.fetch(test)
+              example = lazy_load? ? load_test_lazily(test_id) : index.fetch(test_id)
+              yield example
             else
               # Adding exponential backoff to avoid hammering Redis
               # we just stay online here in case a test gets retried or times out so we can afford to wait
@@ -72,6 +135,22 @@ def poll
         rescue *CONNECTION_ERRORS
         end
 
+        def fetch_manifest_for_lazy_load
+          return unless @lazy_loader
+          return if @manifest_fetched
+
+          @lazy_loader.fetch_manifest(redis, key('manifest'))
+          @manifest_fetched = true
+        end
+
+        def load_test_lazily(test_id)
+          class_name, method_name = LazyLoader.parse_test_id(test_id)
+          @lazy_loader.load_class(class_name)
+          # Return a SingleExample like the index would
+          runnable = @lazy_loader.find_class(class_name)
+          Minitest::Queue::SingleExample.new(runnable, method_name)
+        end
+
         if ::Redis.method_defined?(:exists?)
           def retrying?
             redis.exists?(key('worker', worker_id, 'queue'))
@@ -208,8 +287,11 @@ def try_to_reserve_lost_test
           lost_test
         end
 
-        def push(tests)
-          @total = tests.size
+        # Push test IDs to the queue.
+        # Can be called with test IDs directly, or with a block that returns test IDs.
+        # The block form is used for lazy loading where only the master loads test files.
+        def push(tests = nil, &block)
+          @total = tests.size if tests
 
           # We set a unique value (worker_id) and read it back to make "SET if Not eXists" idempotent in case of a retry.
           value = key('setup', worker_id)
@@ -219,6 +301,10 @@ def push(tests)
           end
 
           if @master = (value == status)
+            # If block given, call it to get test IDs (used for lazy loading)
+            tests = yield if block_given?
+            @total = tests.size
+
             puts "Worker elected as leader, pushing #{@total} tests to the queue."
             puts