Rewrite SnippetsFinder to improve performance

This completely rewrites the SnippetsFinder class from the ground up in
order to improve its performance. The old code was beyond salvaging. It
was complex, included various Rails 5 workarounds, comments that
shouldn't be necessary, and most important of all: it produced a really
poorly performing database query.

As a result, I opted for rewriting the finder from scratch, instead of
trying to patch the existing code. Instead of trying to reuse as many
existing methods as possible, I opted for defining new methods
specifically meant for the SnippetsFinder. This requires some extra code
here and there, but allows us to have much more control over the
resulting SQL queries. It is these changes that then allow us to produce
a _much_ more efficient query.

To illustrate how bad the old query was, we will use my own snippets as
an example. Currently I have 52 snippets, most of which are global ones.
To retrieve these, you would run the following Ruby code:

    user = User.find_by(username: 'yorickpeterse')

    SnippetsFinder.new(user, author: user).execute

On GitLab.com the resulting query will take between 10 and 15 seconds to
run, producing the query plan found at
https://explain.depesz.com/s/Y5IX. Apart from the long execution time,
the total number of buffers (the sum of all shared hits) is around 185
GB, though the real number is probably (hopefully) much lower as I doubt
simply summing these numbers produces the true total number of buffers
used.

The new query's plan can be found at https://explain.depesz.com/s/wHdN,
and this query takes between 10 and 100-ish milliseconds to run. The
total number of buffers used is only about 30 MB.

Fixes https://gitlab.com/gitlab-org/gitlab-ce/issues/52639

app/finders/snippets_finder.rb

 # frozen_string_literal: true
 
-# Snippets Finder
+# Finder for retrieving snippets that a user can see, optionally scoped to a
+# project or snippets author.
 #
-# Used to filter Snippets collections by a set of params
+# Basic usage:
 #
-# Arguments.
+#     user = User.find(1)
 #
-# current_user - The current user, nil also can be used.
-# params:
-#   visibility (integer) - Individual snippet visibility: Public(20), internal(10) or private(0).
-#   project (Project) - Project related.
-#   author (User) - Author related.
+#     SnippetsFinder.new(user).execute
 #
-# params are optional
+# To limit the snippets to a specific project, supply the `project:` option:
+#
+#     user = User.find(1)
+#     project = Project.find(1)
+#
+#     SnippetsFinder.new(user, project: project).execute
+#
+# Limiting snippets to an author can be done by supplying the `author:` option:
+#
+#     user = User.find(1)
+#     project = Project.find(1)
+#
+#     SnippetsFinder.new(user, author: user).execute
+#
+# To filter snippets using a specific visibility level, you can provide the
+# `scope:` option:
+#
+#     user = User.find(1)
+#     project = Project.find(1)
+#
+#     SnippetsFinder.new(user, author: user, scope: :are_public).execute
+#
+# Valid `scope:` values are:
+#
+# * `:are_private`
+# * `:are_internal`
+# * `:are_public`
+#
+# Any other value will be ignored.
 class SnippetsFinder < UnionFinder
-  include Gitlab::Allowable
   include FinderMethods
 
-  attr_accessor :current_user, :project, :params
+  attr_accessor :current_user, :project, :author, :scope
 
-  def initialize(current_user, params = {})
+  def initialize(current_user = nil, params = {})
     @current_user = current_user
-    @params = params
     @project = params[:project]
-  end
-
-  def execute
-    items = init_collection
-    items = by_author(items)
-    items = by_visibility(items)
-
-    items.fresh
-  end
-
-  private
-
-  def init_collection
-    if project.present?
-      authorized_snippets_from_project
-    else
-      authorized_snippets
+    @author = params[:author]
+    @scope = params[:scope].to_s
+
+    if project && author
+      raise(
+        ArgumentError,
+        'Filtering by both an author and a project is not supported, ' \
+          'as this finder is not optimised for this use case'
+      )
     end
   end
 
-  def authorized_snippets_from_project
-    if can?(current_user, :read_project_snippet, project)
-      if project.team.member?(current_user)
-        project.snippets
+  def execute
+    base =
+      if project
+        snippets_for_a_single_project
       else
