Use ExclusiveLeaseGuard where possibletc-backport-exclusive-lease-guard

10 files changed, 87 insertions, 67 deletions

diff --git a/app/services/concerns/exclusive_lease_guard.rb b/app/services/concerns/exclusive_lease_guard.rb
index f94f8d8e731..7ff58805758 100644
--- a/app/services/concerns/exclusive_lease_guard.rb
+++ b/app/services/concerns/exclusive_lease_guard.rb
@@ -1,27 +1,40 @@
 #
 # Concern that helps with getting an exclusive lease for running a block
-# of code.
+# of code. There are flavors:
+#  - `#try_obtain_lease`
+#  - `#try_obtain_lease_for`
 #
 # `#try_obtain_lease` takes a block which will be run if it was able to
 # obtain  the lease. Implement `#lease_timeout` to configure the timeout
 # for the exclusive lease. Optionally override `#lease_key` to set the
 # lease key, it defaults to the class name with underscores.
 #
+# `#try_obtain_lease_for` does about the same, but takes an additional
+# `id` parameter. This `id` is passed to `#lease_key_for`.
+#
 module ExclusiveLeaseGuard
   extend ActiveSupport::Concern
 
-  def try_obtain_lease
-    lease = exclusive_lease.try_obtain
+  def try_obtain_lease(&block)
+    try_obtain_lease_do(exclusive_lease, &block)
+  end
+
+  def try_obtain_lease_for(id, &block)
+    try_obtain_lease_do(exclusive_lease_for(id), &block)
+  end
+
+  def try_obtain_lease_do(lease, &block)
+    uuid = lease.try_obtain
 
-    unless lease
+    unless uuid
       log_error('Cannot obtain an exclusive lease. There must be another instance already in execution.')
       return
     end
 
     begin
-      yield lease
+      yield uuid
     ensure
-      release_lease(lease)
+      release_lease(uuid)
     end
   end
 
@@ -33,6 +46,14 @@ module ExclusiveLeaseGuard
     @lease_key ||= self.class.name.underscore
   end
 
+  def exclusive_lease_for(id)
+    Gitlab::ExclusiveLease.new(lease_key_for(id), timeout: lease_timeout)
+  end
+
+  def lease_key_for(id)
+    self.class.name.underscore.concat(":{#id}")
+  end
+
   def lease_timeout
     raise NotImplementedError,
       "#{self.class.name} does not implement #{__method__}"
@@ -47,6 +68,6 @@ module ExclusiveLeaseGuard
   end
 
   def log_error(message, extra_args = {})
-    logger.error(messages)
+    logger.error(message)
   end
 end
diff --git a/app/workers/gitlab_usage_ping_worker.rb b/app/workers/gitlab_usage_ping_worker.rb
index 6dd281b1147..35c5e408c04 100644
--- a/app/workers/gitlab_usage_ping_worker.rb
+++ b/app/workers/gitlab_usage_ping_worker.rb
@@ -3,17 +3,21 @@ class GitlabUsagePingWorker
 
   include ApplicationWorker
   include CronjobQueue
+  include ExclusiveLeaseGuard
 
   def perform
-    # Multiple Sidekiq workers could run this. We should only do this at most once a day.
-    return unless try_obtain_lease
-
-    SubmitUsagePingService.new.execute
+    try_obtain_lease do
+      SubmitUsagePingService.new.execute
+    end
   end
 
   private
 
-  def try_obtain_lease
-    Gitlab::ExclusiveLease.new('gitlab_usage_ping_worker:ping', timeout: LEASE_TIMEOUT).try_obtain
+  def lease_key
+    'gitlab_usage_ping_worker:ping'
+  end
+
+  def lease_timeout
+    LEASE_TIMEOUT
   end
 end
diff --git a/app/workers/project_cache_worker.rb b/app/workers/project_cache_worker.rb
index a993b4b2680..6d30cce9e66 100644
--- a/app/workers/project_cache_worker.rb
+++ b/app/workers/project_cache_worker.rb
@@ -1,6 +1,7 @@
 # Worker for updating any project specific caches.
 class ProjectCacheWorker
   include ApplicationWorker
+  include ExclusiveLeaseGuard
 
   LEASE_TIMEOUT = 15.minutes.to_i
 
@@ -23,18 +24,20 @@ class ProjectCacheWorker
   end
 
   def update_statistics(project, statistics = [])
-    return unless try_obtain_lease_for(project.id, :update_statistics)
+    try_obtain_lease_for(project.id) do
+      Rails.logger.info("Updating statistics for project #{project.id}")
 
-    Rails.logger.info("Updating statistics for project #{project.id}")
-
-    project.statistics.refresh!(only: statistics)
+      project.statistics.refresh!(only: statistics)
+    end
   end
 
   private
 
