gitlab-test-balancing 0.0.0 → 0.1.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 18b678cbe1614e6ceecbcccc59b3aa07ba7355844bda29ed27e5019f0d2cd20d
-  data.tar.gz: 3d7b2fe157d3d4bcdbe8dc86346c178893e9cd7c72808ff69f767f6104e06b7d
+  metadata.gz: d0847c642c34fc38dbbb8fe6f468c7c906d04a5d756b08d92b0c1fd8f9b715f0
+  data.tar.gz: af8f00f37096a128eb0d0cd0ecc7b190ef15915d660ef64eb7c1c4a0825b9ba5
 SHA512:
-  metadata.gz: 846ded0c9e1ed00910f0e6e7c2d8e2f23537124c7a60db890eea9f03b9cd0752b44d166e435ad70e765328b589cee5976fd1adeecee7eef2e095a15603f346ea
-  data.tar.gz: 482582da2ba877905b12b563aea6b584a2279a212c16519f60ecc0897eee5f147ac019a930d363f438dc4fcef2b77a1bb8b099de31e997f4bb71529eb1a968ff
+  metadata.gz: 4ff4da41f151730ca3d72413bcbce9872c4f3947e0682196b1acdbd1f565d525626287bdd8e055020c6bd9e194b7e5facb0394c918eb9b38faf65437fdfff358
+  data.tar.gz: 2fa0983c324d7d90edcbfbf1d809f7de1a1adfbe3ceb92bf6635ea4ffce6196b0c296d1f5fceb5fdfdd82ecfa3ec741f5f0a864adcb98da6a67361306517562f

data/lib/gitlab/test_balancing/client.rb

