Cancel ExclusiveLease upon completion in ProjectCacheWorker44726-cancel_lease_upon_completion_in_project_cache_worker

3 files changed, 44 insertions, 36 deletions

diff --git a/app/workers/project_cache_worker.rb b/app/workers/project_cache_worker.rb
index a993b4b2680..69d981b0b3e 100644
--- a/app/workers/project_cache_worker.rb
+++ b/app/workers/project_cache_worker.rb
@@ -1,9 +1,12 @@
 # Worker for updating any project specific caches.
 class ProjectCacheWorker
   include ApplicationWorker
+  include ExclusiveLeaseGuard
 
   LEASE_TIMEOUT = 15.minutes.to_i
 
+  attr_reader :project
+
   # project_id - The ID of the project for which to flush the cache.
   # files - An Array containing extra types of files to refresh such as
   #         `:readme` to flush the README and `:changelog` to flush the
@@ -11,30 +14,31 @@ class ProjectCacheWorker
   # statistics - An Array containing columns from ProjectStatistics to
   #              refresh, if empty all columns will be refreshed
   def perform(project_id, files = [], statistics = [])
-    project = Project.find_by(id: project_id)
+    @project = Project.find_by(id: project_id)
 
-    return unless project && project.repository.exists?
+    return unless @project && @project.repository.exists?
 
-    update_statistics(project, statistics.map(&:to_sym))
+    update_statistics(statistics.map(&:to_sym))
 
-    project.repository.refresh_method_caches(files.map(&:to_sym))
+    @project.repository.refresh_method_caches(files.map(&:to_sym))
 
-    project.cleanup
+    @project.cleanup
   end
 
-  def update_statistics(project, statistics = [])
-    return unless try_obtain_lease_for(project.id, :update_statistics)
-
-    Rails.logger.info("Updating statistics for project #{project.id}")
+  private
 
-    project.statistics.refresh!(only: statistics)
+  def update_statistics(statistics = [])
+    try_obtain_lease do
+      Rails.logger.info("Updating statistics for project #{@project.id}")
+      @project.statistics.refresh!(only: statistics)
+    end
   end
 
-  private
+  def lease_timeout
+    LEASE_TIMEOUT
+  end
 
-  def try_obtain_lease_for(project_id, section)
-    Gitlab::ExclusiveLease
-      .new("project_cache_worker:#{project_id}:#{section}", timeout: LEASE_TIMEOUT)
-      .try_obtain
+  def lease_key
+    "project_cache_worker:#{@project.id}:update_statistics"
   end
 end
diff --git a/changelogs/unreleased/44726-cancel_lease_upon_completion_in_project_cache_worker.yml b/changelogs/unreleased/44726-cancel_lease_upon_completion_in_project_cache_worker.yml
new file mode 100644
index 00000000000..664cb7df918
--- /dev/null
+++ b/changelogs/unreleased/44726-cancel_lease_upon_completion_in_project_cache_worker.yml
@@ -0,0 +1,5 @@
+---
+title: Cancel ExclusiveLease upon completion in ProjectCacheWorker
+merge_request:
+author:
+type: fixed
diff --git a/spec/workers/project_cache_worker_spec.rb b/spec/workers/project_cache_worker_spec.rb
index 6b1f2ff3227..06bad4de150 100644
--- a/spec/workers/project_cache_worker_spec.rb
+++ b/spec/workers/project_cache_worker_spec.rb
@@ -5,12 +5,11 @@ describe ProjectCacheWorker do
   let(:project) { create(:project, :repository) }
   let(:statistics) { project.statistics }
 
-  describe '#perform' do
-    before do
-      allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain)
-        .and_return(true)
-    end
+  before do
+    allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain).and_return('uuid')
+  end
 
+  describe '#perform' do
     context 'with a non-existing project' do
       it 'does nothing' do
         expect(worker).not_to receive(:update_statistics)
@@ -32,7 +31,7 @@ describe ProjectCacheWorker do
     context 'with an existing project' do
       it 'updates the project statistics' do
         expect(worker).to receive(:update_statistics)
-          .with(kind_of(Project), %i(repository_size))
+          .with(%i(repository_size))
           .and_call_original
 
         worker.perform(project.id, [], %w(repository_size))
@@ -58,32 +57,32 @@ describe ProjectCacheWorker do
         end
       end
     end
-  end
 
-  describe '#update_statistics' do
     context 'when a lease could not be obtained' do
-      it 'does not update the repository size' do
-        allow(worker).to receive(:try_obtain_lease_for)
-          .with(project.id, :update_statistics)
-          .and_return(false)
+      before do
+        allow(worker).to receive(:try_obtain_lease).and_return(false)
+      end
 
-        expect(statistics).not_to receive(:refresh!)
+      it 'does not update the repository size' do
+        expect_any_instance_of(ProjectStatistics).not_to receive(:refresh!)
 
-        worker.update_statistics(project)
+        worker.perform(project.id, [], %w(repository_size))
       end
     end
 
     context 'when a lease could be obtained' do
       it 'updates the project statistics' do
-        allow(worker).to receive(:try_obtain_lease_for)
-          .with(project.id, :update_statistics)
-          .and_return(true)
+        expect_any_instance_of(ProjectStatistics).to receive(:refresh!)
+                                                       .with(only: %i(repository_size))
+                                                       .and_call_original
 
-        expect(statistics).to receive(:refresh!)
-          .with(only: %i(repository_size))
-          .and_call_original
+        worker.perform(project.id, [], %i(repository_size))
+      end
+
+      it 'cancels the lease after statistics has been updated' do
+        expect(worker).to receive(:release_lease).with('uuid')
 
-        worker.update_statistics(project, %i(repository_size))
+        worker.perform(project.id, [], %i(repository_size))
       end
     end
   end