Merge branch 'ee-da-exclusive-lease-helpers' into 'master'

[EE] Add helper methods to stub Gitlab::ExclusiveLease

See merge request gitlab-org/gitlab-ee!6322

ee/spec/lib/gitlab/geo/log_cursor/daemon_spec.rb

@@ -2,6 +2,7 @@ require 'spec_helper'
 
 describe Gitlab::Geo::LogCursor::Daemon, :postgresql, :clean_gitlab_redis_shared_state do
   include ::EE::GeoHelpers
+  include ExclusiveLeaseHelpers
 
   set(:primary) { create(:geo_node, :primary) }
   set(:secondary) { create(:geo_node) }
@@ -37,9 +38,14 @@ describe Gitlab::Geo::LogCursor::Daemon, :postgresql, :clean_gitlab_redis_shared
     end
 
     it 'skips execution if cannot achieve a lease' do
+      lease = stub_exclusive_lease_taken('geo_log_cursor_processed')
+
+      allow(lease).to receive(:try_obtain_with_ttl).and_return({ ttl: 1, uuid: false })
+      allow(lease).to receive(:same_uuid?).and_return(false)
+      allow(Gitlab::Geo::LogCursor::Lease).to receive(:exclusive_lease).and_return(lease)
+
       is_expected.to receive(:exit?).and_return(false, true)
       is_expected.not_to receive(:run_once!)
-      expect_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain_with_ttl).and_return({ ttl: 1, uuid: false })
 
       daemon.run!
     end

ee/spec/lib/gitlab/geo/log_cursor/lease_spec.rb

 require 'spec_helper'
 
 describe Gitlab::Geo::LogCursor::Lease, :clean_gitlab_redis_shared_state do
+  include ExclusiveLeaseHelpers
+
   describe '.exclusive_lease' do
     it 'returns an exclusive lease instance' do
       expect(described_class.send(:exclusive_lease)).to be_an_instance_of(Gitlab::ExclusiveLease)
@@ -8,15 +10,20 @@ describe Gitlab::Geo::LogCursor::Lease, :clean_gitlab_redis_shared_state do
   end
 
   describe '.renew!' do
+    let(:lease) { stub_exclusive_lease(described_class::LEASE_KEY, renew: true) }
+
+    before do
+      allow(described_class).to receive(:exclusive_lease).and_return(lease)
+    end
+
     it 'returns an exclusive lease instance' do
-      expect_any_instance_of(Gitlab::ExclusiveLease).to receive(:renew)
+      expect(lease).to receive(:renew)
 
       described_class.renew!
     end
 
     it 'logs with the correct caller class' do
       stub_const("Gitlab::Geo::LogCursor::Logger::PID", 111)
-      allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:renew).and_return(true)
 
       expect(::Gitlab::Logger).to receive(:debug).with(pid: 111,
                                                        class: 'Gitlab::Geo::LogCursor::Lease',
@@ -44,7 +51,12 @@ describe Gitlab::Geo::LogCursor::Lease, :clean_gitlab_redis_shared_state do
     end
 
     it 'returns > 0 if there was an error' do
-      expect(Gitlab::ExclusiveLease).to receive(:cancel)
+      lease = stub_exclusive_lease(described_class::LEASE_KEY, 'uuid')
+
+      allow(lease).to receive(:try_obtain_with_ttl).and_return({ ttl: 0, uuid: 'uuid' })
+      allow(described_class).to receive(:exclusive_lease).and_return(lease)
+
+      expect_to_cancel_exclusive_lease(described_class::LEASE_KEY, 'uuid')
 
       result = described_class.try_obtain_with_ttl { raise StandardError }
 

ee/spec/services/geo/file_download_service_spec.rb

@@ -2,21 +2,26 @@ require 'spec_helper'
 
 describe Geo::FileDownloadService do
   include ::EE::GeoHelpers
+  include ExclusiveLeaseHelpers
 
   set(:primary)  { create(:geo_node, :primary) }
   set(:secondary) { create(:geo_node) }
 
   before do
     stub_current_geo_node(secondary)
-
-    allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain).and_return(true)
   end
 
   shared_examples_for 'a service that downloads the file and registers the sync result' do |file_type|
     let(:download_service) { described_class.new(file_type, file.id) }
     let(:registry) { file_type == 'job_artifact' ? Geo::JobArtifactRegistry : Geo::FileRegistry }
+
     subject(:execute!) { download_service.execute }
 
+    before do
+      stub_exclusive_lease("file_download_service:#{file_type}:#{file.id}",
+        timeout: Geo::FileDownloadService::LEASE_TIMEOUT)
+    end
+
     context 'for a new file' do
       context 'when the downloader fails before attempting a transfer' do
         it 'logs that the download failed before attempting a transfer' do

ee/spec/services/geo/file_registry_removal_service_spec.rb

@@ -2,18 +2,18 @@ require 'spec_helper'
 
 describe Geo::FileRegistryRemovalService do
   include ::EE::GeoHelpers
+  include ExclusiveLeaseHelpers
 
   set(:secondary) { create(:geo_node) }
 
   before do
     stub_current_geo_node(secondary)
-
-    allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain).and_return(true)
   end
 
   describe '#execute' do
     it 'delegates log_error to the Geo logger' do
