Add partition pruning to Ci::Runner and Ci::RunnerManager queries

Summary

Both ci_runners and ci_runner_machines tables are LIST-partitioned by runner_type (1 = instance, 2 = group, 3 = project). PostgreSQL can skip entire partitions when WHERE runner_type = X is present in the query. Several scopes, methods, and association lookups that only ever target a known runner type do not include this filter, forcing Postgres to scan all 3 partitions unnecessarily.

Patch
diff --git a/app/models/ci/runner.rb b/app/models/ci/runner.rb
index c8d22c3bfe82..a895eeb3103b 100644
--- a/app/models/ci/runner.rb
+++ b/app/models/ci/runner.rb
@@ -93,14 +93,18 @@ class Runner < Ci::ApplicationRecord
 
     TAG_LIST_MAX_LENGTH = 50
 
-    has_many :runner_managers, inverse_of: :runner
+    has_many :runner_managers,
+      ->(runner) { where(runner_type: runner.runner_type) if runner && Ci::Runner.partition_pruning_enabled? },
+      inverse_of: :runner
     has_many :builds
     has_many :running_builds, inverse_of: :runner
     has_many :runner_projects, inverse_of: :runner, autosave: true, dependent: :destroy # rubocop:disable Cop/ActiveRecordDependent
     has_many :projects, through: :runner_projects, disable_joins: true
     has_many :runner_namespaces, inverse_of: :runner, autosave: true
     has_many :groups, through: :runner_namespaces, disable_joins: true
-    has_many :taggings, class_name: 'Ci::RunnerTagging', inverse_of: :runner
+    has_many :taggings,
+      ->(runner) { where(runner_type: runner.runner_type) if runner && Ci::Runner.partition_pruning_enabled? },
+      class_name: 'Ci::RunnerTagging', inverse_of: :runner
     has_many :tags, class_name: 'Ci::Tag', through: :taggings, source: :tag
 
     # currently we have only 1 namespace assigned, but order is here for consistency
@@ -142,11 +146,15 @@ class Runner < Ci::ApplicationRecord
     scope :deprecated_specific, -> { project_type.or(group_type) }
 
     scope :belonging_to_project, ->(project_id) {
-      joins(:runner_projects).where(ci_runner_projects: { project_id: project_id })
+      scope = joins(:runner_projects).where(ci_runner_projects: { project_id: project_id })
+      scope = scope.project_type if partition_pruning_enabled?
+      scope
     }
 
     scope :belonging_to_group, ->(group_id) {
-      joins(:runner_namespaces).where(ci_runner_namespaces: { namespace_id: group_id })
+      scope = joins(:runner_namespaces).where(ci_runner_namespaces: { namespace_id: group_id })
+      scope = scope.group_type if partition_pruning_enabled?
+      scope
     }
 
     scope :in_organization, ->(organization_id) { where(organization_id: organization_id) }
@@ -260,7 +268,11 @@ class Runner < Ci::ApplicationRecord
     validate :exactly_one_group, if: :group_type?
     validate :no_allowed_plan_ids, unless: :instance_type?
 
-    scope :with_version_prefix, ->(value) { joins(:runner_managers).merge(RunnerManager.with_version_prefix(value)) }
+    scope :with_version_prefix, ->(value) {
+      runner_managers_join = partition_pruning_enabled? ? partitioned_runner_managers_join : :runner_managers
+
+      joins(runner_managers_join).merge(RunnerManager.with_version_prefix(value))
+    }
     scope :with_runner_type, ->(runner_type) do
       return all if AVAILABLE_TYPES.exclude?(runner_type.to_s)
 
@@ -363,10 +375,22 @@ def self.taggings_join_model
       ::Ci::RunnerTagging
     end
 
