Cancel ExclusiveLease upon completion in ProjectCacheWorker

1 files changed, 17 insertions, 16 deletions

diff --git a/app/workers/project_cache_worker.rb b/app/workers/project_cache_worker.rb
index b0e1d8837d9..abe86066fb4 100644
--- a/app/workers/project_cache_worker.rb
+++ b/app/workers/project_cache_worker.rb
@@ -3,6 +3,7 @@
 # Worker for updating any project specific caches.
 class ProjectCacheWorker
   include ApplicationWorker
+  include ExclusiveLeaseGuard
 
   LEASE_TIMEOUT = 15.minutes.to_i
 
@@ -13,30 +14,30 @@ 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&.repository&.exists?
 
-    return unless project && project.repository.exists?
+    update_statistics(statistics)
 
-    update_statistics(project, 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.to_a.map(&:to_sym))
+    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