-      allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain).and_return(false)
+      stub_exclusive_lease_taken("file_registry_removal_service:lfs:99")
+
       expect(Gitlab::Geo::Logger).to receive(:error)
 
       described_class.new(:lfs, 99).execute
@@ -22,6 +22,11 @@ describe Geo::FileRegistryRemovalService do
     shared_examples 'removes' do
       subject(:service) { described_class.new(file_registry.file_type, file_registry.file_id) }
 
+      before do
+        stub_exclusive_lease("file_registry_removal_service:#{file_registry.file_type}:#{file_registry.file_id}",
+          timeout: Geo::FileRegistryRemovalService::LEASE_TIMEOUT)
+      end
+
       it 'file from disk' do
         expect do
           service.execute
@@ -38,6 +43,11 @@ describe Geo::FileRegistryRemovalService do
     shared_examples 'removes artifact' do
       subject(:service) { described_class.new('job_artifact', registry.artifact_id) }
 
+      before do
+        stub_exclusive_lease("file_registry_removal_service:job_artifact:#{registry.artifact_id}",
+          timeout: Geo::FileRegistryRemovalService::LEASE_TIMEOUT)
+      end
+
       it 'file from disk' do
         expect do
           service.execute

ee/spec/services/geo/repository_sync_service_spec.rb

@@ -2,13 +2,15 @@ require 'spec_helper'
 
 describe Geo::RepositorySyncService do
   include ::EE::GeoHelpers
+  include ExclusiveLeaseHelpers
 
   set(:primary) { create(:geo_node, :primary) }
   set(:secondary) { create(:geo_node) }
-
-  let(:lease) { double(try_obtain: true) }
   set(:project) { create(:project_empty_repo) }
+
   let(:repository) { project.repository }
+  let(:lease_key) { "geo_sync_service:repository:#{project.id}" }
+  let(:lease_uuid) { 'uuid'}
 
   subject { described_class.new(project) }
 
@@ -23,7 +25,8 @@ describe Geo::RepositorySyncService do
     let(:url_to_repo) { "#{primary.url}#{project.full_path}.git" }
 
     before do
-      allow(subject).to receive(:exclusive_lease).and_return(lease)
+      stub_exclusive_lease(lease_key, lease_uuid)
+      stub_exclusive_lease("geo_project_housekeeping:#{project.id}")
 
       allow_any_instance_of(Repository).to receive(:fetch_as_mirror)
         .and_return(true)
@@ -47,8 +50,7 @@ describe Geo::RepositorySyncService do
     end
 
     it 'returns the lease when succeed' do
-      expect(Gitlab::ExclusiveLease).to receive(:cancel).once.with(
-        subject.__send__(:lease_key), anything).and_call_original
+      expect_to_cancel_exclusive_lease(lease_key, lease_uuid)
 
       subject.execute
     end
@@ -64,14 +66,13 @@ describe Geo::RepositorySyncService do
         .with(url_to_repo, remote_name: 'geo', forced: true)
         .and_raise(Gitlab::Shell::Error)
 
-      expect(Gitlab::ExclusiveLease).to receive(:cancel).once.with(
-        subject.__send__(:lease_key), anything).and_call_original
+      expect_to_cancel_exclusive_lease(lease_key, lease_uuid)
 
       subject.execute
     end
 
     it 'does not fetch project repository if cannot obtain a lease' do
-      allow(lease).to receive(:try_obtain) { false }
+      stub_exclusive_lease_taken(lease_key)
 
       expect(repository).not_to receive(:fetch_as_mirror)
 

ee/spec/services/geo/wiki_sync_service_spec.rb

@@ -2,13 +2,15 @@ require 'spec_helper'
 
 RSpec.describe Geo::WikiSyncService do
   include ::EE::GeoHelpers
+  include ExclusiveLeaseHelpers
 
   set(:primary) { create(:geo_node, :primary) }
   set(:secondary) { create(:geo_node) }
+  set(:project) { create(:project_empty_repo) }
 
-  let(:lease) { double(try_obtain: true, lease_key: anything ) }
-  let(:project) { create(:project_empty_repo) }
   let(:repository) { project.wiki.repository }
+  let(:lease_key) { "geo_sync_service:wiki:#{project.id}" }
+  let(:lease_uuid) { 'uuid'}
 
   subject { described_class.new(project) }
 
@@ -23,7 +25,7 @@ RSpec.describe Geo::WikiSyncService do
     let(:url_to_repo) { "#{primary.url}#{project.full_path}.wiki.git" }
 
     before do
-      allow(subject).to receive(:exclusive_lease).and_return(lease)
+      stub_exclusive_lease(lease_key, lease_uuid)
 
       allow_any_instance_of(Repository).to receive(:fetch_as_mirror)
         .and_return(true)
@@ -39,8 +41,7 @@ RSpec.describe Geo::WikiSyncService do
     end
 
     it 'releases lease' do
-      expect(Gitlab::ExclusiveLease).to receive(:cancel).at_least(:once).with(
-        subject.__send__(:lease_key), anything).and_call_original
+      expect_to_cancel_exclusive_lease(lease_key, lease_uuid)
 
       subject.execute
     end
