Merge branch 'optimise-owned-runners-index' into 'master'

Make `User#ci_owned_runners` to use unnest index instead of GIN

See merge request gitlab-org/gitlab!83843

app/models/ci/namespace_mirror.rb

@@ -4,6 +4,8 @@ module Ci
   # This model represents a record in a shadow table of the main database's namespaces table.
   # It allows us to navigate the namespace hierarchy on the ci database without resorting to a JOIN.
   class NamespaceMirror < ApplicationRecord
+    include FromUnion
+
     belongs_to :namespace
 
     scope :by_group_and_descendants, -> (id) do
@@ -14,6 +16,24 @@ module Ci
       where('traversal_ids && ARRAY[?]::int[]', ids)
     end
 
+    scope :contains_traversal_ids, -> (traversal_ids) do
+      mirrors = []
+
+      traversal_ids.group_by(&:count).each do |columns_count, traversal_ids_group|
+        columns = Array.new(columns_count) { |i| "(traversal_ids[#{i + 1}])" }
+        pairs = traversal_ids_group.map do |ids|
+          ids = ids.map { |id| Arel::Nodes.build_quoted(id).to_sql }
+          "(#{ids.join(",")})"
+        end
+
+        # Create condition in format:
+        # ((traversal_ids[1]),(traversal_ids[2])) IN ((1,2),(2,3))
+        mirrors << Ci::NamespaceMirror.where("(#{columns.join(",")}) IN (#{pairs.join(",")})") # rubocop:disable GitlabSecurity/SqlInjection
+      end
+
+      self.from_union(mirrors)
+    end
+
     scope :by_namespace_id, -> (namespace_id) { where(namespace_id: namespace_id) }
 
     class << self

app/models/concerns/from_set_operator.rb

@@ -11,7 +11,12 @@ module FromSetOperator
     raise "Trying to redefine method '#{method(method_name)}'" if methods.include?(method_name)
 
     define_method(method_name) do |members, remove_duplicates: true, remove_order: true, alias_as: table_name|
-      operator_sql = operator.new(members, remove_duplicates: remove_duplicates, remove_order: remove_order).to_sql
+      operator_sql =
+        if members.any?
+          operator.new(members, remove_duplicates: remove_duplicates, remove_order: remove_order).to_sql
+        else
+          where("1=0").to_sql
+        end
 
       from(Arel.sql("(#{operator_sql}) #{alias_as}"))
     end

app/models/user.rb

@@ -2286,17 +2286,20 @@ class User < ApplicationRecord
     # to avoid querying descendants since they are already covered
     # by ancestor namespaces. If the FF is not available fallback to
     # inefficient search: https://gitlab.com/gitlab-org/gitlab/-/issues/336436
-    namespace_ids =
-      if Feature.enabled?(:use_traversal_ids, default_enabled: :yaml)
-        Group.joins(:all_group_members)
-          .merge(search_members)
-          .shortest_traversal_ids_prefixes
-          .map(&:last)
-      else
-        search_members.pluck(:source_id)
-      end
+    unless Feature.enabled?(:use_traversal_ids, default_enabled: :yaml)
+      return Ci::NamespaceMirror.contains_any_of_namespaces(search_members.pluck(:source_id))
+    end
 
-    Ci::NamespaceMirror.contains_any_of_namespaces(namespace_ids)
+    traversal_ids = Group.joins(:all_group_members)
+      .merge(search_members)
+      .shortest_traversal_ids_prefixes
+
+    # Use efficient btree index to perform search
+    if Feature.enabled?(:ci_owned_runners_unnest_index, self, default_enabled: :yaml)
+      Ci::NamespaceMirror.contains_traversal_ids(traversal_ids)
+    else
+      Ci::NamespaceMirror.contains_any_of_namespaces(traversal_ids.map(&:last))
+    end
   end
 end
 

config/feature_flags/development/ci_owned_runners_unnest_index.yml

+---
+name: ci_owned_runners_unnest_index
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/83843
+rollout_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/357869
+milestone: '14.10'
+type: development
+group: group::sharding
+default_enabled: false

db/migrate/20220329110630_add_ci_namespace_mirrors_unnest_index_on_traversal_ids.rb