-        project.snippets.public_to_user(current_user)
+        snippets_for_multiple_projects
       end
-    else
-      Snippet.none
-    end
-  end
 
-  # rubocop: disable CodeReuse/ActiveRecord
-  def authorized_snippets
-    # This query was intentionally converted to a raw one to get it work in Rails 5.0.
-    # In Rails 5.0 and 5.1 there's a bug: https://github.com/rails/arel/issues/531
-    # Please convert it back when on rails 5.2 as it works again as expected since 5.2.
-    Snippet.where("#{feature_available_projects} OR #{not_project_related}")
-      .public_or_visible_to_user(current_user)
+    base.with_optional_visibility(visibility_from_scope).fresh
   end
-  # rubocop: enable CodeReuse/ActiveRecord
 
-  # Returns a collection of projects that is either public or visible to the
-  # logged in user.
+  # Produces a query that retrieves snippets from multiple projects.
   #
-  # A caller must pass in a block to modify individual parts of
-  # the query, e.g. to apply .with_feature_available_for_user on top of it.
-  # This is useful for performance as we can stick those additional filters
-  # at the bottom of e.g. the UNION.
-  # rubocop: disable CodeReuse/ActiveRecord
-  def projects_for_user
-    return yield(Project.public_to_user) unless current_user
-
-    # If the current_user is allowed to see all projects,
-    # we can shortcut and just return.
-    return yield(Project.all) if current_user.full_private_access?
-
-    authorized_projects = yield(Project.where('EXISTS (?)', current_user.authorizations_for_projects))
-
-    levels = Gitlab::VisibilityLevel.levels_for_user(current_user)
-    visible_projects = yield(Project.where(visibility_level: levels))
-
-    # We use a UNION here instead of OR clauses since this results in better
-    # performance.
-    Project.from_union([authorized_projects, visible_projects])
-  end
-  # rubocop: enable CodeReuse/ActiveRecord
-
-  def feature_available_projects
-    # Don't return any project related snippets if the user cannot read cross project
-    return table[:id].eq(nil).to_sql unless Ability.allowed?(current_user, :read_cross_project)
-
-    projects = projects_for_user do |part|
-      part.with_feature_available_for_user(:snippets, current_user)
-    end.select(:id)
+  # The resulting query will, depending on the user's permissions, include the
+  # following collections of snippets:
+  #
+  # 1. Snippets that don't belong to any project.
+  # 2. Snippets of projects that are visible to the current user (e.g. snippets
+  #    in public projects).
+  # 3. Snippets of projects that the current user is a member of.
+  #
+  # Each collection is constructed in isolation, allowing for greater control
+  # over the resulting SQL query.
+  def snippets_for_multiple_projects
+    queries = [global_snippets]
+
+    if Ability.allowed?(current_user, :read_cross_project)
+      queries << snippets_of_visible_projects
+      queries << snippets_of_authorized_projects if current_user
+    end
 
-    # This query was intentionally converted to a raw one to get it work in Rails 5.0.
-    # In Rails 5.0 and 5.1 there's a bug: https://github.com/rails/arel/issues/531
-    # Please convert it back when on rails 5.2 as it works again as expected since 5.2.
-    "snippets.project_id IN (#{projects.to_sql})"
+    find_union(queries, Snippet)
   end
 
-  def not_project_related
-    table[:project_id].eq(nil).to_sql
+  def snippets_for_a_single_project
+    Snippet.for_project_with_user(project, current_user)
   end
 
-  def table
-    Snippet.arel_table
+  def global_snippets
+    snippets_for_author_or_visible_to_user.only_global_snippets
   end
 
-  # rubocop: disable CodeReuse/ActiveRecord
-  def by_visibility(items)
-    visibility = params[:visibility] || visibility_from_scope
+  # Returns the snippets that the current user (logged in or not) can view.
+  def snippets_of_visible_projects
+    snippets_for_author_or_visible_to_user
+      .only_include_projects_visible_to(current_user)
+      .only_include_projects_with_snippets_enabled
+  end
 