@@ -52,7 +53,7 @@ RSpec.describe Geo::WikiSyncService do
     end
 
     it 'does not fetch wiki repository if cannot obtain a lease' do
-      allow(lease).to receive(:try_obtain) { false }
+      stub_exclusive_lease_taken(lease_key)
 
       expect(repository).not_to receive(:fetch_as_mirror)
 

ee/spec/workers/geo/prune_event_log_worker_spec.rb

@@ -2,13 +2,18 @@ require 'spec_helper'
 
 describe Geo::PruneEventLogWorker, :geo do
   include ::EE::GeoHelpers
+  include ExclusiveLeaseHelpers
 
   subject(:worker) { described_class.new }
+
   set(:primary) { create(:geo_node, :primary) }
   set(:secondary) { create(:geo_node) }
 
+  let(:lease_key) { 'geo/prune_event_log_worker' }
+  let(:lease_timeout) { Geo::PruneEventLogWorker::LEASE_TIMEOUT }
+
   before do
-    allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain).and_return(true)
+    stub_exclusive_lease(lease_key, timeout: lease_timeout)
   end
 
   describe '#perform' do
@@ -30,7 +35,7 @@ describe Geo::PruneEventLogWorker, :geo do
       end
 
       it 'logs error when it cannot obtain lease' do
-        allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain) { nil }
+        stub_exclusive_lease_taken(lease_key, timeout: lease_timeout)
 
         expect(worker).to receive(:log_error).with(/^Cannot obtain an exclusive lease/)
 

ee/spec/workers/geo/repositories_clean_up_worker_spec.rb

@@ -2,10 +2,12 @@ require 'spec_helper'
 
 describe Geo::RepositoriesCleanUpWorker do
   describe '#perform' do
+    include ExclusiveLeaseHelpers
+
     let(:geo_node) { create(:geo_node) }
 
     before do
-      allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain) { true }
+      stub_exclusive_lease
     end
 
     context 'when node has namespace restrictions' do
@@ -68,7 +70,7 @@ describe Geo::RepositoriesCleanUpWorker do
     end
 
     it 'does not perform GeoRepositoryDestroyWorker when cannnot obtain a lease' do
-      allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain) { false }
+      stub_exclusive_lease_taken
 
       expect(GeoRepositoryDestroyWorker).not_to receive(:perform_async)
 

ee/spec/workers/geo/repository_shard_sync_worker_spec.rb

@@ -4,6 +4,7 @@ require 'spec_helper'
 # can't see changes inside a transaction of a different connection.
 describe Geo::RepositoryShardSyncWorker, :geo, :delete, :clean_gitlab_redis_cache do
   include ::EE::GeoHelpers
+  include ExclusiveLeaseHelpers
 
   let!(:primary) { create(:geo_node, :primary) }
   let!(:secondary) { create(:geo_node) }
@@ -25,8 +26,7 @@ describe Geo::RepositoryShardSyncWorker, :geo, :delete, :clean_gitlab_redis_cach
     end
 
     before do
-      allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain) { true }
-      allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:renew) { true }
+      stub_exclusive_lease(renew: true)
 
       Gitlab::ShardHealthCache.update([shard_name])
     end
@@ -48,6 +48,7 @@ describe Geo::RepositoryShardSyncWorker, :geo, :delete, :clean_gitlab_redis_cach
 
     it 'does not perform Geo::ProjectSyncWorker when shard becomes unhealthy' do
       Gitlab::ShardHealthCache.update([])
+
       expect(Geo::ProjectSyncWorker).not_to receive(:perform_async)
 
       subject.perform(shard_name)

ee/spec/workers/geo/repository_verification/primary/shard_worker_spec.rb

@@ -2,6 +2,7 @@ require 'spec_helper'
 
 describe Geo::RepositoryVerification::Primary::ShardWorker, :postgresql, :clean_gitlab_redis_cache do
   include ::EE::GeoHelpers
+  include ExclusiveLeaseHelpers
 
   let!(:primary)   { create(:geo_node, :primary) }
   let(:shard_name) { Gitlab.config.repositories.storages.keys.first }
@@ -13,8 +14,7 @@ describe Geo::RepositoryVerification::Primary::ShardWorker, :postgresql, :clean_
 
   describe '#perform' do
     before do
-      allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain) { true }
-      allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:renew) { true }
+      stub_exclusive_lease(renew: true)
 
       Gitlab::ShardHealthCache.update([shard_name])
     end

ee/spec/workers/geo/repository_verification/primary/single_worker_spec.rb

@@ -2,6 +2,7 @@ require 'spec_helper'
 
 describe Geo::RepositoryVerification::Primary::SingleWorker, :postgresql, :clean_gitlab_redis_cache do
   include ::EE::GeoHelpers
+  include ExclusiveLeaseHelpers
 
   set(:project) { create(:project) }
 
@@ -11,13 +12,10 @@ describe Geo::RepositoryVerification::Primary::SingleWorker, :postgresql, :clean
 
   before do
     stub_current_geo_node(primary)
+    stub_exclusive_lease
   end
 
   describe '#perform' do