+# frozen_string_literal: true
+
+class AddCiNamespaceMirrorsUnnestIndexOnTraversalIds < Gitlab::Database::Migration[1.0]
+  disable_ddl_transaction!
+
+  INDEX_NAME = 'index_ci_namespace_mirrors_on_traversal_ids_unnest'
+
+  def up
+    return if index_exists_by_name?(:ci_namespace_mirrors, INDEX_NAME)
+
+    # We add only 4-levels since on average it is not expected that namespaces
+    # will be so granular beyond this point
+    disable_statement_timeout do
+      execute <<-SQL
+        CREATE INDEX CONCURRENTLY #{INDEX_NAME} ON ci_namespace_mirrors
+          USING btree ((traversal_ids[1]), (traversal_ids[2]), (traversal_ids[3]), (traversal_ids[4]))
+          INCLUDE (traversal_ids, namespace_id)
+      SQL
+    end
+  end
+
+  def down
+    remove_concurrent_index_by_name(:ci_namespace_mirrors, INDEX_NAME)
+  end
+end

db/schema_migrations/20220329110630

+47664025bee8b3728411bde53aa597cb962d859402b0504a1962a3e4f117782f
\ No newline at end of file

db/structure.sql

@@ -27035,6 +27035,8 @@ CREATE INDEX index_ci_minutes_additional_packs_on_namespace_id_purchase_xid ON c
 
 CREATE UNIQUE INDEX index_ci_namespace_mirrors_on_namespace_id ON ci_namespace_mirrors USING btree (namespace_id);
 
+CREATE INDEX index_ci_namespace_mirrors_on_traversal_ids_unnest ON ci_namespace_mirrors USING btree ((traversal_ids[1]), (traversal_ids[2]), (traversal_ids[3]), (traversal_ids[4])) INCLUDE (traversal_ids, namespace_id);
+
 CREATE UNIQUE INDEX index_ci_namespace_monthly_usages_on_namespace_id_and_date ON ci_namespace_monthly_usages USING btree (namespace_id, date);
 
 CREATE INDEX index_ci_pending_builds_id_on_protected_partial ON ci_pending_builds USING btree (id) WHERE (protected = true);

spec/models/ci/namespace_mirror_spec.rb

@@ -44,6 +44,53 @@ RSpec.describe Ci::NamespaceMirror do
       end
     end
 
+    describe '.contains_traversal_ids' do
+      let!(:other_group1) { create(:group) }
+      let!(:other_group2) { create(:group, parent: other_group1) }
+      let!(:other_group3) { create(:group, parent: other_group2) }
+      let!(:other_group4) { create(:group) }
+
+      subject(:result) { described_class.contains_traversal_ids(all_traversal_ids) }
+
+      context 'when passing a top-level group' do
+        let(:all_traversal_ids) do
+          [
+            [other_group1.id]
+          ]
+        end
+
+        it 'returns only itself and children of that group' do
+          expect(result.map(&:namespace)).to contain_exactly(other_group1, other_group2, other_group3)
+        end
+      end
+
+      context 'when passing many levels of groups' do
+        let(:all_traversal_ids) do
+          [
+            [other_group2.parent_id, other_group2.id],
+            [other_group3.parent_id, other_group3.id],
+            [other_group4.id]
+          ]
+        end
+
+        it 'returns only the asked group' do
+          expect(result.map(&:namespace)).to contain_exactly(other_group2, other_group3, other_group4)
+        end
+      end
+
+      context 'when passing invalid data ' do
+        let(:all_traversal_ids) do
+          [
+            ["; UPDATE"]
+          ]
+        end
+
+        it 'data is properly sanitised' do
+          expect(result.to_sql).to include "((traversal_ids[1])) IN (('; UPDATE'))"
+        end
+      end
+    end
+
     describe '.by_namespace_id' do
       subject(:result) { described_class.by_namespace_id(group2.id) }
 

spec/models/user_spec.rb

@@ -4271,6 +4271,14 @@ RSpec.describe User do
 
       it_behaves_like '#ci_owned_runners'
     end
+
+    context 'when FF ci_owned_runners_unnest_index is disabled uses GIN index' do
+      before do
+        stub_feature_flags(ci_owned_runners_unnest_index: false)
+      end
+
+      it_behaves_like '#ci_owned_runners'
+    end
   end
 
   describe '#projects_with_reporter_access_limited_to' do

spec/support/shared_examples/models/concerns/from_set_operator_shared_examples.rb

@@ -20,6 +20,12 @@ RSpec.shared_examples 'from set operator' do |sql_klass|
       expect(query.to_sql).to match(/FROM \(\(SELECT.+\)\n#{operator_keyword}\n\(SELECT.+\)\) users/m)
     end
 
+    it "returns empty set when passing empty array" do
+      query = model.public_send(operator_method, [])
+
+      expect(query.to_sql).to match(/WHERE \(1=0\)/m)
+    end
+
     it 'supports the use of a custom alias for the sub query' do
       query = model.public_send(operator_method,
         [model.where(id: 1), model.where(id: 2)],