-    return items unless visibility
+  # Returns the snippets that the currently logged in user has access to by
+  # being a member of the project the snippets belong to.
+  #
+  # This method requires that `current_user` returns a `User` instead of `nil`,
+  # and is optimised for this specific scenario.
+  def snippets_of_authorized_projects
+    base = author ? snippets_for_author : Snippet.all
+
+    base
+      .only_include_projects_with_snippets_enabled(include_private: true)
+      .only_include_authorized_projects(current_user)
+  end
 
-    items.where(visibility_level: visibility)
+  def snippets_for_author_or_visible_to_user
+    if author
+      snippets_for_author
+    elsif current_user
+      Snippet.visible_to_or_authored_by(current_user)
+    else
+      Snippet.public_to_user
+    end
   end
-  # rubocop: enable CodeReuse/ActiveRecord
 
-  # rubocop: disable CodeReuse/ActiveRecord
-  def by_author(items)
-    return items unless params[:author]
+  def snippets_for_author
+    base = author.snippets
 
-    items.where(author_id: params[:author].id)
+    if author == current_user
+      # If the current user is also the author of all snippets, then we can
+      # include private snippets.
+      base
+    else
+      base.public_to_user(current_user)
+    end
   end
-  # rubocop: enable CodeReuse/ActiveRecord
 
   def visibility_from_scope
-    case params[:scope].to_s
+    case scope
     when 'are_private'
       Snippet::PRIVATE
     when 'are_internal'

app/models/project.rb

@@ -2073,6 +2073,10 @@ class Project < ActiveRecord::Base
     storage_version != LATEST_STORAGE_VERSION
   end
 
+  def snippets_visible?(user = nil)
+    Ability.allowed?(user, :read_project_snippet, self)
+  end
+
   private
 
   def use_hashed_storage

app/models/snippet.rb

@@ -63,6 +63,62 @@ class Snippet < ActiveRecord::Base
   attr_spammable :title, spam_title: true
   attr_spammable :content, spam_description: true
 
+  def self.with_optional_visibility(value = nil)
+    if value
+      where(visibility_level: value)
+    else
+      all
+    end
+  end
+
+  def self.only_global_snippets
+    where(project_id: nil)
+  end
+
+  def self.only_include_projects_visible_to(current_user = nil)
+    levels = Gitlab::VisibilityLevel.levels_for_user(current_user)
+
+    joins(:project).where('projects.visibility_level IN (?)', levels)
+  end
+
+  def self.only_include_projects_with_snippets_enabled(include_private: false)
+    column = ProjectFeature.access_level_attribute(:snippets)
+    levels = [ProjectFeature::ENABLED, ProjectFeature::PUBLIC]
+
+    levels << ProjectFeature::PRIVATE if include_private
+
+    joins(project: :project_feature)
+      .where(project_features: { column => levels })
+  end
+
+  def self.only_include_authorized_projects(current_user)
+    where(
+      'EXISTS (?)',
+      ProjectAuthorization
+        .select(1)
+        .where('project_id = snippets.project_id')
+        .where(user_id: current_user.id)
+    )
+  end
+
+  def self.for_project_with_user(project, user = nil)
+    return none unless project.snippets_visible?(user)
+
+    if user && project.team.member?(user)
+      project.snippets
+    else
+      project.snippets.public_to_user(user)
+    end
+  end
+
+  def self.visible_to_or_authored_by(user)
+    where(
+      'snippets.visibility_level IN (?) OR snippets.author_id = ?',
+      Gitlab::VisibilityLevel.levels_for_user(user),
+      user.id
+    )
+  end
+
   def self.reference_prefix
     '$'
   end
@@ -81,27 +137,6 @@ class Snippet < ActiveRecord::Base
     @link_reference_pattern ||= super("snippets", /(?<snippet>\d+)/)
   end
 
-  # Returns a collection of snippets that are either public or visible to the
-  # logged in user.
-  #
-  # This method does not verify the user actually has the access to the project
-  # the snippet is in, so it should be only used on a relation that's already scoped
-  # for project access
-  def self.public_or_visible_to_user(user = nil)
-    if user
-      authorized = user
-        .project_authorizations
-        .select(1)
-        .where('project_authorizations.project_id = snippets.project_id')
-
-      levels = Gitlab::VisibilityLevel.levels_for_user(user)
-
-      where('EXISTS (?) OR snippets.visibility_level IN (?) or snippets.author_id = (?)', authorized, levels, user.id)
-    else
-      public_to_user
-    end
-  end
-
   def to_reference(from = nil, full: false)
     reference = "#{self.class.reference_prefix}#{id}"
 

