Refactor sleep logic into state object

This adds a new class, Gitlab::ExclusiveLeaseHelpers::SleepingLock,
which abstracts the sleep and retry logic needed to obtain a
Gitlab::ExclusiveLease successfully. This abstraction transforms an
imperative loop into OO messages, aiming to maximise correctness and
intelligibility.

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

@@ -94,6 +94,13 @@ module Gitlab
         ttl if ttl.positive?
       end
     end
+
+    # Gives up this lease, allowing it to be obtained by others.
+    def cancel
+      Gitlab::Redis::SharedState.with do |redis|
+        redis.eval(LUA_CANCEL_SCRIPT, keys: [@redis_shared_state_key], argv: [@uuid])
+      end
+    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.attempts > 1)
     ensure
-      Gitlab::ExclusiveLease.cancel(key, uuid)
+      lease&.release
     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
+      attr_reader :attempts
+
+      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 release
+        @lease.cancel
+      end
+
+      private
+
+      attr_reader :delay
+
+      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 '#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.attempts).to eq(1)
+      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 'records the successful attempt number' 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.attempts).to eq(4)
+        end
+      end
+    end
+
+    describe 'release' do
+      let!(:lease) { stub_exclusive_lease(key, 'uuid') }
+
+      it 'cancels the lease' do
+        expect(lease).to receive(:cancel)
+
+        subject.release
+      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

@@ -76,6 +76,31 @@ describe Gitlab::ExclusiveLease, :clean_gitlab_redis_shared_state do
     end
   end
 
+  describe '#cancel' do
+    it 'can cancel a lease' do
+      lease = new_lease(unique_key)
+      uuid = lease.try_obtain
+      expect(uuid).to be_present
+      expect(new_lease(unique_key).try_obtain).to eq(false)
+
+      lease.cancel
+
+      expect(new_lease(unique_key).try_obtain).to be_present
+    end
+
+    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
+
+    def new_lease(key)
+      described_class.new(key, timeout: 3600)
+    end
+  end
+
   describe '#ttl' do
     it 'returns the TTL of the Redis key' do
       lease = described_class.new('kittens', timeout: 100)

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)