+    def self.partitioned_runner_managers_join
+      Arel.sql(<<~SQL)
+        INNER JOIN #{RunnerManager.table_name}
+          ON #{RunnerManager.table_name}.runner_id = #{table_name}.id
+          AND #{RunnerManager.table_name}.runner_type = #{table_name}.runner_type
+      SQL
+    end
+
     def self.created_runner_prefix
       ::Authn::TokenField::PrefixHelper.prepend_instance_prefix(CREATED_RUNNER_TOKEN_PREFIX)
     end
 
+    def self.partition_pruning_enabled?
+      Feature.enabled?(:ci_runner_partition_pruning, Feature.current_request)
+    end
+
     def runner_matcher
       Gitlab::Ci::Matching::RunnerMatcher.new({
         runner_ids: [id],
@@ -564,8 +588,11 @@ def compute_token_expiration(**)
 
     def ensure_manager(system_xid)
       # rubocop: disable Performance/ActiveRecordSubtransactionMethods -- This is used only in API endpoints outside of transactions
-      RunnerManager.safe_find_or_create_by!(runner_id: id, system_xid: system_xid.to_s) do |m|
-        m.runner_type = runner_type
+      attrs = { runner_id: id, system_xid: system_xid.to_s }
+      attrs[:runner_type] = runner_type if self.class.partition_pruning_enabled?
+
+      RunnerManager.safe_find_or_create_by!(**attrs) do |m|
+        m.runner_type = runner_type unless self.class.partition_pruning_enabled?
         m.organization_id = organization_id
       end
       # rubocop: enable Performance/ActiveRecordSubtransactionMethods
@@ -596,7 +623,9 @@ def partition_id
     private
 
     scope :with_upgrade_status, ->(upgrade_status) do
-      joins(:runner_managers).merge(RunnerManager.with_upgrade_status(upgrade_status))
+      runner_managers_join = partition_pruning_enabled? ? partitioned_runner_managers_join : :runner_managers
+
+      joins(runner_managers_join).merge(RunnerManager.with_upgrade_status(upgrade_status))
     end
 
     def fallback_owner_project
diff --git a/app/models/ci/runner_manager.rb b/app/models/ci/runner_manager.rb
index 9e3e7eb914ef..d7ea7818d4f6 100644
--- a/app/models/ci/runner_manager.rb
+++ b/app/models/ci/runner_manager.rb
@@ -12,6 +12,8 @@ class RunnerManager < Ci::ApplicationRecord
     self.table_name = 'ci_runner_machines'
     self.primary_key = :id
 
+    query_constraints :runner_id, :runner_type
+
     AVAILABLE_STATUSES = %w[online offline never_contacted stale].freeze
     AVAILABLE_STATUSES_INCL_DEPRECATED = AVAILABLE_STATUSES
 
diff --git a/config/feature_flags/gitlab_com_derisk/ci_runner_partition_pruning.yml b/config/feature_flags/gitlab_com_derisk/ci_runner_partition_pruning.yml
new file mode 100644
index 000000000000..446558b03459
--- /dev/null
+++ b/config/feature_flags/gitlab_com_derisk/ci_runner_partition_pruning.yml
@@ -0,0 +1,10 @@
+---
+name: ci_runner_partition_pruning
+description: 'Add runner_type filters to Ci::Runner scopes and associations to enable partition pruning'
+feature_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/594861
+introduced_by_url:
+rollout_issue_url:
+milestone: '18.11'
+group: group::runner
+type: gitlab_com_derisk
+default_enabled: false

Problem

Ci::Runner scopes

The following scopes join against association tables that are exclusive to a specific runner type, but do not include runner_type in the WHERE clause:

Scope File Implicitly returns Missing filter
belonging_to_project app/models/ci/runner.rb project runners only .project_type
belonging_to_group app/models/ci/runner.rb group runners only .group_type

These two scopes are used as building blocks in several compound scopes, all of which would automatically benefit from this fix:

  • belonging_to_group_and_ancestors
  • belonging_to_parent_groups_of_project
  • belonging_to_group_or_project_descendants (union of belonging_to_group + belonging_to_project)
  • owned_or_instance_wide (union including belonging_to_project + belonging_to_parent_groups_of_project + shared_runners)
  • group_or_instance_wide (union including belonging_to_group_and_ancestors + shared_runners)
  • usable_from_scope (union including belonging_to_group + belonging_to_group_or_project_descendants + shared_runners)

EE callers that also benefit (no changes needed)

  • Ci::Runners::StaleGroupRunnersPruneService — calls belonging_to_group
  • Analytics::DevopsAdoption::SnapshotCalculator — calls belonging_to_group_or_project_descendants

Ci::RunnerManager (via ci_runner_machines)

The ci_runner_machines table is also LIST-partitioned by runner_type. Several queries miss the runner_type filter:

Location Issue
has_many :runner_managers association Generates WHERE runner_id = ? without runner_type; affects all callers via runner.runner_managers
with_version_prefix scope joins(:runner_managers) generates a JOIN on runner_id only, without runner_type correlation
with_upgrade_status scope Same — joins(:runner_managers) without runner_type
ensure_manager method safe_find_or_create_by! uses runner_id + system_xid without runner_type in the lookup

Proposal

1. Add type filters to Ci::Runner base scopes

scope :belonging_to_project, ->(project_id) {
  project_type
    .joins(:runner_projects)
    .where(ci_runner_projects: { project_id: project_id })
}

scope :belonging_to_group, ->(group_id) {
  group_type
    .joins(:runner_namespaces)
    .where(ci_runner_namespaces: { namespace_id: group_id })
}

2. Add query_constraints to Ci::RunnerManager

Add query_constraints :runner_id, :runner_type to Ci::RunnerManager, matching the pattern already used by Ci::RunnerTagging. This automatically includes runner_type in all association lookups via runner.runner_managers (e.g. find_by_system_xid, order_contacted_at_desc.first, API listing).

This is safe for Ci::RunnerManagerBuild.belongs_to :runner_manager because that association uses an explicit foreign_key: :runner_machine_id pointing to RunnerManager's primary_key (:id), which is unaffected by query_constraints.

3. Add runner_type correlation to Ci::RunnerManager joins

Replace joins(:runner_managers) with an explicit Arel join that correlates on both runner_id AND runner_type:

def self.partitioned_runner_managers_join
  runner_managers_table = RunnerManager.arel_table
  join_condition = arel_table[:id].eq(runner_managers_table[:runner_id])
    .and(arel_table[:runner_type].eq(runner_managers_table[:runner_type]))

  Arel::Nodes::InnerJoin.new(runner_managers_table, Arel::Nodes::On.new(join_condition))
end

4. Add runner_type to ensure_manager find conditions

RunnerManager.safe_find_or_create_by!(runner_id: id, runner_type: runner_type, system_xid: system_xid.to_s) do |m|
  m.organization_id = organization_id
end

NOTE: This particular change should definitely be behind a FF since this is a high-risk code path.

Safety

All changes are logically safe:

  • ci_runner_namespaces only has entries for group runners
  • ci_runner_projects only has entries for project runners
  • ci_runner_machines records always share the same runner_type as their parent ci_runners record
  • Adding the type filter is redundant from a data perspective but enables partition pruning

Why not query_constraints on Ci::Runner?

Unlike RunnerManager and RunnerTagging, several models that belongs_to :runner — notably Ci::Build, Ci::RunnerNamespace, and Ci::RunnerProject — do not have a runner_type column. Adding query_constraints to Ci::Runner would break those associations.

Risk assessment

Low risk. If any caller chains these scopes on an already-typed relation (e.g., Ci::Runner.instance_type.belonging_to_group(...)), it would produce WHERE runner_type = 1 AND runner_type = 2 which returns nothing — but it already returns nothing today since instance runners have no ci_runner_namespaces entries.

Verification

  1. Existing specs should pass unchanged
  2. Partition pruning can be verified with EXPLAIN — the query plan should show only the relevant partition being scanned
Edited by 🤖 GitLab Bot 🤖