changelogs/unreleased/refactor-snippets-finder.yml

+---
+title: Rewrite SnippetsFinder to improve performance by a factor of 1500
+merge_request:
+author:
+type: performance

lib/api/snippets.rb

@@ -14,7 +14,7 @@ module API
         end
 
         def public_snippets
-          SnippetsFinder.new(current_user, visibility: Snippet::PUBLIC).execute
+          SnippetsFinder.new(current_user, scope: :are_public).execute
         end
       end
 

spec/finders/snippets_finder_spec.rb

@@ -4,16 +4,13 @@ describe SnippetsFinder do
   include Gitlab::Allowable
   using RSpec::Parameterized::TableSyntax
 
-  context 'filter by visibility' do
-    let!(:snippet1) { create(:personal_snippet, :private) }
-    let!(:snippet2) { create(:personal_snippet, :internal) }
-    let!(:snippet3) { create(:personal_snippet, :public) }
+  describe '#initialize' do
+    it 'raises ArgumentError when a project and author are given' do
+      user = build(:user)
+      project = build(:project)
 
-    it "returns public snippets when visibility is PUBLIC" do
-      snippets = described_class.new(nil, visibility: Snippet::PUBLIC).execute
-
-      expect(snippets).to include(snippet3)
-      expect(snippets).not_to include(snippet1, snippet2)
+      expect { described_class.new(user, author: user, project: project) }
+        .to raise_error(ArgumentError)
     end
   end
 
@@ -66,21 +63,21 @@ describe SnippetsFinder do
     end
 
     it "returns internal snippets" do
-      snippets = described_class.new(user, author: user, visibility: Snippet::INTERNAL).execute
+      snippets = described_class.new(user, author: user, scope: :are_internal).execute
 
       expect(snippets).to include(snippet2)
       expect(snippets).not_to include(snippet1, snippet3)
     end
 
     it "returns private snippets" do
-      snippets = described_class.new(user, author: user, visibility: Snippet::PRIVATE).execute
+      snippets = described_class.new(user, author: user, scope: :are_private).execute
 
       expect(snippets).to include(snippet1)
       expect(snippets).not_to include(snippet2, snippet3)
     end
 
     it "returns public snippets" do
-      snippets = described_class.new(user, author: user, visibility: Snippet::PUBLIC).execute
+      snippets = described_class.new(user, author: user, scope: :are_public).execute
 
       expect(snippets).to include(snippet3)
       expect(snippets).not_to include(snippet1, snippet2)
@@ -98,6 +95,13 @@ describe SnippetsFinder do
       expect(snippets).to include(snippet3)
       expect(snippets).not_to include(snippet2, snippet1)
     end
+
+    it 'returns all snippets for an admin' do
+      admin = create(:user, :admin)
+      snippets = described_class.new(admin, author: user).execute
+
+      expect(snippets).to include(snippet1, snippet2, snippet3)
+    end
   end
 
   context 'filter by project' do
@@ -126,21 +130,21 @@ describe SnippetsFinder do
     end
 
     it "returns public snippets for non project members" do
-      snippets = described_class.new(user, project: project1, visibility: Snippet::PUBLIC).execute
+      snippets = described_class.new(user, project: project1, scope: :are_public).execute
 
       expect(snippets).to include(@snippet3)
       expect(snippets).not_to include(@snippet1, @snippet2)
     end
 
     it "returns internal snippets for non project members" do
-      snippets = described_class.new(user, project: project1, visibility: Snippet::INTERNAL).execute
+      snippets = described_class.new(user, project: project1, scope: :are_internal).execute
 
       expect(snippets).to include(@snippet2)
       expect(snippets).not_to include(@snippet1, @snippet3)
     end
 
     it "does not return private snippets for non project members" do
-      snippets = described_class.new(user, project: project1, visibility: Snippet::PRIVATE).execute
+      snippets = described_class.new(user, project: project1, scope: :are_private).execute
 
       expect(snippets).not_to include(@snippet1, @snippet2, @snippet3)
     end