-    before do
-      allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain) { true }
-    end
-
     it 'does not calculate the checksum when not running on a primary' do
       stub_project_repository(project, repository)
       stub_wiki_repository(project.wiki, wiki)

ee/spec/workers/geo/repository_verification/secondary/shard_worker_spec.rb

@@ -2,6 +2,7 @@ require 'spec_helper'
 
 describe Geo::RepositoryVerification::Secondary::ShardWorker, :postgresql, :clean_gitlab_redis_cache do
   include ::EE::GeoHelpers
+  include ExclusiveLeaseHelpers
 
   let!(:secondary) { create(:geo_node) }
   let(:shard_name) { Gitlab.config.repositories.storages.keys.first }
@@ -16,8 +17,7 @@ describe Geo::RepositoryVerification::Secondary::ShardWorker, :postgresql, :clea
 
   describe '#perform' do
     before do
-      allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain) { true }
-      allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:renew) { true }
+      stub_exclusive_lease(renew: true)
 
       Gitlab::ShardHealthCache.update([shard_name])
     end

ee/spec/workers/geo/repository_verification/secondary/single_worker_spec.rb

@@ -2,6 +2,7 @@ require 'spec_helper'
 
 describe Geo::RepositoryVerification::Secondary::SingleWorker, :postgresql, :clean_gitlab_redis_cache do
   include ::EE::GeoHelpers
+  include ExclusiveLeaseHelpers
 
   let!(:secondary) { create(:geo_node) }
   let(:project)    { create(:project, :repository, :wiki_repo) }
@@ -9,13 +10,10 @@ describe Geo::RepositoryVerification::Secondary::SingleWorker, :postgresql, :cle
 
   before do
     stub_current_geo_node(secondary)
+    stub_exclusive_lease
   end
 
   describe '#perform' do
-    before do
-      allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain) { true }
-    end
-
     it 'does not calculate the checksum when not running on a secondary' do
       allow(Gitlab::Geo).to receive(:secondary?) { false }
 

ee/spec/workers/repository_update_mirror_worker_spec.rb

 require 'rails_helper'
 
 describe RepositoryUpdateMirrorWorker do
+  include ExclusiveLeaseHelpers
+
   describe '#perform' do
     let(:jid) { '12345678' }
     let!(:project) { create(:project) }
     let!(:import_state) { create(:import_state, :mirror, :scheduled, project: project) }
 
     before do
-      allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain).and_return(true)
       allow(subject).to receive(:jid).and_return(jid)
     end
 
@@ -56,6 +57,10 @@ describe RepositoryUpdateMirrorWorker do
       end
 
       context 'when we obtain the lease' do
+        before do
+          allow(stub_exclusive_lease).to receive(:exists?).and_return(false)
+        end
+
         it 'performs UpdateAllMirrorsWorker when reschedule_immediately? returns true' do
           allow(Gitlab::Mirror).to receive(:reschedule_immediately?).and_return(true)
 
@@ -74,7 +79,7 @@ describe RepositoryUpdateMirrorWorker do
       end
 
       it 'does not perform UpdateAllMirrorsWorker when we cannot obtain the lease' do
-        allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain).and_return(false)
+        stub_exclusive_lease_taken
 
         expect(UpdateAllMirrorsWorker).not_to receive(:perform_async)
 
@@ -82,7 +87,7 @@ describe RepositoryUpdateMirrorWorker do
       end
 
       it 'does not perform UpdateAllMirrorsWorker when the lease already exists' do
-        allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:exists?).and_return(true)
+        stub_exclusive_lease_taken
 
         expect(UpdateAllMirrorsWorker).not_to receive(:perform_async)
 

ee/spec/workers/update_all_mirrors_worker_spec.rb

 require 'spec_helper'
 
 describe UpdateAllMirrorsWorker do
+  include ExclusiveLeaseHelpers
+
   subject(:worker) { described_class.new }
 
   before do
-    allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain).and_return(true)
+    stub_exclusive_lease
   end
 
   describe '#perform' do
     it 'does not execute if cannot get the lease' do
-      create(:project, :mirror)
-
-      allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain).and_return(false)
+      stub_exclusive_lease_taken
 
       expect(worker).not_to receive(:schedule_mirrors!)
 

spec/lib/mattermost/session_spec.rb

 require 'spec_helper'
 
 describe Mattermost::Session, type: :request do
+  include ExclusiveLeaseHelpers
+
   let(:user) { create(:user) }
 
   let(:gitlab_url) { "http://gitlab.com" }
@@ -97,26 +99,20 @@ describe Mattermost::Session, type: :request do
       end
     end
 
-    context 'with lease' do
-      before do
-        allow(subject).to receive(:lease_try_obtain).and_return('aldkfjsldfk')
-      end
+    context 'exclusive lease' do
+      let(:lease_key) { 'mattermost:session' }
 
       it 'tries to obtain a lease' do
-        expect(subject).to receive(:lease_try_obtain)
-        expect(Gitlab::ExclusiveLease).to receive(:cancel)
+        expect_to_obtain_exclusive_lease(lease_key, 'uuid')
+        expect_to_cancel_exclusive_lease(lease_key, 'uuid')
 
         # Cannot setup a session, but we should still cancel the lease
         expect { subject.with_session }.to raise_error(Mattermost::NoSessionError)
       end
