From 93a964d449ad29e94e785209d7ecde217a8e9b25 Mon Sep 17 00:00:00 2001
From: Shinya Maeda <shinya@gitlab.com>
Date: Tue, 3 Jul 2018 16:20:27 +0900
Subject: [PATCH] Add spec for ExclusiveLeaseHelpers

---
 app/models/ci/build_trace_chunk.rb            |  2 +-
 app/services/concerns/exclusive_lease_lock.rb | 22 ------
 lib/gitlab/exclusive_lease_helpers.rb         | 29 +++++++
 .../gitlab/exclusive_lease_helpers_spec.rb    | 76 +++++++++++++++++++
 4 files changed, 106 insertions(+), 23 deletions(-)
 delete mode 100644 app/services/concerns/exclusive_lease_lock.rb
 create mode 100644 lib/gitlab/exclusive_lease_helpers.rb
 create mode 100644 spec/lib/gitlab/exclusive_lease_helpers_spec.rb

diff --git a/app/models/ci/build_trace_chunk.rb b/app/models/ci/build_trace_chunk.rb
index b724d0cb517..43bede7638c 100644
--- a/app/models/ci/build_trace_chunk.rb
+++ b/app/models/ci/build_trace_chunk.rb
@@ -1,7 +1,7 @@
 module Ci
   class BuildTraceChunk < ActiveRecord::Base
     include FastDestroyAll
-    include ExclusiveLeaseLock
+    include ::Gitlab::ExclusiveLeaseHelpers
     extend Gitlab::Ci::Model
 
     belongs_to :build, class_name: "Ci::Build", foreign_key: :build_id
diff --git a/app/services/concerns/exclusive_lease_lock.rb b/app/services/concerns/exclusive_lease_lock.rb
deleted file mode 100644
index 8da54b0d15d..00000000000
--- a/app/services/concerns/exclusive_lease_lock.rb
+++ /dev/null
@@ -1,22 +0,0 @@
-module ExclusiveLeaseLock
-  extend ActiveSupport::Concern
-
-  FailedToObtainLockError = Class.new(StandardError)
-
-  def in_lock(key, ttl: 1.minute, retries: 10, sleep_sec: 0.01.seconds)
-    lease = Gitlab::ExclusiveLease.new(key, timeout: ttl)
-
-    until uuid = lease.try_obtain
-      # Keep trying until we obtain the lease. To prevent hammering Redis too
-      # much we'll wait for a bit.
-      sleep(sleep_sec)
-      break if (retries -= 1) < 0
-    end
-
-    raise FailedToObtainLockError, 'Failed to obtain a lock' unless uuid
-
-    return yield
-  ensure
-    Gitlab::ExclusiveLease.cancel(key, uuid)
-  end
-end
diff --git a/lib/gitlab/exclusive_lease_helpers.rb b/lib/gitlab/exclusive_lease_helpers.rb
new file mode 100644
index 00000000000..ab6838adc6d
--- /dev/null
+++ b/lib/gitlab/exclusive_lease_helpers.rb
@@ -0,0 +1,29 @@
+module Gitlab
+  # This module provides helper methods which are intregrated with GitLab::ExclusiveLease
+  module ExclusiveLeaseHelpers
+    FailedToObtainLockError = Class.new(StandardError)
+
+    ##
+    # This helper method blocks a process/thread until the other process cancel the obrainted lease key.
+    #
+    # Note: It's basically discouraged to use this method in the unicorn's thread,
+    #       because it holds the connection until all `retries` is consumed.
+    #       This could potentially eat up all connection pools.
+    def in_lock(key, ttl: 1.minute, retries: 10, sleep_sec: 0.01.seconds)
+      lease = Gitlab::ExclusiveLease.new(key, timeout: ttl)
+
+      until uuid = lease.try_obtain
+        # Keep trying until we obtain the lease. To prevent hammering Redis too
+        # much we'll wait for a bit.
+        sleep(sleep_sec)
+        break if (retries -= 1) < 0
+      end
+
+      raise FailedToObtainLockError, 'Failed to obtain a lock' unless uuid
+
+      return yield
+    ensure
+      Gitlab::ExclusiveLease.cancel(key, uuid)
+    end
+  end
+end
diff --git a/spec/lib/gitlab/exclusive_lease_helpers_spec.rb b/spec/lib/gitlab/exclusive_lease_helpers_spec.rb
new file mode 100644
index 00000000000..eb2143f9d1d
--- /dev/null
+++ b/spec/lib/gitlab/exclusive_lease_helpers_spec.rb
@@ -0,0 +1,76 @@
+require 'spec_helper'
+
+describe Gitlab::ExclusiveLeaseHelpers, :clean_gitlab_redis_shared_state do
+  include ::ExclusiveLeaseHelpers
+
+  let(:class_instance) { (Class.new { include ::Gitlab::ExclusiveLeaseHelpers }).new }
+  let(:unique_key) { SecureRandom.hex(10) }
+
+  describe '#in_lock' do
+    subject { class_instance.in_lock(unique_key, **options) { } }
+
+    let(:options) { { } }
+
+    context 'when the lease is not obtained yet' do
+      before do
+        stub_exclusive_lease(unique_key, 'uuid')
+      end
+
+      it 'calls the given block' do
+        expect { |b| class_instance.in_lock(unique_key, &b) }.to yield_control.once
+      end
+
+      it 'calls the given block continuously' do
+        expect { |b| class_instance.in_lock(unique_key, &b) }.to yield_control.once
+        expect { |b| class_instance.in_lock(unique_key, &b) }.to yield_control.once
+        expect { |b| class_instance.in_lock(unique_key, &b) }.to yield_control.once
+      end
+
+      it 'cancels the exclusive lease after the block' do
+        expect_to_cancel_exclusive_lease(unique_key, 'uuid')
+
+        subject
+      end
+    end
+
+    context 'when the lease is obtained already' do
+      let!(:lease) { stub_exclusive_lease_taken(unique_key) }
+
+      it 'retries to obtain a lease and raises an error' do
+        expect(lease).to receive(:try_obtain).exactly(11).times
+
+        expect { subject }.to raise_error('Failed to obtain a lock')
+      end
+
+      context 'when ttl is specified' do
+        let(:options) { { ttl: 10.minute } }
+
+        it 'receives the specified argument' do
+          expect(Gitlab::ExclusiveLease).to receive(:new).with(unique_key, { timeout: 10.minute } )
+  
+          expect { subject }.to raise_error('Failed to obtain a lock')
+        end
+      end
+
+      context 'when retry count is specified' do
+        let(:options) { { retries: 3 } }
+
+        it 'retries for the specified times' do
+          expect(lease).to receive(:try_obtain).exactly(4).times
+  
+          expect { subject }.to raise_error('Failed to obtain a lock')
+        end
+      end
+
+      context 'when sleep second is specified' do
+        let(:options) { { retries: 0, sleep_sec: 0.05.second } }
+
+        it 'receives the specified argument' do
+          expect(class_instance).to receive(:sleep).with(0.05.second).once
+  
+          expect { subject }.to raise_error('Failed to obtain a lock')
+        end
+      end
+    end
+  end
+end