-  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_for(project_id)
+    "project_cache_worker:#{project_id}:update_statistics"
+  end
+
+  def lease_timeout
+    LEASE_TIMEOUT
   end
 end
diff --git a/app/workers/project_migrate_hashed_storage_worker.rb b/app/workers/project_migrate_hashed_storage_worker.rb
index d01eb744e5d..eb1217a8b1f 100644
--- a/app/workers/project_migrate_hashed_storage_worker.rb
+++ b/app/workers/project_migrate_hashed_storage_worker.rb
@@ -1,5 +1,6 @@
 class ProjectMigrateHashedStorageWorker
   include ApplicationWorker
+  include ExclusiveLeaseGuard
 
   LEASE_TIMEOUT = 30.seconds.to_i
 
@@ -7,28 +8,18 @@ class ProjectMigrateHashedStorageWorker
     project = Project.find_by(id: project_id)
     return if project.nil? || project.pending_delete?
 
-    uuid = lease_for(project_id).try_obtain
-    if uuid
+    try_obtain_lease_for(project_id) do
       ::Projects::HashedStorageMigrationService.new(project, logger).execute
-    else
-      false
     end
-  rescue => ex
-    cancel_lease_for(project_id, uuid) if uuid
-    raise ex
-  end
-
-  def lease_for(project_id)
-    Gitlab::ExclusiveLease.new(lease_key(project_id), timeout: LEASE_TIMEOUT)
   end
 
   private
 
-  def lease_key(project_id)
+  def lease_key_for(project_id)
     "project_migrate_hashed_storage_worker:#{project_id}"
   end
 
-  def cancel_lease_for(project_id, uuid)
-    Gitlab::ExclusiveLease.cancel(lease_key(project_id), uuid)
+  def lease_timeout
+    LEASE_TIMEOUT
   end
 end
diff --git a/app/workers/propagate_service_template_worker.rb b/app/workers/propagate_service_template_worker.rb
index 635a97c99af..d9a7041a27e 100644
--- a/app/workers/propagate_service_template_worker.rb
+++ b/app/workers/propagate_service_template_worker.rb
@@ -1,20 +1,19 @@
 # Worker for updating any project specific caches.
 class PropagateServiceTemplateWorker
   include ApplicationWorker
+  include ExclusiveLeaseGuard
 
   LEASE_TIMEOUT = 4.hours.to_i
 
   def perform(template_id)
-    return unless try_obtain_lease_for(template_id)
-
-    Projects::PropagateServiceTemplate.propagate(Service.find_by(id: template_id))
+    try_obtain_lease_for(template_id) do
+      Projects::PropagateServiceTemplate.propagate(Service.find_by(id: template_id))
+    end
   end
 
   private
 
-  def try_obtain_lease_for(template_id)
-    Gitlab::ExclusiveLease
-      .new("propagate_service_template_worker:#{template_id}", timeout: LEASE_TIMEOUT)
-      .try_obtain
+  def lease_timeout
+    LEASE_TIMEOUT
   end
 end
diff --git a/app/workers/repository_check/batch_worker.rb b/app/workers/repository_check/batch_worker.rb
index 76688cf51c1..75900020fa7 100644
--- a/app/workers/repository_check/batch_worker.rb
+++ b/app/workers/repository_check/batch_worker.rb
@@ -2,6 +2,7 @@ module RepositoryCheck
   class BatchWorker
     include ApplicationWorker
     include CronjobQueue
+    include ExclusiveLeaseGuard
 
     RUN_TIME = 3600
 
@@ -17,9 +18,9 @@ module RepositoryCheck
         break if Time.now - start >= RUN_TIME
         break unless current_settings.repository_checks_enabled
 
-        next unless try_obtain_lease(project_id)
-
-        SingleRepositoryWorker.new.perform(project_id)
+        try_obtain_lease_for(project_id) do
+          SingleRepositoryWorker.new.perform(project_id)
+        end
       end
     end
 
@@ -39,13 +40,14 @@ module RepositoryCheck
       never_checked_projects + old_check_projects
     end
 
-    def try_obtain_lease(id)
+    def lease_key_for(id)
+      "project_repository_check:#{id}"
+    end
+
+    def lease_timeout
       # Use a 24-hour timeout because on servers/projects where 'git fsck' is
       # super slow we definitely do not want to run it twice in parallel.
-      Gitlab::ExclusiveLease.new(
-        "project_repository_check:#{id}",
-        timeout: 24.hours
-      ).try_obtain
+      24.hours
     end
 
     def current_settings