-    end
 
-    context 'without lease' do
-      before do
-        allow(subject).to receive(:lease_try_obtain).and_return(nil)
-      end
+      it 'returns a NoSessionError error without lease' do
+        stub_exclusive_lease_taken(lease_key)
 
-      it 'returns a NoSessionError error' do
         expect { subject.with_session }.to raise_error(Mattermost::NoSessionError)
       end
     end

spec/models/ci/build_trace_chunk_spec.rb

 require 'spec_helper'
 
 describe Ci::BuildTraceChunk, :clean_gitlab_redis_shared_state do
+  include ExclusiveLeaseHelpers
+
   set(:build) { create(:ci_build, :running) }
   let(:chunk_index) { 0 }
   let(:data_store) { :redis }
@@ -322,7 +324,7 @@ describe Ci::BuildTraceChunk, :clean_gitlab_redis_shared_state do
 
   describe 'ExclusiveLock' do
     before do
-      allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain) { nil }
+      stub_exclusive_lease_taken
       stub_const('Ci::BuildTraceChunk::WRITE_LOCK_RETRY', 1)
     end
 

spec/models/concerns/reactive_caching_spec.rb

 require 'spec_helper'
 
 describe ReactiveCaching, :use_clean_rails_memory_store_caching do
+  include ExclusiveLeaseHelpers
   include ReactiveCachingHelpers
 
   class CacheTest
@@ -106,8 +107,8 @@ describe ReactiveCaching, :use_clean_rails_memory_store_caching do
       end
 
       it 'takes and releases the lease' do
-        expect_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain).and_return("000000")
-        expect(Gitlab::ExclusiveLease).to receive(:cancel).with(cache_key, "000000")
+        expect_to_obtain_exclusive_lease(cache_key, 'uuid')
+        expect_to_cancel_exclusive_lease(cache_key, 'uuid')
 
         go!
       end
@@ -160,11 +161,9 @@ describe ReactiveCaching, :use_clean_rails_memory_store_caching do
     end
 
     context 'when the lease is already taken' do
-      before do
-        expect_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain).and_return(nil)
-      end
-
       it 'skips the calculation' do
+        stub_exclusive_lease_taken(cache_key)
+
         expect(instance).to receive(:calculate_reactive_cache).never
 
         go!

spec/services/clusters/applications/check_ingress_ip_address_service_spec.rb

 require 'spec_helper'
 
 describe Clusters::Applications::CheckIngressIpAddressService do
+  include ExclusiveLeaseHelpers
+
   let(:application) { create(:clusters_applications_ingress, :installed) }
   let(:service) { described_class.new(application) }
   let(:kubeclient) { double(::Kubeclient::Client, get_service: kube_service) }
   let(:ingress) { [{ ip: '111.222.111.222' }] }
-  let(:exclusive_lease) { instance_double(Gitlab::ExclusiveLease, try_obtain: true) }
+  let(:lease_key) { "check_ingress_ip_address_service:#{application.id}" }
 
   let(:kube_service) do
     ::Kubeclient::Resource.new(
@@ -22,11 +24,8 @@ describe Clusters::Applications::CheckIngressIpAddressService do
   subject { service.execute }
 
   before do
+    stub_exclusive_lease(lease_key, timeout: 15.seconds.to_i)
     allow(application.cluster).to receive(:kubeclient).and_return(kubeclient)
-    allow(Gitlab::ExclusiveLease)
-      .to receive(:new)
-      .with("check_ingress_ip_address_service:#{application.id}", timeout: 15.seconds.to_i)
-      .and_return(exclusive_lease)
   end
 
   describe '#execute' do
@@ -47,13 +46,9 @@ describe Clusters::Applications::CheckIngressIpAddressService do
     end
 
     context 'when the exclusive lease cannot be obtained' do
-      before do
-        allow(exclusive_lease)
-          .to receive(:try_obtain)
-          .and_return(false)
-      end
-
       it 'does not call kubeclient' do
+        stub_exclusive_lease_taken(lease_key, timeout: 15.seconds.to_i)
+
         subject
 
         expect(kubeclient).not_to have_received(:get_service)

spec/services/users/refresh_authorized_projects_service_spec.rb

 require 'spec_helper'
 
 describe Users::RefreshAuthorizedProjectsService do
+  include ExclusiveLeaseHelpers
+
   # We're using let! here so that any expectations for the service class are not
   # triggered twice.
   let!(:project) { create(:project) }
@@ -10,12 +12,10 @@ describe Users::RefreshAuthorizedProjectsService do
 
   describe '#execute', :clean_gitlab_redis_shared_state do
     it 'refreshes the authorizations using a lease' do
-      expect_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain)
-        .and_return('foo')
-
-      expect(Gitlab::ExclusiveLease).to receive(:cancel)
-        .with(an_instance_of(String), 'foo')
+      lease_key = "refresh_authorized_projects:#{user.id}"
 
+      expect_to_obtain_exclusive_lease(lease_key, 'uuid')
+      expect_to_cancel_exclusive_lease(lease_key, 'uuid')
       expect(service).to receive(:execute_without_lease)
 
       service.execute