@@ -156,10 +160,17 @@ describe SnippetsFinder do
     it "returns private snippets for project members" do
       project1.add_developer(user)
 
-      snippets = described_class.new(user, project: project1, visibility: Snippet::PRIVATE).execute
+      snippets = described_class.new(user, project: project1, scope: :are_private).execute
 
       expect(snippets).to include(@snippet1)
     end
+
+    it 'returns all snippets for an admin' do
+      admin = create(:user, :admin)
+      snippets = described_class.new(admin, project: project1).execute
+
+      expect(snippets).to include(@snippet1, @snippet2, @snippet3)
+    end
   end
 
   describe '#execute' do
@@ -184,4 +195,6 @@ describe SnippetsFinder do
       end
     end
   end
+
+  it_behaves_like 'snippet visibility'
 end

spec/models/project_spec.rb

@@ -3963,6 +3963,28 @@ describe Project do
     end
   end
 
+  describe '#snippets_visible?' do
+    it 'returns true when a logged in user can read snippets' do
+      project = create(:project, :public)
+      user = create(:user)
+
+      expect(project.snippets_visible?(user)).to eq(true)
+    end
+
+    it 'returns true when an anonymous user can read snippets' do
+      project = create(:project, :public)
+
+      expect(project.snippets_visible?).to eq(true)
+    end
+
+    it 'returns false when a user can not read snippets' do
+      project = create(:project, :private)
+      user = create(:user)
+
+      expect(project.snippets_visible?(user)).to eq(false)
+    end
+  end
+
   def rugged_config
     rugged_repo(project.repository).config
   end

spec/models/snippet_spec.rb

@@ -131,6 +131,217 @@ describe Snippet do
     end
   end
 
