Fix performance issue finding snippet replicables

Fix a performance issue when looking up snippet repositories
when replication is enabled by namespace

Relates to https://gitlab.com/gitlab-org/gitlab/issues/287763

changelogs/unreleased/alexives-287763-fix_performance_of_snippet_repository_queries.yml

+---
+title: Update snippet repository finder for namespace replication
+merge_request: 49518
+author:
+type: performance

db/migrate/20201210175044_add_index_to_snippet_on_project_id.rb

+# frozen_string_literal: true
+
+class AddIndexToSnippetOnProjectId < ActiveRecord::Migration[6.0]
+  include Gitlab::Database::MigrationHelpers
+  INDEX_NAME = "index_snippet_on_id_and_project_id"
+  DOWNTIME = false
+
+  disable_ddl_transaction!
+
+  def up
+    add_concurrent_index :snippets, [:id, :project_id], name: INDEX_NAME
+  end
+
+  def down
+    remove_concurrent_index_by_name :snippets, INDEX_NAME
+  end
+end

db/schema_migrations/20201210175044

+32f0889266a05c12f1bba6d3d8646c6cd5c27ffcc01e4cc0cb1721c495e17237
\ No newline at end of file

db/structure.sql

@@ -22565,6 +22565,8 @@ CREATE UNIQUE INDEX index_smartcard_identities_on_subject_and_issuer ON smartcar
 
 CREATE INDEX index_smartcard_identities_on_user_id ON smartcard_identities USING btree (user_id);
 
+CREATE INDEX index_snippet_on_id_and_project_id ON snippets USING btree (id, project_id);
+
 CREATE UNIQUE INDEX index_snippet_repositories_on_disk_path ON snippet_repositories USING btree (disk_path);
 
 CREATE INDEX index_snippet_repositories_on_shard_id ON snippet_repositories USING btree (shard_id);

ee/app/models/ee/snippet.rb

@@ -8,6 +8,8 @@ module EE
     prepended do
       include Elastic::SnippetsSearch
       include UsageStatistics
+
+      scope :for_projects, ->(projects) { where(project: projects) }
     end
 
     override :repository_size_checker

ee/app/models/ee/snippet_repository.rb

@@ -17,32 +17,33 @@ module EE
       def replicables_for_current_secondary(primary_key_in)
         node = ::Gitlab::Geo.current_node
 
-        replicables = if !node.selective_sync?
-                        all
-                      elsif node.selective_sync_by_namespaces?
-                        snippet_repositories_for_selected_namespaces
-                      elsif node.selective_sync_by_shards?
-                        snippet_repositories_for_selected_shards
-                      else
-                        self.none
-                      end
-
-        replicables.primary_key_in(primary_key_in)
+        if !node.selective_sync?
+          all.primary_key_in(primary_key_in)
+        elsif node.selective_sync_by_namespaces?
+          snippet_repositories_for_selected_namespaces(primary_key_in)
+        elsif node.selective_sync_by_shards?
+          snippet_repositories_for_selected_shards(primary_key_in)
+        else
+          self.none
+        end
       end
 
-      def snippet_repositories_for_selected_namespaces
-        personal_snippets = self.joins(:snippet).where(snippet: ::Snippet.only_personal_snippets)
+      def snippet_repositories_for_selected_namespaces(primary_key_in)
+        personal_snippets = self.where(snippet: ::Snippet.only_personal_snippets.primary_key_in(primary_key_in))
 
-        project_snippets = self.joins(snippet: :project).where(
-          'snippets.project_id IN(?)',
-          ::Gitlab::Geo.current_node.projects.select(:id)
-        )
+        project_snippets = self.where(snippet: ::Snippet.for_projects(::Gitlab::Geo.current_node.projects.select(:id))
+                                                   .primary_key_in(primary_key_in))
 
-        self.from_union([project_snippets, personal_snippets])
+        # We use `find_by_sql` here in order to perform the union operation and cast the results as
+        # Snippet repositories. If we use a `from_union`, it wraps the query in in a sub-select and
+        # it increases the query time by a surprising amount.
+        self.find_by_sql(::Gitlab::SQL::Union.new([project_snippets, personal_snippets]).to_sql)
       end
 
-      def snippet_repositories_for_selected_shards
-        self.for_repository_storage(::Gitlab::Geo.current_node.selective_sync_shards)
+      def snippet_repositories_for_selected_shards(primary_key_in)
+        self
+          .for_repository_storage(::Gitlab::Geo.current_node.selective_sync_shards)
+          .primary_key_in(primary_key_in)
       end
     end