Merge branch 'ajk-fix-extra-sleep-exclusive-lease' into 'master'

Remove extra sleep in exclusive lease helper

Closes #216050

See merge request gitlab-org/gitlab!30654

changelogs/unreleased/ajk-fix-extra-sleep-exclusive-lease.yml

+---
+title: Remove extra sleep when obtaining exclusive lease
+merge_request: 30654
+author:
+type: fixed

lib/gitlab/exclusive_lease.rb

@@ -12,6 +12,9 @@ module Gitlab
   # ExclusiveLease.
   #
   class ExclusiveLease
+    PREFIX = 'gitlab:exclusive_lease'
+    NoKey = Class.new(ArgumentError)
+
     LUA_CANCEL_SCRIPT = <<~EOS.freeze
       local key, uuid = KEYS[1], ARGV[1]
       if redis.call("get", key) == uuid then
@@ -34,13 +37,21 @@ module Gitlab
     end
 
     def self.cancel(key, uuid)
+      return unless key.present?
+
       Gitlab::Redis::SharedState.with do |redis|
-        redis.eval(LUA_CANCEL_SCRIPT, keys: [redis_shared_state_key(key)], argv: [uuid])
+        redis.eval(LUA_CANCEL_SCRIPT, keys: [ensure_prefixed_key(key)], argv: [uuid])
       end
     end
 
     def self.redis_shared_state_key(key)
-      "gitlab:exclusive_lease:#{key}"
+      "#{PREFIX}:#{key}"
+    end
+
+    def self.ensure_prefixed_key(key)
+      raise NoKey unless key.present?
+
+      key.start_with?(PREFIX) ? key : redis_shared_state_key(key)
     end
 
     # Removes any existing exclusive_lease from redis
@@ -94,6 +105,11 @@ module Gitlab
         ttl if ttl.positive?
       end
     end
+
+    # Gives up this lease, allowing it to be obtained by others.
+    def cancel
+      self.class.cancel(@redis_shared_state_key, @uuid)
+    end
   end
 end
 

lib/gitlab/exclusive_lease_helpers.rb

@@ -6,33 +6,38 @@ module Gitlab
     FailedToObtainLockError = Class.new(StandardError)
 
     ##
-    # This helper method blocks a process/thread until the other process cancel the obrainted lease key.
+    # This helper method blocks a process/thread until the lease can be acquired, either due to
+    # the lease TTL expiring, or due to the current holder explicitly releasing
+    # their hold.
     #
-    # Note: It's basically discouraged to use this method in the unicorn's thread,
-    #       because it holds the connection until all `retries` is consumed.
+    # If the lease cannot be obtained, raises `FailedToObtainLockError`.
+    #
+    # @param [String] key The lock the thread will try to acquire. Only one thread
+    #                     in one process across all Rails instances can hold this named lock at any
+    #                     one time.
+    # @param [Float] ttl: The length of time the lock will be valid for. The lock
+    #                    will be automatically be released after this time, so any work should be
+    #                    completed within this time.
+    # @param [Integer] retries: The maximum number of times we will re-attempt
+    #                           to acquire the lock. The maximum number of attempts will be `retries + 1`:
+    #                           one for the initial attempt, and then one for every re-try.
+    # @param [Float|Proc] sleep_sec: Either a number of seconds to sleep, or
+    #                                a proc that computes the sleep time given the number of preceding attempts
+    #                               (from 1 to retries - 1)
+    #
+    # Note: It's basically discouraged to use this method in a unicorn thread,
+    #       because this ties up all thread related resources until all `retries` are consumed.
     #       This could potentially eat up all connection pools.
     def in_lock(key, ttl: 1.minute, retries: 10, sleep_sec: 0.01.seconds)
       raise ArgumentError, 'Key needs to be specified' unless key
 
-      lease = Gitlab::ExclusiveLease.new(key, timeout: ttl)
-      retried = false
-      max_attempts = 1 + retries
-
-      until uuid = lease.try_obtain
-        # Keep trying until we obtain the lease. To prevent hammering Redis too
-        # much we'll wait for a bit.
-        attempt_number = max_attempts - retries
-        delay = sleep_sec.respond_to?(:call) ? sleep_sec.call(attempt_number) : sleep_sec
-
-        sleep(delay)
-        (retries -= 1) < 0 ? break : retried ||= true
-      end
+      lease = SleepingLock.new(key, timeout: ttl, delay: sleep_sec)
 
-      raise FailedToObtainLockError, 'Failed to obtain a lock' unless uuid
+      lease.obtain(1 + retries)
 
-      yield(retried)
+      yield(lease.retried?)
     ensure
-      Gitlab::ExclusiveLease.cancel(key, uuid)
+      lease&.cancel
     end
   end
 end

