Add minor updates

2 files changed, 32 insertions, 9 deletions

diff --git a/app/workers/stuck_ci_builds_worker.rb b/app/workers/stuck_ci_builds_worker.rb
index 1c1e71cca23..e9bd3f180b7 100644
--- a/app/workers/stuck_ci_builds_worker.rb
+++ b/app/workers/stuck_ci_builds_worker.rb
@@ -2,6 +2,8 @@ class StuckCiBuildsWorker
   include Sidekiq::Worker
   include CronjobQueue
 
+  EXCLUSIVE_LEASE_KEY = 'stuck_ci_builds_worker_lease'
+
   BUILD_RUNNING_OUTDATED_TIMEOUT = 1.hour
   BUILD_PENDING_OUTDATED_TIMEOUT = 1.day
   BUILD_PENDING_STUCK_TIMEOUT = 1.hour
@@ -11,20 +13,26 @@ class StuckCiBuildsWorker
 
     Rails.logger.info "#{self.class}: Cleaning stuck builds"
 
-    drop       :running, BUILD_RUNNING_OUTDATED_TIMEOUT
-    drop       :pending, BUILD_PENDING_OUTDATED_TIMEOUT
+    drop :running, BUILD_RUNNING_OUTDATED_TIMEOUT
+    drop :pending, BUILD_PENDING_OUTDATED_TIMEOUT
     drop_stuck :pending, BUILD_PENDING_STUCK_TIMEOUT
+
+    remove_lease
   end
 
   private
 
   def try_obtain_lease
-    Gitlab::ExclusiveLease.new(
-      'stuck_ci_builds_worker_lease',
-      timeout: 30.minutes
+    @uuid = Gitlab::ExclusiveLease.new(
+        EXCLUSIVE_LEASE_KEY,
+        timeout: 30.minutes
     ).try_obtain
   end
 
+  def remove_lease
+    Gitlab::ExclusiveLease.cancel(EXCLUSIVE_LEASE_KEY, @uuid)
+  end
+
   def drop(status, timeout)
     search(status, timeout) do |build|
       drop_build :outdated, build, status, timeout
diff --git a/spec/workers/stuck_ci_builds_worker_spec.rb b/spec/workers/stuck_ci_builds_worker_spec.rb
index c88a0dfa4a6..5b87718a795 100644
--- a/spec/workers/stuck_ci_builds_worker_spec.rb
+++ b/spec/workers/stuck_ci_builds_worker_spec.rb
@@ -4,6 +4,7 @@ describe StuckCiBuildsWorker do
   let!(:runner) { create :ci_runner }
   let!(:build) { create :ci_build, runner: runner }
   let(:worker) { described_class.new }
+  let(:exclusive_lease_uuid) { SecureRandom.uuid }
 
   subject do
     build.reload
@@ -12,7 +13,7 @@ describe StuckCiBuildsWorker do
 
   before do
     build.update!(status: status, updated_at: updated_at)
-    allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain).and_return(true)
+    allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain).and_return(exclusive_lease_uuid)
   end
 
   shared_examples 'build is dropped' do
@@ -103,12 +104,26 @@ describe StuckCiBuildsWorker do
   describe 'exclusive lease' do
     let(:status) { 'running' }
     let(:updated_at) { 2.days.ago }
+    let(:worker2) { described_class.new }
 
-    it 'is guard by exclusive lease' do
+    it 'is guard by exclusive lease when executed concurrently' do
+      expect(worker).to receive(:drop).at_least(:once)
+      expect(worker2).not_to receive(:drop)
       worker.perform
       allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain).and_return(false)
-      expect(worker).not_to receive(:drop)
+      worker2.perform
+    end
+
+    it 'can be executed in sequence' do
+      expect(worker).to receive(:drop).at_least(:once)
+      expect(worker2).to receive(:drop).at_least(:once)
+      worker.perform
+      worker2.perform
+    end
+
+    it 'cancels exclusive lease after worker perform' do
+      expect(Gitlab::ExclusiveLease).to receive(:cancel).with(described_class::EXCLUSIVE_LEASE_KEY, exclusive_lease_uuid)
       worker.perform
     end
   end
-end
+end
\ No newline at end of file