Preload path locks for TreeSummary

This will allow us to avoid N + 1 request

ee/app/models/path_lock.rb

@@ -9,6 +9,8 @@ class PathLock < ApplicationRecord
   validates :path, presence: true, uniqueness: { scope: :project_id }
   validate :path_unique_validation
 
+  scope :for_paths, ->(paths) { where(path: paths) }
+
   def downstream?(path)
     self.path.start_with?(path) && !exact?(path)
   end

ee/changelogs/unreleased/id-preload-path-locks.yml

+---
+title: Preload path locks for TreeSummary
+merge_request: 29949
+author:
+type: performance

ee/lib/ee/gitlab/tree_summary.rb

@@ -17,12 +17,16 @@ module EE
 
       private
 
-      # FIXME: Loading the path locks from the database is an N+1 problem
-      # https://gitlab.com/gitlab-org/gitlab/issues/7481
       def fill_path_locks!(entries)
+        return unless project.feature_available?(:file_locks)
+
+        finder = ::Gitlab::PathLocksFinder.new(project)
+        paths = entries.map { |entry| entry_path(entry) }
+        finder.preload_for_paths(paths)
+
         entries.each do |entry|
           path = entry_path(entry)
-          path_lock = project.find_path_lock(path)
+          path_lock = finder.find_by_path(path)
 
           entry[:lock_label] = path_lock && text_label_for_lock(path_lock, path)
         end

ee/lib/gitlab/path_locks_finder.rb

@@ -35,6 +35,22 @@ class Gitlab::PathLocksFinder
     end
   end
 
+  def preload_for_paths(paths)
+    paths.each do |path|
+      tokenize(path).each do |token|
+        lazy_find_by_token(token)
+      end
+    end
+  end
+
+  def find_by_path(path)
+    tokenize(path).find do |token|
+      if lock = lazy_find_by_token(token)&.itself
+        break lock
+      end
+    end
+  end
+
   private
 
   # This returns hierarchy tokens for path
@@ -67,6 +83,14 @@ class Gitlab::PathLocksFinder
   end
   # rubocop: enable CodeReuse/ActiveRecord
 
+  def lazy_find_by_token(token)
+    BatchLoader.for(token).batch do |tokens, loader|
+      @project.path_locks.for_paths(tokens).each do |path_lock|
+        loader.call(path_lock.path, path_lock)
+      end
+    end
+  end
+
   # rubocop: disable CodeReuse/ActiveRecord
   def find_downstream(path)
     @project.path_locks.find_by("path LIKE ?", "#{sanitize_sql_like(path)}%")

ee/spec/lib/gitlab/path_locks_finder_spec.rb

@@ -3,17 +3,32 @@
 require 'spec_helper'
 
 describe Gitlab::PathLocksFinder do
-  let(:project) { create :project }
-  let(:user) { create :user }
+  let_it_be(:project) { create :project }
+  let_it_be(:user) { create :user }
+  let_it_be(:lock1) { create(:path_lock, project: project, path: 'app') }
+  let_it_be(:lock2) { create :path_lock, project: project, path: 'lib/gitlab/repo.rb' }
+
   let(:finder) { described_class.new(project) }
 
   it "returns correct lock information" do
-    lock1 = create :path_lock, project: project, path: 'app'
-    lock2 = create :path_lock, project: project, path: 'lib/gitlab/repo.rb'
-
     expect(finder.find('app')).to eq(lock1)
     expect(finder.find('app/models/project.rb')).to eq(lock1)
     expect(finder.find('lib')).to be_falsey
     expect(finder.find('lib/gitlab/repo.rb')).to eq(lock2)
   end
+
+  describe '#preload_for_paths' do
+    it 'does not perform N + 1 requests' do
+      finder.preload_for_paths(['app/models/project.rb', 'lib/gitlab/repo.rb'])
+
+      count = ActiveRecord::QueryRecorder.new do
+        expect(finder.find_by_path('app')).to eq(lock1)
+        expect(finder.find_by_path('app/models/project.rb')).to eq(lock1)
+        expect(finder.find_by_path('lib')).to be_falsey
+        expect(finder.find_by_path('lib/gitlab/repo.rb')).to eq(lock2)
+      end.count
+
+      expect(count).to eq(1)
+    end
+  end
 end

ee/spec/lib/gitlab/tree_summary_spec.rb

@@ -3,17 +3,28 @@
 require 'spec_helper'
 
 describe Gitlab::TreeSummary do
-  let(:project) { create(:project, :custom_repo, files: { 'a.txt' => '' }) }
+  let_it_be(:project) { create(:project, :custom_repo, files: { 'a.txt' => '' }) }
+  let_it_be(:path_lock) { create(:path_lock, project: project, path: 'a.txt') }
+
   let(:commit) { project.repository.head_commit }
-  let!(:path_lock) { create(:path_lock, project: project, path: 'a.txt') }
 
-  describe '#summarize (entries)' do
-    subject { described_class.new(commit, project).summarize.first }
+  subject { described_class.new(commit, project).summarize.first }
 
+  describe '#summarize (entries)' do
     it 'includes path locks in entries' do
       is_expected.to contain_exactly(
         a_hash_including(file_name: 'a.txt', lock_label: "Locked by #{path_lock.user.name}")
       )
     end
   end
+
+  context 'when file_locks feature is unavailable' do
+    before do
+      stub_feature_flags(file_locks: false)
+    end
+
+    it 'does not fill lock labels' do
+      expect(subject.first.keys).not_to include(:lock_label)
+    end
+  end
 end

ee/spec/models/path_lock_spec.rb

@@ -69,4 +69,13 @@ describe PathLock do
       expect(path_lock.exact?("app")).to be_falsey
     end
   end
+
+  describe '.for_paths' do
+    let!(:another_path_lock) { create(:path_lock, path: 'app') }
+
+    it 'filters path locks by passed' do
+      expect(described_class.for_paths(['app'])).to eq([another_path_lock])
+      expect(described_class.for_paths(['app/models'])).to eq([path_lock])
+    end
+  end
 end