lib/gitlab/exclusive_lease_helpers/sleeping_lock.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module ExclusiveLeaseHelpers
+    # Wrapper around ExclusiveLease that adds retry logic
+    class SleepingLock
+      delegate :cancel, to: :@lease
+
+      def initialize(key, timeout:, delay:)
+        @lease = ::Gitlab::ExclusiveLease.new(key, timeout: timeout)
+        @delay = delay
+        @attempts = 0
+      end
+
+      def obtain(max_attempts)
+        until held?
+          raise FailedToObtainLockError, 'Failed to obtain a lock' if attempts >= max_attempts
+
+          sleep(sleep_sec) unless first_attempt?
+          try_obtain
+        end
+      end
+
+      def retried?
+        attempts > 1
+      end
+
+      private
+
+      attr_reader :delay, :attempts
+
+      def held?
+        @uuid.present?
+      end
+
+      def try_obtain
+        @uuid ||= @lease.try_obtain
+        @attempts += 1
+      end
+
+      def first_attempt?
+        attempts.zero?
+      end
+
+      def sleep_sec
+        delay.respond_to?(:call) ? delay.call(attempts) : delay
+      end
+    end
+  end
+end

spec/lib/gitlab/exclusive_lease_helpers/sleeping_lock_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe Gitlab::ExclusiveLeaseHelpers::SleepingLock, :clean_gitlab_redis_shared_state do
+  include ::ExclusiveLeaseHelpers
+
+  let(:timeout) { 1.second }
+  let(:delay) { 0.1.seconds }
+  let(:key) { SecureRandom.hex(10) }
+
+  subject { described_class.new(key, timeout: timeout, delay: delay) }
+
+  describe '#retried?' do
+    before do
+      stub_exclusive_lease(key, 'uuid')
+    end
+
+    context 'we have not made any attempts' do
+      it { is_expected.not_to be_retried }
+    end
+
+    context 'we just made a single (initial) attempt' do
+      it 'is not considered a retry' do
+        subject.send(:try_obtain)
+
+        is_expected.not_to be_retried
+      end
+    end
+
+    context 'made multiple attempts' do
+      it 'is considered a retry' do
+        2.times { subject.send(:try_obtain) }
+
+        is_expected.to be_retried
+      end
+    end
+  end
+
+  describe '#obtain' do
+    context 'when the lease is not held' do
+      before do
+        stub_exclusive_lease(key, 'uuid')
+      end
+
+      it 'obtains the lease on the first attempt, without sleeping' do
+        expect(subject).not_to receive(:sleep)
+
+        subject.obtain(10)
+
+        expect(subject).not_to be_retried
+      end
+    end
+
+    context 'when the lease is held elsewhere' do
+      let!(:lease) { stub_exclusive_lease_taken(key) }
+      let(:max_attempts) { 7 }
+
+      it 'retries to obtain a lease and raises an error' do
+        expect(subject).to receive(:sleep).with(delay).exactly(max_attempts - 1).times
+        expect(lease).to receive(:try_obtain).exactly(max_attempts).times
+
+        expect { subject.obtain(max_attempts) }.to raise_error('Failed to obtain a lock')
+      end
+
+      context 'when the delay is computed from the attempt number' do
+        let(:delay) { ->(n) { 3 * n } }
+
+        it 'uses the computation to determine the sleep length' do
+          expect(subject).to receive(:sleep).with(3).once
+          expect(subject).to receive(:sleep).with(6).once
+          expect(subject).to receive(:sleep).with(9).once
+          expect(lease).to receive(:try_obtain).exactly(4).times
+
+          expect { subject.obtain(4) }.to raise_error('Failed to obtain a lock')
+        end
+      end
+
+      context 'when lease is granted after retry' do
+        it 'knows that it retried' do
+          expect(subject).to receive(:sleep).with(delay).exactly(3).times
+          expect(lease).to receive(:try_obtain).exactly(3).times { nil }
+          expect(lease).to receive(:try_obtain).once { 'obtained' }
+
+          subject.obtain(max_attempts)
+
+          expect(subject).to be_retried
+        end
+      end
+    end
+
+    describe 'cancel' do
+      let!(:lease) { stub_exclusive_lease(key, 'uuid') }
+
+      it 'cancels the lease' do
+        expect(lease).to receive(:cancel)
+
+        subject.cancel
+      end
+    end
+  end
+end

spec/lib/gitlab/exclusive_lease_helpers_spec.rb

@@ -22,9 +22,7 @@ describe Gitlab::ExclusiveLeaseHelpers, :clean_gitlab_redis_shared_state do
     end
 
     context 'when the lease is not obtained yet' do
-      before do
-        stub_exclusive_lease(unique_key, 'uuid')
-      end
+      let!(:lease) { stub_exclusive_lease(unique_key, 'uuid') }
 
       it 'calls the given block' do
         expect { |b| class_instance.in_lock(unique_key, &b) }.to yield_with_args(false)
@@ -37,7 +35,7 @@ describe Gitlab::ExclusiveLeaseHelpers, :clean_gitlab_redis_shared_state do
       end
 
       it 'cancels the exclusive lease after the block' do
-        expect_to_cancel_exclusive_lease(unique_key, 'uuid')
+        expect(lease).to receive(:cancel).once
 
         subject
       end