+  describe '.with_optional_visibility' do
+    context 'when a visibility level is provided' do
+      it 'returns snippets with the given visibility' do
+        create(:snippet, :private)
+
+        snippet = create(:snippet, :public)
+        snippets = described_class
+          .with_optional_visibility(Gitlab::VisibilityLevel::PUBLIC)
+
+        expect(snippets).to eq([snippet])
+      end
+    end
+
+    context 'when a visibility level is not provided' do
+      it 'returns all snippets' do
+        snippet1 = create(:snippet, :public)
+        snippet2 = create(:snippet, :private)
+        snippets = described_class.with_optional_visibility
+
+        expect(snippets).to include(snippet1, snippet2)
+      end
+    end
+  end
+
+  describe '.only_global_snippets' do
+    it 'returns snippets not associated with any projects' do
+      create(:project_snippet)
+
+      snippet = create(:snippet)
+      snippets = described_class.only_global_snippets
+
+      expect(snippets).to eq([snippet])
+    end
+  end
+
+  describe '.only_include_projects_visible_to' do
+    let!(:project1) { create(:project, :public) }
+    let!(:project2) { create(:project, :internal) }
+    let!(:project3) { create(:project, :private) }
+    let!(:snippet1) { create(:project_snippet, project: project1) }
+    let!(:snippet2) { create(:project_snippet, project: project2) }
+    let!(:snippet3) { create(:project_snippet, project: project3) }
+
+    context 'when a user is provided' do
+      it 'returns snippets visible to the user' do
+        user = create(:user)
+
+        snippets = described_class.only_include_projects_visible_to(user)
+
+        expect(snippets).to include(snippet1, snippet2)
+        expect(snippets).not_to include(snippet3)
+      end
+    end
+
+    context 'when a user is not provided' do
+      it 'returns snippets visible to anonymous users' do
+        snippets = described_class.only_include_projects_visible_to
+
+        expect(snippets).to include(snippet1)
+        expect(snippets).not_to include(snippet2, snippet3)
+      end
+    end
+  end
+
+  describe 'only_include_projects_with_snippets_enabled' do
+    context 'when the include_private option is enabled' do
+      it 'includes snippets for projects with snippets set to private' do
+        project = create(:project)
+
+        project.project_feature
+          .update(snippets_access_level: ProjectFeature::PRIVATE)
+
+        snippet = create(:project_snippet, project: project)
+
+        snippets = described_class
+          .only_include_projects_with_snippets_enabled(include_private: true)
+
+        expect(snippets).to eq([snippet])
+      end
+    end
+
+    context 'when the include_private option is not enabled' do
+      it 'does not include snippets for projects that have snippets set to private' do
+        project = create(:project)
+
+        project.project_feature
+          .update(snippets_access_level: ProjectFeature::PRIVATE)
+
+        create(:project_snippet, project: project)
+
+        snippets = described_class.only_include_projects_with_snippets_enabled
+
+        expect(snippets).to be_empty
+      end
+    end
+
+    it 'includes snippets for projects with snippets enabled' do
+      project = create(:project)
+
+      project.project_feature
+        .update(snippets_access_level: ProjectFeature::ENABLED)
+
+      snippet = create(:project_snippet, project: project)
+      snippets = described_class.only_include_projects_with_snippets_enabled
+
+      expect(snippets).to eq([snippet])
+    end
+  end
+
+  describe '.only_include_authorized_projects' do
+    it 'only includes snippets for projects the user is authorized to see' do
+      user = create(:user)
+      project1 = create(:project, :private)
+      project2 = create(:project, :private)
+
+      project1.team.add_developer(user)
+
+      create(:project_snippet, project: project2)
+
+      snippet = create(:project_snippet, project: project1)
+      snippets = described_class.only_include_authorized_projects(user)
+
+      expect(snippets).to eq([snippet])
+    end
+  end
+
+  describe '.for_project_with_user' do
+    context 'when a user is provided' do
+      it 'returns an empty collection if the user can not view the snippets' do
+        project = create(:project, :private)
+        user = create(:user)
+
+        project.project_feature
+          .update(snippets_access_level: ProjectFeature::ENABLED)
+
+        create(:project_snippet, :public, project: project)
+
+        expect(described_class.for_project_with_user(project, user)).to be_empty
+      end
+
+      it 'returns the snippets if the user is a member of the project' do
+        project = create(:project, :private)
+        user = create(:user)
+        snippet = create(:project_snippet, project: project)
+
+        project.team.add_developer(user)
+
+        snippets = described_class.for_project_with_user(project, user)
+
+        expect(snippets).to eq([snippet])
+      end
+
+      it 'returns public snippets for a public project the user is not a member of' do
+        project = create(:project, :public)
+
+        project.project_feature
+          .update(snippets_access_level: ProjectFeature::ENABLED)
+
+        user = create(:user)
+        snippet = create(:project_snippet, :public, project: project)
+
+        create(:project_snippet, :private, project: project)
+
+        snippets = described_class.for_project_with_user(project, user)
+
+        expect(snippets).to eq([snippet])
+      end
+    end
+
+    context 'when a user is not provided' do
+      it 'returns an empty collection for a private project' do
+        project = create(:project, :private)
+
+        project.project_feature
+          .update(snippets_access_level: ProjectFeature::ENABLED)
+
+        create(:project_snippet, :public, project: project)
+
+        expect(described_class.for_project_with_user(project)).to be_empty
+      end
+
+      it 'returns public snippets for a public project' do
+        project = create(:project, :public)
+        snippet = create(:project_snippet, :public, project: project)
+
+        project.project_feature
+          .update(snippets_access_level: ProjectFeature::PUBLIC)
+
+        create(:project_snippet, :private, project: project)
+
+        snippets = described_class.for_project_with_user(project)
+
+        expect(snippets).to eq([snippet])
+      end
+    end
+  end
+
+  describe '.visible_to_or_authored_by' do
+    it 'returns snippets visible to the user' do
+      user = create(:user)
+      snippet1 = create(:snippet, :public)
+      snippet2 = create(:snippet, :private, author: user)
+      snippet3 = create(:snippet, :private)
+
+      snippets = described_class.visible_to_or_authored_by(user)
+
+      expect(snippets).to include(snippet1, snippet2)
+      expect(snippets).not_to include(snippet3)
+    end
+  end
+
   describe '#participants' do
     let(:project) { create(:project, :public) }
     let(:snippet) { create(:snippet, content: 'foo', project: project) }