spec/support/helpers/exclusive_lease_helpers.rb

+module ExclusiveLeaseHelpers
+  def stub_exclusive_lease(key = nil, uuid = 'uuid', renew: false, timeout: nil)
+    key     ||= instance_of(String)
+    timeout ||= instance_of(Integer)
+
+    lease = instance_double(
+      Gitlab::ExclusiveLease,
+      try_obtain: uuid,
+      exists?: true,
+      renew: renew
+    )
+
+    allow(Gitlab::ExclusiveLease)
+      .to receive(:new)
+      .with(key, timeout: timeout)
+      .and_return(lease)
+
+    lease
+  end
+
+  def stub_exclusive_lease_taken(key = nil, timeout: nil)
+    stub_exclusive_lease(key, nil, timeout: timeout)
+  end
+
+  def expect_to_obtain_exclusive_lease(key, uuid = 'uuid', timeout: nil)
+    lease = stub_exclusive_lease(key, uuid, timeout: timeout)
+
+    expect(lease).to receive(:try_obtain)
+  end
+
+  def expect_to_cancel_exclusive_lease(key, uuid)
+    expect(Gitlab::ExclusiveLease)
+      .to receive(:cancel)
+      .with(key, uuid)
+  end
+end

spec/support/shared_examples/ci_trace_shared_examples.rb

@@ -247,8 +247,10 @@ shared_examples_for 'common trace features' do
         end
 
         context 'when another process has already been archiving', :clean_gitlab_redis_shared_state do
+          include ExclusiveLeaseHelpers
+
           before do
-            Gitlab::ExclusiveLease.new("trace:archive:#{trace.job.id}", timeout: 1.hour).try_obtain
+            stub_exclusive_lease_taken("trace:archive:#{trace.job.id}", timeout: 1.hour)
           end
 
           it 'blocks concurrent archiving' do

spec/support/shared_examples/uploaders/object_storage_shared_examples.rb

@@ -76,8 +76,10 @@ shared_examples "migrates" do |to_store:, from_store: nil|
   end
 
   context 'when migrate! is occupied by another process' do
+    include ExclusiveLeaseHelpers
+
     before do
-      @uuid = Gitlab::ExclusiveLease.new(subject.exclusive_lease_key, timeout: 1.hour.to_i).try_obtain
+      stub_exclusive_lease_taken(subject.exclusive_lease_key, timeout: 1.hour.to_i)
     end
 
     it 'does not execute migrate!' do
@@ -91,10 +93,6 @@ shared_examples "migrates" do |to_store:, from_store: nil|
 
       expect { subject.use_file }.to raise_error(ObjectStorage::ExclusiveLeaseTaken)
     end
-
-    after do
-      Gitlab::ExclusiveLease.cancel(subject.exclusive_lease_key, @uuid)
-    end
   end
 
   context 'migration is unsuccessful' do

spec/workers/project_cache_worker_spec.rb

 require 'spec_helper'
 
 describe ProjectCacheWorker do
+  include ExclusiveLeaseHelpers
+
   let(:worker) { described_class.new }
   let(:project) { create(:project, :repository) }
   let(:statistics) { project.statistics }
+  let(:lease_key) { "project_cache_worker:#{project.id}:update_statistics" }
+  let(:lease_timeout) { ProjectCacheWorker::LEASE_TIMEOUT }
 
   describe '#perform' do
     before do
-      allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain)
-        .and_return(true)
+      stub_exclusive_lease(lease_key, timeout: lease_timeout)
     end
 
     context 'with a non-existing project' do
@@ -79,9 +82,7 @@ describe ProjectCacheWorker do
   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)
+        stub_exclusive_lease_taken(lease_key, timeout: lease_timeout)
 
         expect(statistics).not_to receive(:refresh!)
 
@@ -91,9 +92,7 @@ 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)
-          .and_return(true)
+        stub_exclusive_lease(lease_key, timeout: lease_timeout)
 
         expect(statistics).to receive(:refresh!)
           .with(only: %i(repository_size))

spec/workers/project_migrate_hashed_storage_worker_spec.rb

 require 'spec_helper'
 
 describe ProjectMigrateHashedStorageWorker, :clean_gitlab_redis_shared_state do
+  include ExclusiveLeaseHelpers
+
   describe '#perform' do
     let(:project) { create(:project, :empty_repo) }
-    let(:pending_delete_project) { create(:project, :empty_repo, pending_delete: true) }
+    let(:lease_key) { "project_migrate_hashed_storage_worker:#{project.id}" }
+    let(:lease_timeout) { ProjectMigrateHashedStorageWorker::LEASE_TIMEOUT }
+
+    it 'skips when project no longer exists' do
+      expect(::Projects::HashedStorageMigrationService).not_to receive(:new)
+
+      subject.perform(-1)
+    end
 
-    context 'when have exclusive lease' do
-      before do
-        lease = subject.lease_for(project.id)
+    it 'skips when project is pending delete' do
+      pending_delete_project = create(:project, :empty_repo, pending_delete: true)
 
-        allow(Gitlab::ExclusiveLease).to receive(:new).and_return(lease)
-        allow(lease).to receive(:try_obtain).and_return(true)
-      end
+      expect(::Projects::HashedStorageMigrationService).not_to receive(:new)
 