@@ -0,0 +1,148 @@
+# frozen_string_literal: true
+
+require 'net/http'
+require 'openssl'
+require 'json'
+require 'securerandom'
+require 'uri'
+
+module Gitlab
+  module TestBalancing
+    # HTTP client for the test balancing REST API.
+    #
+    # Talks to the `/job/test_balancing` endpoints on a GitLab instance,
+    # authenticated with a CI/CD job token. A parallel job seeds the job group's
+    # shared pending pool with its static test split, then repeatedly requests
+    # duration-budgeted batches of test splits until the pool is drained.
+    class Client
+      Error = Class.new(StandardError)
+
+      # Raised when the API responds 404, which the server returns when test
+      # balancing is unavailable for the project (feature flag off or unlicensed).
+      # Callers use this to fall back to a static split.
+      FeatureUnavailableError = Class.new(Error)
+
+      ServerError = Class.new(Error)
+      private_constant :ServerError
+
+      Result = Struct.new(:mode, :test_splits, keyword_init: true)
+
+      OPEN_TIMEOUT = 10
+      READ_TIMEOUT = 30
+      MAX_ATTEMPTS = 3
+      RETRY_BACKOFF = 2 # seconds, multiplied by the attempt number
+
+      # Transient errors worth retrying: network failures and 5xx responses.
+      # Other HTTP error responses (e.g. 404/422) are not retried.
+      RETRIABLE_ERRORS = [
+        Net::OpenTimeout,
+        Net::ReadTimeout,
+        Errno::ECONNRESET,
+        Errno::ECONNREFUSED,
+        Errno::EHOSTUNREACH,
+        Errno::ENETUNREACH,
+        SocketError,
+        IOError,
+        OpenSSL::SSL::SSLError,
+        ServerError
+      ].freeze
+
+      def initialize(api_url: ENV['CI_API_V4_URL'], job_token: ENV['CI_JOB_TOKEN'], logger: nil)
+        @api_url = api_url&.chomp('/')
+        @job_token = job_token
+        @logger = logger
+        @run_id = SecureRandom.uuid
+        @seq = 0
+      end
+
+      # test_splits: array of { path:, expected_duration: }. expected_duration may
+      # be omitted; the server defaults it. Returns a Result with mode ("seed" or
+      # "retry") and the test splits the node should run first.
+      def initialize_balancing(test_splits)
+        response = post('initialize', { test_splits: test_splits })
+
+        Result.new(mode: response['mode'], test_splits: symbolize_test_splits(response))
+      end
+
+      # Returns the next batch of test splits (array), empty when the queue is drained.
+      #
+      # Retries of this call reuse the run_id and seq, so the server returns the
+      # same batch if it already claimed one but the response never arrived, and
+      # ignores a late attempt from an earlier call.
+      def request
+        @seq += 1
+
+        symbolize_test_splits(post('request', { run_id: run_id, seq: @seq }))
+      end
+
+      private
+
+      attr_reader :api_url, :job_token, :logger, :run_id
+
+      def post(action, body)
+        uri = URI.parse("#{api_url}/job/test_balancing/#{action}")
+
+        request = Net::HTTP::Post.new(uri)
+        request['Content-Type'] = 'application/json'
+        request['JOB-TOKEN'] = job_token
+        request.body = JSON.dump(body)
+
+        response = with_retries(action) do
+          net_http_class.start(
+            uri.hostname, uri.port,
+            use_ssl: uri.scheme == 'https', open_timeout: OPEN_TIMEOUT, read_timeout: READ_TIMEOUT
+          ) do |http|
+            http.request(request).tap do |res|
+              raise ServerError, "#{res.code} #{res.body}" if res.is_a?(Net::HTTPServerError)
+            end
+          end
+        end
+
+        raise FeatureUnavailableError, "test balancing is unavailable (404) for #{action}" if
+          response.is_a?(Net::HTTPNotFound)
+
+        unless response.is_a?(Net::HTTPSuccess)
+          raise Error, "test balancing #{action} failed: #{response.code} #{response.body}"
+        end
+
+        JSON.parse(response.body)
+      end
+
+      # Retry the HTTP request on transient errors with linear backoff.
+      # Network failures and 5xx responses are retried; other HTTP error
+      # responses are handled by the caller and are not retried here.
+      def with_retries(action)
+        attempt = 0
+
+        begin
+          attempt += 1
+          yield
+        rescue *RETRIABLE_ERRORS => e
+          raise Error, "test balancing #{action} failed after #{MAX_ATTEMPTS} attempts: #{e.class}: #{e.message}" if
+            attempt >= MAX_ATTEMPTS
+
+          logger&.info("test balancing #{action} attempt #{attempt} failed (#{e.class}), retrying...")
+          sleep(RETRY_BACKOFF * attempt)
+          retry
+        end
+      end
+
+      def symbolize_test_splits(response)
+        Array(response['test_splits']).map { |split| split.transform_keys(&:to_sym) }
+      end
+
+      # When WebMock is loaded (e.g. the client runs inside a spec suite that
+      # enabled WebMock), use the original, un-stubbed Net::HTTP it saved so our
+      # own API calls hit the network without touching WebMock's global config or
+      # tripping the suite's request stubs. Falls back to plain Net::HTTP when
+      # WebMock is absent.
+      def net_http_class
+        if defined?(::WebMock::HttpLibAdapters::NetHttpAdapter::OriginalNetHTTP)
+          ::WebMock::HttpLibAdapters::NetHttpAdapter::OriginalNetHTTP
+        else
+          Net::HTTP
+        end
+      end
+    end
+  end
+end

data/lib/gitlab/test_balancing/runner/rspec.rb

