WIP test-bot: parallelize dependent test with sharding

Library/Homebrew/github_runner_matrix.rb

@@ -64,28 +64,40 @@ def initialize(testing_formulae, deleted_formulae, all_supported:, dependent_mat
     @dependent_matrix = dependent_matrix
     @compatible_testing_formulae = T.let({}, T::Hash[GitHubRunner, T::Array[TestRunnerFormula]])
     @formulae_with_untested_dependents = T.let({}, T::Hash[GitHubRunner, T::Array[TestRunnerFormula]])
-
+    @compatible_untested_dependent_names = T.let({}, T::Hash[GitHubRunner, T::Array[String]])
     @runners = T.let([], T::Array[GitHubRunner])
     generate_runners!
 
     freeze
   end
 
-  sig { returns(T::Array[RunnerSpecHash]) }
+  sig { returns(T::Array[T::Hash[Symbol, T.untyped]]) }
   def active_runner_specs_hash
-    runners.select(&:active)
-           .map(&:spec)
-           .map(&:to_h)
+    runners.filter(&:active).flat_map do |r|
+      Array.new(shard_count = selected_runner_count_for(r)) do |i|
+        (spec = r.spec.to_h).merge(
+          name:        (shard_count > 1) ? "#{spec[:name]} (shard #{i + 1}/#{shard_count})" : spec[:name],
+          shard_index: i,
+          shard_total: shard_count,
+        )
+      end
+    end
   end
 
   private
 
   SELF_HOSTED_LINUX_RUNNER = "linux-self-hosted-1"
+  DEPS_SHARDING_ENV = "HOMEBREW_DEPS_SHARDING"
+  DEPS_SHARD_MAX_RUNNERS_ENV = "HOMEBREW_DEPS_SHARD_MAX_RUNNERS"
+  DEPS_SHARD_BASE_THRESHOLD_ENV = "HOMEBREW_DEPS_SHARD_BASE_THRESHOLD"
+  DEPS_SHARD_RUNNER_PENALTY_ENV = "HOMEBREW_DEPS_SHARD_RUNNER_PENALTY"
   # ARM macOS timeout, keep this under 1/2 of GitHub's job execution time limit for self-hosted runners.
   # https://docs.github.com/en/actions/hosting-your-own-runners/managing-self-hosted-runners/about-self-hosted-runners#usage-limits
   GITHUB_ACTIONS_LONG_TIMEOUT = 2160 # 36 hours
   GITHUB_ACTIONS_SHORT_TIMEOUT = 60
-  private_constant :SELF_HOSTED_LINUX_RUNNER, :GITHUB_ACTIONS_LONG_TIMEOUT, :GITHUB_ACTIONS_SHORT_TIMEOUT
+  private_constant :SELF_HOSTED_LINUX_RUNNER, :DEPS_SHARDING_ENV, :DEPS_SHARD_MAX_RUNNERS_ENV,
+                   :DEPS_SHARD_BASE_THRESHOLD_ENV, :DEPS_SHARD_RUNNER_PENALTY_ENV,
+                   :GITHUB_ACTIONS_LONG_TIMEOUT, :GITHUB_ACTIONS_SHORT_TIMEOUT
 
   sig { params(arch: Symbol).returns(LinuxRunnerSpec) }
   def linux_runner_spec(arch)
@@ -261,6 +273,62 @@ def generate_runners!
     @runners.freeze
   end
 
+  sig { params(runner: GitHubRunner).returns(Integer) }
+  def selected_runner_count_for(runner)
+    return 1 if !@dependent_matrix || %w[1 true].exclude?(ENV.fetch(DEPS_SHARDING_ENV, "false").downcase)
+
+    max_available_runners = T.must(sharding_integer_env(DEPS_SHARD_MAX_RUNNERS_ENV, runner, 1))
+    base_spillover_threshold = sharding_integer_env(DEPS_SHARD_BASE_THRESHOLD_ENV, runner, nil)
+    additional_runner_reluctance = sharding_integer_env(DEPS_SHARD_RUNNER_PENALTY_ENV, runner, nil)
+    raise ArgumentError, "#{DEPS_SHARD_MAX_RUNNERS_ENV} must be positive" unless max_available_runners.positive?
+    raise ArgumentError, "#{DEPS_SHARD_BASE_THRESHOLD_ENV} must be positive" if base_spillover_threshold&.<= 0
+
+    if additional_runner_reluctance&.negative?
+      raise ArgumentError,
+            "#{DEPS_SHARD_RUNNER_PENALTY_ENV} must be non-negative"
+    end
+    return 1 if max_available_runners <= 1 || base_spillover_threshold.nil? || additional_runner_reluctance.nil?
+
+    dependent_names = @compatible_untested_dependent_names[runner] ||= compatible_testing_formulae(runner)
+                                                                       .flat_map do |formula|
+      compatible_untested_dependent_names_for_formula(
+        formula, runner
+      )
+    end.uniq.sort
+    dependent_count = dependent_names.count
+    return 1 if dependent_count.zero?
+
+    runner_count = if additional_runner_reluctance.zero?
+      dependent_count.to_f / base_spillover_threshold
+    else
+      threshold_difference = base_spillover_threshold - additional_runner_reluctance
+      discriminant = (threshold_difference**2) + (4 * additional_runner_reluctance * dependent_count)
+      (Math.sqrt(discriminant.to_f) - threshold_difference) / (2 * additional_runner_reluctance)
+    end
+
+    runner_count.ceil.clamp(1, max_available_runners)
+  end
+
+  sig { params(base_env_name: String, runner: GitHubRunner, default: T.nilable(Integer)).returns(T.nilable(Integer)) }
+  def sharding_integer_env(base_env_name, runner, default)
+    platform = runner.platform.to_s.upcase
+    arch = runner.arch.to_s.upcase
+    version = runner.macos_version&.to_sym
+    version = version.to_s.upcase if version
+
+    [
+      ("#{base_env_name}_#{platform}_#{arch}_#{version}" if version),
+      "#{base_env_name}_#{platform}_#{arch}",
+      "#{base_env_name}_#{platform}",
+      base_env_name,
+    ].compact.each do |env_name|
+      env_value = ENV.fetch(env_name, nil)
+      return Integer(env_value, 10) if env_value.present?
+    end
+
+    default
+  end
+
   sig { params(runner: GitHubRunner).returns(T::Array[String]) }
   def testable_formulae(runner)
     formulae = if @dependent_matrix
@@ -301,27 +369,30 @@ def compatible_testing_formulae(runner)
 
   sig { params(runner: GitHubRunner).returns(T::Array[TestRunnerFormula]) }
   def formulae_with_untested_dependents(runner)
-    @formulae_with_untested_dependents[runner] ||= begin
-      platform = runner.platform
-      arch = runner.arch
-      macos_version = runner.macos_version
+    @formulae_with_untested_dependents[runner] ||= compatible_testing_formulae(runner).select do |formula|
+      compatible_untested_dependent_names_for_formula(formula, runner).present?
+    end
+  end
 
-      compatible_testing_formulae(runner).select do |formula|
-        compatible_dependents = formula.dependents(platform:, arch:, macos_version: macos_version&.to_sym)
-                                       .select do |dependent_f|
-          Homebrew::SimulateSystem.with(os: platform, arch: Homebrew::SimulateSystem.arch_symbols.fetch(arch)) do
-            simulated_dependent_f = TestRunnerFormula.new(Formulary.factory(dependent_f.name))
-            next false if macos_version && !simulated_dependent_f.compatible_with?(macos_version)
+  sig { params(formula: TestRunnerFormula, runner: GitHubRunner).returns(T::Array[String]) }
+  def compatible_untested_dependent_names_for_formula(formula, runner)
+    platform = runner.platform
+    arch = runner.arch
+    macos_version = runner.macos_version
 
-            simulated_dependent_f.public_send(:"#{platform}_compatible?") &&
-              simulated_dependent_f.public_send(:"#{arch}_compatible?")
-          end
-        end
+    compatible_dependents = formula.dependents(platform:, arch:, macos_version: macos_version&.to_sym)
+                                   .select do |dependent_f|
+      Homebrew::SimulateSystem.with(os: platform, arch: Homebrew::SimulateSystem.arch_symbols.fetch(arch)) do
+        simulated_dependent_f = TestRunnerFormula.new(Formulary.factory(dependent_f.name))
+        next false if macos_version && !simulated_dependent_f.compatible_with?(macos_version)
 
-        # These arrays will generally have been generated by different Formulary caches,
-        # so we can only compare them by name and not directly.
-        (compatible_dependents.map(&:name) - @testing_formulae.map(&:name)).present?
+        simulated_dependent_f.public_send(:"#{platform}_compatible?") &&
+          simulated_dependent_f.public_send(:"#{arch}_compatible?")
       end
     end
+
+    # These arrays will generally have been generated by different Formulary caches,
+    # so we can only compare them by name and not directly.
+    compatible_dependents.map(&:name) - @testing_formulae.map(&:name)
   end
 end

Library/Homebrew/test/github_runner_matrix_spec.rb

@@ -6,6 +6,7 @@
 
 RSpec.describe GitHubRunnerMatrix, :no_api do
   before do
+    allow(ENV).to receive(:[]).and_call_original
     allow(ENV).to receive(:fetch).and_call_original
     allow(ENV).to receive(:fetch).with("HOMEBREW_LINUX_RUNNER").and_return("ubuntu-latest")
     allow(ENV).to receive(:fetch).with("HOMEBREW_MACOS_LONG_TIMEOUT", "false").and_return("false")
@@ -292,6 +293,95 @@
               matrix = described_class.new([testball], ["deleted"], all_supported: false, dependent_matrix: true)
               expect(get_runner_names(matrix)).to eq(["Linux x86_64"])
             end
+
+            it "shards dependent runners using the most specific sharding settings" do
+              allow(Homebrew::EnvConfig).to receive(:eval_all?).and_return(true)
+              allow(Formula).to receive(:all).and_return([
+                testball,
+                testball_depender_linux,
+                setup_test_runner_formula("testball-depender-linux-two", ["testball", :linux]),
+              ].map(&:formula))
+              allow(ENV).to receive(:[]).with("HOMEBREW_LINUX_RUNNER").and_return("linux-self-hosted-1")
+              env_fetch_overrides = {
+                ["HOMEBREW_LINUX_RUNNER", "ubuntu-latest"]               => "linux-self-hosted-1",
+                ["HOMEBREW_DEPS_SHARDING", "false"]                      => "1",
+                ["HOMEBREW_DEPS_SHARD_BASE_THRESHOLD_LINUX_X86_64", nil] => "1",
+                ["HOMEBREW_DEPS_SHARD_BASE_THRESHOLD_LINUX", nil]        => "10",
+                ["HOMEBREW_DEPS_SHARD_RUNNER_PENALTY", nil]              => "0",
+                ["HOMEBREW_DEPS_SHARD_MAX_RUNNERS_LINUX", nil]           => "2",
+              }
+              allow(ENV).to receive(:fetch).and_wrap_original do |original, key, default = nil|
+                if env_fetch_overrides.key?([key, default])
+                  env_fetch_overrides[[key, default]]
+                else
+                  original.call(key, default)
+                end
+              end
+              matrix = described_class.new([testball], [], all_supported: false, dependent_matrix: true)
+
+              expect(matrix.active_runner_specs_hash)
+                .to eq([
+                  {
+                    name:             "Linux x86_64 (shard 1/2)",
+                    runner:           "linux-self-hosted-1",
+                    container:        {
+                      image:   "ghcr.io/homebrew/brew:main",
+                      options: "--user=linuxbrew -e GITHUB_ACTIONS_HOMEBREW_SELF_HOSTED",
+                    },
+                    workdir:          "/github/home",
+                    timeout:          2160,
+                    cleanup:          true,
+                    testing_formulae: "testball",
+                    shard_index:      0,
+                    shard_total:      2,
+                  },
+                  {
+                    name:             "Linux x86_64 (shard 2/2)",
+                    runner:           "linux-self-hosted-1",
+                    container:        {
+                      image:   "ghcr.io/homebrew/brew:main",
+                      options: "--user=linuxbrew -e GITHUB_ACTIONS_HOMEBREW_SELF_HOSTED",
+                    },
+                    workdir:          "/github/home",
+                    timeout:          2160,
+                    cleanup:          true,
+                    testing_formulae: "testball",
+                    shard_index:      1,
+                    shard_total:      2,
+                  },
+                ])
+            end
+
+            it "counts shared dependents once when sizing shards" do
+              shared_testing_formula = setup_test_runner_formula("testball-user")
+              allow(Homebrew::EnvConfig).to receive(:eval_all?).and_return(true)
+              allow(Formula).to receive(:all).and_return([
+                testball,
+                shared_testing_formula,
+                setup_test_runner_formula("testball-shared-depender-linux", ["testball", "testball-user", :linux]),
+                setup_test_runner_formula("testball-unique-depender-linux", ["testball", :linux]),
+              ].map(&:formula))
+              allow(ENV).to receive(:[]).with("HOMEBREW_LINUX_RUNNER").and_return("linux-self-hosted-1")
+              env_fetch_overrides = {
+                ["HOMEBREW_LINUX_RUNNER", "ubuntu-latest"]        => "linux-self-hosted-1",
+                ["HOMEBREW_DEPS_SHARDING", "false"]               => "1",
+                ["HOMEBREW_DEPS_SHARD_BASE_THRESHOLD_LINUX", nil] => "1",
+                ["HOMEBREW_DEPS_SHARD_RUNNER_PENALTY", nil]       => "0",
+                ["HOMEBREW_DEPS_SHARD_MAX_RUNNERS_LINUX", nil]    => "3",
+              }
+              allow(ENV).to receive(:fetch).and_wrap_original do |original, key, default = nil|
+                if env_fetch_overrides.key?([key, default])
+                  env_fetch_overrides[[key, default]]
+                else
+                  original.call(key, default)
+                end
+              end
+
+              matrix = described_class.new([testball, shared_testing_formula], [], all_supported:    false,
+                                                                                   dependent_matrix: true)
+
+              expect(matrix.active_runner_specs_hash.count).to eq(2)
+            end
           end
 
           context "when dependent formulae require macOS" do

Library/Homebrew/test/test_bot/formulae_dependents_spec.rb

@@ -0,0 +1,104 @@
+# typed: false
+# frozen_string_literal: true
+
+require "dev-cmd/test-bot"
+
+RSpec.describe Homebrew::TestBot::FormulaeDependents do
+  subject(:formulae_dependents) do
+    described_class.new(tap: nil, git: "git", dry_run: true, fail_fast: false, verbose: false)
+  end
+
+  let(:args) { double(skip_recursive_dependents?: false) }
+
+  describe "#configure_dependent_sharding!" do
+    before do
+      formulae_dependents.instance_variable_set(:@dependent_testing_formulae, %w[testball testball-user])
+      allow(Formulary).to receive(:factory).with("testball").and_return(instance_double(Formula))
+      allow(Formulary).to receive(:factory).with("testball-user").and_return(instance_double(Formula))
+      allow(formulae_dependents).to receive(:skip_recursive_dependents_for).and_return(false)
+      allow(formulae_dependents).to receive(:dependent_formula_names)
+        .with("testball", skip_recursive_dependents: false)
+        .and_return(%w[alpha shared])
+      allow(formulae_dependents).to receive(:dependent_formula_names)
+        .with("testball-user", skip_recursive_dependents: false)
+        .and_return(%w[beta shared])
+    end
+
+    it "deduplicates shared dependents across testing formulae before sharding" do
+      with_env(
+        "HOMEBREW_DEPS_SHARD_INDEX" => "0",
+        "HOMEBREW_DEPS_SHARD_TOTAL" => "2",
+      ) do
+        formulae_dependents.send(:configure_dependent_sharding!, args:)
+      end
+
+      expect(formulae_dependents.instance_variable_get(:@assigned_dependent_formula_names))
+        .to eq(Set.new(%w[alpha shared]))
+    end
+
+    it "assigns dependents across shards without duplicates" do
+      shard_zero = described_class.new(tap: nil, git: "git", dry_run: true, fail_fast: false, verbose: false)
+      shard_one = described_class.new(tap: nil, git: "git", dry_run: true, fail_fast: false, verbose: false)
+
+      [shard_zero, shard_one].each do |deps|
+        deps.instance_variable_set(:@dependent_testing_formulae, %w[testball testball-user])
+        allow(deps).to receive(:skip_recursive_dependents_for).and_return(false)
+        allow(Formulary).to receive(:factory).with("testball").and_return(instance_double(Formula))
+        allow(Formulary).to receive(:factory).with("testball-user").and_return(instance_double(Formula))
+        allow(deps).to receive(:dependent_formula_names)
+          .with("testball", skip_recursive_dependents: false)
+          .and_return(%w[alpha gamma epsilon])
+        allow(deps).to receive(:dependent_formula_names)
+          .with("testball-user", skip_recursive_dependents: false)
+          .and_return(%w[beta delta])
+      end
+
+      with_env("HOMEBREW_DEPS_SHARD_INDEX" => "0", "HOMEBREW_DEPS_SHARD_TOTAL" => "2") do
+        shard_zero.send(:configure_dependent_sharding!, args:)
+      end
+      with_env("HOMEBREW_DEPS_SHARD_INDEX" => "1", "HOMEBREW_DEPS_SHARD_TOTAL" => "2") do
+        shard_one.send(:configure_dependent_sharding!, args:)
+      end
+
+      assigned_zero = shard_zero.instance_variable_get(:@assigned_dependent_formula_names)
+      assigned_one = shard_one.instance_variable_get(:@assigned_dependent_formula_names)
+
+      expect(assigned_zero | assigned_one).to eq(Set.new(%w[alpha beta delta epsilon gamma]))
+      expect(assigned_zero & assigned_one).to eq(Set.new)
+      expect(assigned_zero.count - assigned_one.count).to be <= 1
+    end
+
+    it "preserves existing behavior for a single shard" do
+      dependent = instance_double(Formula, full_name: "alpha")
+      formulae_dependents.instance_variable_set(:@handled_dependent_formula_names, Set["alpha"])
+
+      with_env("HOMEBREW_DEPS_SHARD_TOTAL" => "1") do
+        formulae_dependents.send(:configure_dependent_sharding!, args:)
+      end
+
+      expect(formulae_dependents.instance_variable_get(:@assigned_dependent_formula_names)).to be_nil
+      expect(formulae_dependents.instance_variable_get(:@handled_dependent_formula_names)).to eq(Set.new)
+      expect(formulae_dependents.send(:sharded_dependents, [dependent])).to eq([dependent])
+    end
+
+    it "tests a shared dependent once when multiple changed formulae reach it on the same shard" do
+      alpha = instance_double(Formula, full_name: "alpha")
+      shared = instance_double(Formula, full_name: "shared")
+
+      with_env(
+        "HOMEBREW_DEPS_SHARD_INDEX" => "0",
+        "HOMEBREW_DEPS_SHARD_TOTAL" => "2",
+      ) do
+        formulae_dependents.send(:configure_dependent_sharding!, args:)
+      end
+
+      first_formula_dependents = formulae_dependents.send(:sharded_dependents, [alpha, shared])
+      first_formula_dependents.each do |dependent|
+        formulae_dependents.instance_variable_get(:@handled_dependent_formula_names).add(dependent.full_name)
+      end
+
+      expect(first_formula_dependents).to eq([alpha, shared])
+      expect(formulae_dependents.send(:sharded_dependents, [shared])).to eq([])
+    end
+  end
+end