-      it 'skips when project no longer exists' do
-        nonexistent_id = 999999999999
+      subject.perform(pending_delete_project.id)
+    end
 
-        expect(::Projects::HashedStorageMigrationService).not_to receive(:new)
-        subject.perform(nonexistent_id)
-      end
+    it 'delegates removal to service class when have exclusive lease' do
+      stub_exclusive_lease(lease_key, 'uuid', timeout: lease_timeout)
 
-      it 'skips when project is pending delete' do
-        expect(::Projects::HashedStorageMigrationService).not_to receive(:new)
+      migration_service = spy
 
-        subject.perform(pending_delete_project.id)
-      end
+      allow(::Projects::HashedStorageMigrationService)
+        .to receive(:new).with(project, subject.logger)
+        .and_return(migration_service)
 
-      it 'delegates removal to service class' do
-        service = double('service')
-        expect(::Projects::HashedStorageMigrationService).to receive(:new).with(project, subject.logger).and_return(service)
-        expect(service).to receive(:execute)
+      subject.perform(project.id)
 
-        subject.perform(project.id)
-      end
+      expect(migration_service).to have_received(:execute)
     end
 
-    context 'when dont have exclusive lease' do
-      before do
-        lease = subject.lease_for(project.id)
-
-        allow(Gitlab::ExclusiveLease).to receive(:new).and_return(lease)
-        allow(lease).to receive(:try_obtain).and_return(false)
-      end
+    it 'skips when dont have lease when dont have exclusive lease' do
+      stub_exclusive_lease_taken(lease_key, timeout: lease_timeout)
 
-      it 'skips when dont have lease' do
-        expect(::Projects::HashedStorageMigrationService).not_to receive(:new)
+      expect(::Projects::HashedStorageMigrationService).not_to receive(:new)
 
-        subject.perform(project.id)
-      end
+      subject.perform(project.id)
     end
   end
 end

spec/workers/propagate_service_template_worker_spec.rb

 require 'spec_helper'
 
 describe PropagateServiceTemplateWorker do
-  let!(:service_template) do
-    PushoverService.create(
-      template: true,
-      active: true,
-      properties: {
-        device: 'MyDevice',
-        sound: 'mic',
-        priority: 4,
-        user_key: 'asdf',
-        api_key: '123456789'
-      })
-  end
-
-  before do
-    allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain)
-      .and_return(true)
-  end
+  include ExclusiveLeaseHelpers
 
   describe '#perform' do
     it 'calls the propagate service with the template' do
-      expect(Projects::PropagateServiceTemplate).to receive(:propagate).with(service_template)
+      template = PushoverService.create(
+        template: true,
+        active: true,
+        properties: {
+          device: 'MyDevice',
+          sound: 'mic',
+          priority: 4,
+          user_key: 'asdf',
+          api_key: '123456789'
+        })
+
+      stub_exclusive_lease("propagate_service_template_worker:#{template.id}",
+        timeout: PropagateServiceTemplateWorker::LEASE_TIMEOUT)
+
+      expect(Projects::PropagateServiceTemplate)
+        .to receive(:propagate)
+        .with(template)
 
-      subject.perform(service_template.id)
+      subject.perform(template.id)
     end
   end
 end

spec/workers/repository_remove_remote_worker_spec.rb

 require 'rails_helper'
 
 describe RepositoryRemoveRemoteWorker do
-  subject(:worker) { described_class.new }
+  include ExclusiveLeaseHelpers
 
   describe '#perform' do
-    let(:remote_name) { 'joe'}
     let!(:project) { create(:project, :repository) }
+    let(:remote_name) { 'joe'}
+    let(:lease_key) { "remove_remote_#{project.id}_#{remote_name}" }
+    let(:lease_timeout) { RepositoryRemoveRemoteWorker::LEASE_TIMEOUT }
 
-    context 'when it cannot obtain lease' do
-      it 'logs error' do
-        allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain) { nil }
-
-        expect_any_instance_of(Repository).not_to receive(:remove_remote)
-        expect(worker).to receive(:log_error).with('Cannot obtain an exclusive lease. There must be another instance already in execution.')
-
-        worker.perform(project.id, remote_name)
-      end
+    it 'returns nil when project does not exist' do
+      expect(subject.perform(-1, 'remote_name')).to be_nil
     end
 
-    context 'when it gets the lease' do
+    context 'when project exists' do
       before do
-        allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain).and_return(true)
+        allow(Project)
+          .to receive(:find_by)
+          .with(id: project.id)
+          .and_return(project)
       end
 
-      context 'when project does not exist' do
-        it 'returns nil' do
-          expect(worker.perform(-1, 'remote_name')).to be_nil
-        end
-      end
+      it 'does not remove remote when cannot obtain lease' do
+        stub_exclusive_lease_taken(lease_key, timeout: lease_timeout)
+
+        expect(project.repository)
+          .not_to receive(:remove_remote)
 
-      context 'when project exists' do
-        it 'removes remote from repository' do
-          masterrev = project.repository.find_branch('master').dereferenced_target
+        expect(subject)
+          .to receive(:log_error)
+          .with('Cannot obtain an exclusive lease. There must be another instance already in execution.')
 