diff --git a/app/workers/stuck_ci_jobs_worker.rb b/app/workers/stuck_ci_jobs_worker.rb
index fb26fa4c515..5d3b73f5247 100644
--- a/app/workers/stuck_ci_jobs_worker.rb
+++ b/app/workers/stuck_ci_jobs_worker.rb
@@ -1,6 +1,7 @@
 class StuckCiJobsWorker
   include ApplicationWorker
   include CronjobQueue
+  include ExclusiveLeaseGuard
 
   EXCLUSIVE_LEASE_KEY = 'stuck_ci_builds_worker_lease'.freeze
 
@@ -9,25 +10,23 @@ class StuckCiJobsWorker
   BUILD_PENDING_STUCK_TIMEOUT = 1.hour
 
   def perform
-    return unless try_obtain_lease
+    try_obtain_lease do
+      Rails.logger.info "#{self.class}: Cleaning stuck builds"
 
-    Rails.logger.info "#{self.class}: Cleaning stuck builds"
-
-    drop :running, BUILD_RUNNING_OUTDATED_TIMEOUT
-    drop :pending, BUILD_PENDING_OUTDATED_TIMEOUT
-    drop_stuck :pending, BUILD_PENDING_STUCK_TIMEOUT
-
-    remove_lease
+      drop :running, BUILD_RUNNING_OUTDATED_TIMEOUT
+      drop :pending, BUILD_PENDING_OUTDATED_TIMEOUT
+      drop_stuck :pending, BUILD_PENDING_STUCK_TIMEOUT
+    end
   end
 
   private
 
-  def try_obtain_lease
-    @uuid = Gitlab::ExclusiveLease.new(EXCLUSIVE_LEASE_KEY, timeout: 30.minutes).try_obtain
+  def lease_key
+    EXCLUSIVE_LEASE_KEY
   end
 
-  def remove_lease
-    Gitlab::ExclusiveLease.cancel(EXCLUSIVE_LEASE_KEY, @uuid)
+  def lease_timeout
+    30.minutes
   end
 
   def drop(status, timeout)
diff --git a/spec/workers/gitlab_usage_ping_worker_spec.rb b/spec/workers/gitlab_usage_ping_worker_spec.rb
index 49b4e04dc7c..3a4bd9d2066 100644
--- a/spec/workers/gitlab_usage_ping_worker_spec.rb
+++ b/spec/workers/gitlab_usage_ping_worker_spec.rb
@@ -4,7 +4,7 @@ describe GitlabUsagePingWorker do
   subject { described_class.new }
 
   it 'delegates to SubmitUsagePingService' do
-    allow(subject).to receive(:try_obtain_lease).and_return(true)
+    allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain).and_return(true)
 
     expect_any_instance_of(SubmitUsagePingService).to receive(:execute)
 
diff --git a/spec/workers/project_cache_worker_spec.rb b/spec/workers/project_cache_worker_spec.rb
index 6b1f2ff3227..c164bfb39c7 100644
--- a/spec/workers/project_cache_worker_spec.rb
+++ b/spec/workers/project_cache_worker_spec.rb
@@ -64,7 +64,7 @@ describe ProjectCacheWorker 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)
+          .with(project.id)
           .and_return(false)
 
         expect(statistics).not_to receive(:refresh!)
@@ -76,8 +76,9 @@ describe ProjectCacheWorker do
     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)
+          .with(project.id)
           .and_return(true)
+          .and_call_original
 
         expect(statistics).to receive(:refresh!)
           .with(only: %i(repository_size))
diff --git a/spec/workers/project_migrate_hashed_storage_worker_spec.rb b/spec/workers/project_migrate_hashed_storage_worker_spec.rb
index 2e3951e7afc..b06f3b60d45 100644
--- a/spec/workers/project_migrate_hashed_storage_worker_spec.rb
+++ b/spec/workers/project_migrate_hashed_storage_worker_spec.rb
@@ -7,7 +7,7 @@ describe ProjectMigrateHashedStorageWorker, :clean_gitlab_redis_shared_state do
 
     context 'when have exclusive lease' do
       before do
-        lease = subject.lease_for(project.id)
+        lease = subject.exclusive_lease_for(project.id)
 
         allow(Gitlab::ExclusiveLease).to receive(:new).and_return(lease)
         allow(lease).to receive(:try_obtain).and_return(true)
@@ -37,7 +37,7 @@ describe ProjectMigrateHashedStorageWorker, :clean_gitlab_redis_shared_state do
 
     context 'when dont have exclusive lease' do
       before do
-        lease = subject.lease_for(project.id)
+        lease = subject.exclusive_lease_for(project.id)
 
         allow(Gitlab::ExclusiveLease).to receive(:new).and_return(lease)
         allow(lease).to receive(:try_obtain).and_return(false)