@@ -0,0 +1,198 @@
+# frozen_string_literal: true
+
+require 'gitlab/test_balancing/client'
+
+module Gitlab
+  module TestBalancing
+    module Runner
+      # Drives a dynamic, API-balanced RSpec run for a single parallel node.
+      #
+      # Seeds this node's static test split into the shared pool via the API, then
+      # pulls and runs batches until the queue drains, keeping ONE reporter open
+      # for the whole lane so counts, the `Failures:` block and rerun snippets
+      # print once at the end.
+      #
+      # This class owns only the RSpec lifecycle and the queue-drain loop; it is
+      # project-agnostic. The caller computes the static split (however it likes)
+      # and passes it in. `#run` returns a Result and never calls `exit`, so the
+      # caller controls process termination.
+      #
+      # Each pulled batch is driven through the full `RSpec::Core::Runner#run_specs`
+      # lifecycle
+      # (https://github.com/rspec/rspec-core/blob/main/lib/rspec/core/runner.rb),
+      # loading the batch's files together (the server caps a claim at a small
+      # number of files).
+      #
+      # The only deviation from `run_specs` is keeping ONE outer `reporter.report`
+      # across all batches, and `with_suite_hooks` wrapping the whole lane so
+      # before(:suite)/after(:suite) run once.
+      class Rspec
+        # status: process exit status (0 on success).
+        # unavailable: true when the feature was unavailable (404)
+        Result = Struct.new(:status, :unavailable, keyword_init: true) do
+          def passed?
+            status.zero?
+          end
+
+          def unavailable?
+            unavailable
+          end
+        end
+
+        NULL_LOGGER = Object.new.tap do |logger|
+          def logger.info(*); end
+          def logger.warn(*); end
+        end
+
+        # test_splits: array of { path:, expected_duration: } for this node.
+        # rspec_args: array of RSpec CLI args;
+        #   any trailing file list (after `--`) is stripped since files are supplied dynamically per batch.
+        def initialize(test_splits:, rspec_args: [], client: nil, logger: NULL_LOGGER)
+          @test_splits = test_splits
+          @rspec_args = rspec_args
+          @logger = logger
+          @client = client || Client.new(logger: logger)
+        end
+
+        def run
+          first_batch = seed_and_claim_first_batch
+          return Result.new(status: 0, unavailable: true) if first_batch.nil?
+
+          boot_rspec
+          all_passed = drain_with_report(first_batch)
+
+          Result.new(status: exit_code(all_passed), unavailable: false)
+        ensure
+          # Persist even if the run raised (e.g. a profiler blowing up during
+          # teardown), so scripts can still --only-failures retry the specs that
+          # actually ran. Guarded by @config so an early return/crash is safe.
+          persist_example_statuses if @config
+        end
+
+        private
+
+        attr_reader :test_splits, :rspec_args, :client, :logger
+
+        def seed_and_claim_first_batch
+          logger.info "Seeding #{test_splits.size} test splits into the test balancing pool."
+          result = client.initialize_balancing(test_splits)
+          logger.info "Test balancing initialized in '#{result.mode}' mode."
+
+          result.test_splits
+        rescue Client::FeatureUnavailableError
+          logger.info 'Test balancing is unavailable, falling back to static split.'
+          nil
+        end
+
+        def boot_rspec
+          require 'rspec/core'
+
+          @options = ::RSpec::Core::ConfigurationOptions.new(base_rspec_args)
+          ::RSpec::Core::Runner.new(@options).configure($stderr, $stdout)
+
+          @config = ::RSpec.configuration
+          @world  = ::RSpec.world
+
+          # We need to alter world.example_groups for each iteration and world.all_examples
+          # only ever reflects the current batch. So we keep our own record of executed examples.
+          @executed_examples = []
+        end
+
+        def drain_with_report(first_batch)
+          # One report cycle for the entire lane. The true example count is unknown
+          # up front (files arrive dynamically), so we pass 0 as the expected count;
+          # progress/summary formatters treat this as "unknown total".
+          @config.reporter.report(0) do |reporter|
+            @reporter = reporter
+            @config.with_suite_hooks do # before(:suite)/after(:suite) run once for the lane
+              drain_queue(first_batch)
+            end
+          end
+        end
+
+        # Run the first batch, then keep requesting and running batches until the
+        # queue is drained. Returns true when every example across every batch passed.
+        #
+        # Stops pulling once RSpec sets world.wants_to_quit (fail-fast limit met or
+        # an interrupt), mirroring how Runner#run guards run_specs.
+        def drain_queue(first_batch)
+          all_passed = true
+          batch = first_batch
+
+          loop do
+            all_passed = false unless run_batch(batch)
+
+            break if @world.wants_to_quit
+
+            batch = client.request
+            break if batch.empty?
+          end
+
+          all_passed
+        end
+
+        # Run a whole pulled batch through the full `run_specs` lifecycle against
+        # the shared reporter. Returns true if every example in the batch passed.
+        def run_batch(batch)
+          paths = batch.map { |split| split[:path] }
+          expected_duration = batch.sum { |split| split[:expected_duration].to_f }
+          logger.info "Pulled #{paths.size} test splits (expected duration: #{expected_duration.round(2)}s) " \
+            "from queue:\n#{paths.join("\n")}"
+
+          @world.reset
+
+          # Setup a fresh FilterManager for every batch as
+          # Rspec's announce_filters may mutate the filter manager when
+          # run_all_when_everything_filtered is enabled.
+          filter_manager = ::RSpec::Core::FilterManager.new
+          @options.configure_filter_manager(filter_manager)
+          @config.filter_manager = filter_manager
+
+          @config.files_or_directories_to_run = paths
+          @config.load_spec_files
+          @world.announce_filters
+
+          groups = @world.ordered_example_groups
+          @world.example_count(groups) # force filtered_examples computation, as run_specs does
+
+          passed = groups.map { |group| group.run(@reporter) }.all?
+
+          # Capture this batch's examples before the next iteration's world.reset
+          @executed_examples.concat(@world.all_examples)
+
+          passed
+        end
+
+        # Mirrors RSpec::Core::Runner#exit_code: errors raised outside of examples
+        # (e.g. a before(:suite) hook, or "errors occurred outside of examples")
+        # fail the run even when every example passed.
+        def exit_code(examples_passed)
+          return @config.error_exit_code || @config.failure_exit_code if @world.non_example_failure
+          return @config.failure_exit_code unless examples_passed
+
+          0
+        end
+
+        # Reproduces RSpec::Core::Runner#persist_example_statuses, which the queue
+        # loop bypasses (it never calls Runner#run). Writes the accumulated example
+        # statuses to config.example_status_persistence_file_path in RSpec's native
+        # format so --only-failures retries work unchanged. ExampleStatusPersister
+        # merges with, and file-locks, the existing file, so a single end-of-lane
+        # call is safe.
+        def persist_example_statuses
+          path = @config.example_status_persistence_file_path
+          return if path.nil? || @executed_examples.empty?
+
+          ::RSpec::Core::ExampleStatusPersister.persist(@executed_examples, path)
+        rescue SystemCallError => e
+          logger.warn "Could not write example statuses to #{path}: #{e.inspect}"
+        end
+
+        def base_rspec_args
+          idx = rspec_args.index('--')
+          idx ? rspec_args[0...idx] : rspec_args
+        end
+      end
+    end
+  end
+end