@@ -81,11 +79,21 @@ describe Gitlab::ExclusiveLeaseHelpers, :clean_gitlab_redis_shared_state do
         end
       end
 
+      context 'when we specify no retries' do
+        let(:options) { { retries: 0 } }
+
+        it 'never sleeps' do
+          expect(class_instance).not_to receive(:sleep)
+
+          expect { subject }.to raise_error('Failed to obtain a lock')
+        end
+      end
+
       context 'when sleep second is specified' do
         let(:options) { { retries: 1, sleep_sec: 0.05.seconds } }
 
         it 'receives the specified argument' do
-          expect(class_instance).to receive(:sleep).with(0.05.seconds).twice
+          expect_any_instance_of(Object).to receive(:sleep).with(0.05.seconds).once
 
           expect { subject }.to raise_error('Failed to obtain a lock')
         end
@@ -95,9 +103,8 @@ describe Gitlab::ExclusiveLeaseHelpers, :clean_gitlab_redis_shared_state do
         let(:options) { { retries: 2, sleep_sec: ->(num) { 0.1 + num } } }
 
         it 'receives the specified argument' do
-          expect(class_instance).to receive(:sleep).with(1.1.seconds).once
-          expect(class_instance).to receive(:sleep).with(2.1.seconds).once
-          expect(class_instance).to receive(:sleep).with(3.1.seconds).once
+          expect_any_instance_of(Object).to receive(:sleep).with(1.1.seconds).once
+          expect_any_instance_of(Object).to receive(:sleep).with(2.1.seconds).once
 
           expect { subject }.to raise_error('Failed to obtain a lock')
         end

spec/lib/gitlab/exclusive_lease_spec.rb

@@ -21,6 +21,27 @@ describe Gitlab::ExclusiveLease, :clean_gitlab_redis_shared_state do
     end
   end
 
+  describe '.redis_shared_state_key' do
+    it 'provides a namespaced key' do
+      expect(described_class.redis_shared_state_key(unique_key))
+        .to start_with(described_class::PREFIX)
+        .and include(unique_key)
+    end
+  end
+
+  describe '.ensure_prefixed_key' do
+    it 'does not double prefix a key' do
+      prefixed = described_class.redis_shared_state_key(unique_key)
+
+      expect(described_class.ensure_prefixed_key(unique_key))
+        .to eq(described_class.ensure_prefixed_key(prefixed))
+    end
+
+    it 'raises errors when there is no key' do
+      expect { described_class.ensure_prefixed_key(nil) }.to raise_error(described_class::NoKey)
+    end
+  end
+
   describe '#renew' do
     it 'returns true when we have the existing lease' do
       lease = described_class.new(unique_key, timeout: 3600)
@@ -61,18 +82,61 @@ describe Gitlab::ExclusiveLease, :clean_gitlab_redis_shared_state do
     end
   end
 
-  describe '.cancel' do
-    it 'can cancel a lease' do
-      uuid = new_lease(unique_key)
+  describe 'cancellation' do
+    def new_lease(key)
+      described_class.new(key, timeout: 3600)
+    end
+
+    shared_examples 'cancelling a lease' do
+      let(:lease) { new_lease(unique_key) }
+
+      it 'releases the held lease' do
+        uuid = lease.try_obtain
         expect(uuid).to be_present
-      expect(new_lease(unique_key)).to eq(false)
+        expect(new_lease(unique_key).try_obtain).to eq(false)
 
-      described_class.cancel(unique_key, uuid)
-      expect(new_lease(unique_key)).to be_present
+        cancel_lease(uuid)
+
+        expect(new_lease(unique_key).try_obtain).to be_present
+      end
     end
 
-    def new_lease(key)
-      described_class.new(key, timeout: 3600).try_obtain
+    describe '.cancel' do
+      def cancel_lease(uuid)
+        described_class.cancel(release_key, uuid)
+      end
+
+      context 'when called with the unprefixed key' do
+        it_behaves_like 'cancelling a lease' do
+          let(:release_key) { unique_key }
+        end
+      end
+
+      context 'when called with the prefixed key' do
+        it_behaves_like 'cancelling a lease' do
+          let(:release_key) { described_class.redis_shared_state_key(unique_key) }
+        end
+      end
+
+      it 'does not raise errors when given a nil key' do
+        expect { described_class.cancel(nil, nil) }.not_to raise_error
+      end
+    end
+
+    describe '#cancel' do
+      def cancel_lease(_uuid)
+        lease.cancel
+      end
+
+      it_behaves_like 'cancelling a lease'
+
+      it 'is safe to call even if the lease was never obtained' do
+        lease = new_lease(unique_key)
+
+        lease.cancel
+
+        expect(new_lease(unique_key).try_obtain).to be_present
+      end
     end
   end
 

spec/support/helpers/exclusive_lease_helpers.rb

@@ -9,7 +9,8 @@ module ExclusiveLeaseHelpers
       Gitlab::ExclusiveLease,
       try_obtain: uuid,
       exists?: true,
-      renew: renew
+      renew: renew,
+      cancel: nil
     )
 
     allow(Gitlab::ExclusiveLease)