-          create_remote_branch(remote_name, 'remote_branch', masterrev)
+        subject.perform(project.id, remote_name)
+      end
 
-          expect_any_instance_of(Repository).to receive(:remove_remote).with(remote_name).and_call_original
+      it 'removes remote from repository when obtain a lease' do
+        stub_exclusive_lease(lease_key, timeout: lease_timeout)
+        masterrev = project.repository.find_branch('master').dereferenced_target
+        create_remote_branch(remote_name, 'remote_branch', masterrev)
 
-          worker.perform(project.id, remote_name)
-        end
+        expect(project.repository)
+          .to receive(:remove_remote)
+          .with(remote_name)
+          .and_call_original
+
+        subject.perform(project.id, remote_name)
       end
     end
   end
@@ -47,6 +53,7 @@ describe RepositoryRemoveRemoteWorker do
     rugged = Gitlab::GitalyClient::StorageSettings.allow_disk_access do
       project.repository.rugged
     end
+
     rugged.references.create("refs/remotes/#{remote_name}/#{branch_name}", target.id)
   end
 end

spec/workers/stuck_ci_jobs_worker_spec.rb

 require 'spec_helper'
 
 describe StuckCiJobsWorker do
+  include ExclusiveLeaseHelpers
+
   let!(:runner) { create :ci_runner }
   let!(:job) { create :ci_build, runner: runner }
-  let(:worker) { described_class.new }
-  let(:exclusive_lease_uuid) { SecureRandom.uuid }
+  let(:trace_lease_key) { "trace:archive:#{job.id}" }
+  let(:trace_lease_uuid) { SecureRandom.uuid }
+  let(:worker_lease_key) { StuckCiJobsWorker::EXCLUSIVE_LEASE_KEY }
+  let(:worker_lease_uuid) { SecureRandom.uuid }
+
+  subject(:worker) { described_class.new }
 
   before do
+    stub_exclusive_lease(worker_lease_key, worker_lease_uuid)
+    stub_exclusive_lease(trace_lease_key, trace_lease_uuid)
     job.update!(status: status, updated_at: updated_at)
-    allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain).and_return(exclusive_lease_uuid)
   end
 
   shared_examples 'job is dropped' do
@@ -44,16 +51,19 @@ describe StuckCiJobsWorker do
 
       context 'when job was not updated for more than 1 day ago' do
         let(:updated_at) { 2.days.ago }
+
         it_behaves_like 'job is dropped'
       end
 
       context 'when job was updated in less than 1 day ago' do
         let(:updated_at) { 6.hours.ago }
+
         it_behaves_like 'job is unchanged'
       end
 
       context 'when job was not updated for more than 1 hour ago' do
         let(:updated_at) { 2.hours.ago }
+
         it_behaves_like 'job is unchanged'
       end
     end
@@ -65,11 +75,14 @@ describe StuckCiJobsWorker do
 
       context 'when job was not updated for more than 1 hour ago' do
         let(:updated_at) { 2.hours.ago }
+
         it_behaves_like 'job is dropped'
       end
 
-      context 'when job was updated in less than 1 hour ago' do
+      context 'when job was updated in less than 1
+       hour ago' do
         let(:updated_at) { 30.minutes.ago }
+
         it_behaves_like 'job is unchanged'
       end
     end
@@ -80,11 +93,13 @@ describe StuckCiJobsWorker do
 
     context 'when job was not updated for more than 1 hour ago' do
       let(:updated_at) { 2.hours.ago }
+
       it_behaves_like 'job is dropped'
     end
 
     context 'when job was updated in less than 1 hour ago' do
       let(:updated_at) { 30.minutes.ago }
+
       it_behaves_like 'job is unchanged'
     end
   end
@@ -93,6 +108,7 @@ describe StuckCiJobsWorker do
     context "when job is #{status}" do
       let(:status) { status }
       let(:updated_at) { 2.days.ago }
+
       it_behaves_like 'job is unchanged'
     end
   end
@@ -119,23 +135,27 @@ describe StuckCiJobsWorker do
     it 'is guard by exclusive lease when executed concurrently' do
       expect(worker).to receive(:drop).at_least(:once).and_call_original
       expect(worker2).not_to receive(:drop)
+
       worker.perform
-      allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain).and_return(false)
+
+      stub_exclusive_lease_taken(worker_lease_key)
+
       worker2.perform
     end
 
     it 'can be executed in sequence' do
       expect(worker).to receive(:drop).at_least(:once).and_call_original
       expect(worker2).to receive(:drop).at_least(:once).and_call_original
+
       worker.perform
       worker2.perform
     end
 
-    it 'cancels exclusive lease after worker perform' do
-      worker.perform
+    it 'cancels exclusive leases after worker perform' do
+      expect_to_cancel_exclusive_lease(trace_lease_key, trace_lease_uuid)
+      expect_to_cancel_exclusive_lease(worker_lease_key, worker_lease_uuid)
 
-      expect(Gitlab::ExclusiveLease.new(described_class::EXCLUSIVE_LEASE_KEY, timeout: 1.hour))
-        .not_to be_exists
+      worker.perform
     end
   end
 end