data/lib/gitlab/test_balancing/version.rb

@@ -0,0 +1,7 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module TestBalancing
+    VERSION = '0.1.0'
+  end
+end

data/lib/gitlab-test-balancing.rb

@@ -1 +1,9 @@
-raise "Reserved for GitLab"
+# rubocop:disable Naming/FileName -- required to match the gem name for the default require
+
+# frozen_string_literal: true
+
+require_relative 'gitlab/test_balancing/version'
+require_relative 'gitlab/test_balancing/client'
+require_relative 'gitlab/test_balancing/runner/rspec'
+
+# rubocop:enable Naming/FileName

metadata

@@ -1,15 +1,73 @@
 --- !ruby/object:Gem::Specification
 name: gitlab-test-balancing
 version: !ruby/object:Gem::Version
-  version: 0.0.0
+  version: 0.1.0
 platform: ruby
 authors:
-- GitLab
+- GitLab Engineers
+autorequire:
 bindir: bin
 cert_chain: []
-date: 1980-01-02 00:00:00.000000000 Z
-dependencies: []
-description: Reserved for GitLab
+date: 2026-11-06 00:00:00.000000000 Z
+dependencies:
+- !ruby/object:Gem::Dependency
+  name: gitlab-styles
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: 13.0.1
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: 13.0.1
+- !ruby/object:Gem::Dependency
+  name: rspec
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '3.0'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '3.0'
+- !ruby/object:Gem::Dependency
+  name: webmock
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '3.18'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '3.18'
+- !ruby/object:Gem::Dependency
+  name: webrick
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '1.8'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '1.8'
+description: Distributes tests across the nodes of a parallel CI/CD job using the
+  GitLab test balancing API
 email:
 - engineering@gitlab.com
 executables: []
@@ -17,11 +75,15 @@ extensions: []
 extra_rdoc_files: []
 files:
 - lib/gitlab-test-balancing.rb
-homepage: https://gitlab.com/gitlab-org/gitlab/-/tree/master/gems/gitlab-test-balancing
+- lib/gitlab/test_balancing/client.rb
+- lib/gitlab/test_balancing/runner/rspec.rb
+- lib/gitlab/test_balancing/version.rb
+homepage: https://gitlab.com/gitlab-org/ruby/gems/gitlab-test-balancing
 licenses:
 - MIT
 metadata:
   rubygems_mfa_required: 'true'
+post_install_message:
 rdoc_options: []
 require_paths:
 - lib
@@ -29,14 +91,15 @@ required_ruby_version: !ruby/object:Gem::Requirement
   requirements:
   - - ">="
     - !ruby/object:Gem::Version
-      version: '3.0'
+      version: '3.2'
 required_rubygems_version: !ruby/object:Gem::Requirement
   requirements:
   - - ">="
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 4.0.20
+rubygems_version: 3.5.22
+signing_key:
 specification_version: 4
-summary: Reserved for GitLab
+summary: Client for the GitLab test balancing API
 test_files: []