Add tests for optimistic locking

5 files changed, 38 insertions, 7 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index e61165b2826..0e237ac885a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -66,6 +66,7 @@ Please view this file on the master branch, on stable branches it's out of date.
   - Updating verbiage on git basics to be more intuitive
   - Fix project_feature record not generated on project creation
   - Clarify documentation for Runners API (Gennady Trafimenkov)
+  - Use optimistic locking for pipelines and builds
   - The instrumentation for Banzai::Renderer has been restored
   - Change user & group landing page routing from /u/:username to /:username
   - Added documentation for .gitattributes files
diff --git a/app/models/ci/pipeline.rb b/app/models/ci/pipeline.rb
index 27178e2a6c1..d3432632899 100644
--- a/app/models/ci/pipeline.rb
+++ b/app/models/ci/pipeline.rb
@@ -260,7 +260,7 @@ module Ci
     end
 
     def update_status
-      Gitlab::OptimisticLocking.retry_lock(build) do
+      Gitlab::OptimisticLocking.retry_lock(self) do
         case latest_builds_status
         when 'pending' then enqueue
         when 'running' then run
diff --git a/app/services/ci/register_build_service.rb b/app/services/ci/register_build_service.rb
index 8d3bc8e2dee..74b5ebf372b 100644
--- a/app/services/ci/register_build_service.rb
+++ b/app/services/ci/register_build_service.rb
@@ -35,7 +35,7 @@ module Ci
 
       build
 
-    rescue StateMachines::InvalidTransition, StaleObjectError
+    rescue StateMachines::InvalidTransition, ActiveRecord::StaleObjectError
       nil
     end
 
diff --git a/lib/gitlab/optimistic_locking.rb b/lib/gitlab/optimistic_locking.rb
index 7668c518bc5..17010d73c57 100644
--- a/lib/gitlab/optimistic_locking.rb
+++ b/lib/gitlab/optimistic_locking.rb
@@ -1,10 +1,12 @@
 module Gitlab
-  module OptimisticLocking
-    def retry_lock(subject, &block)
-      while true do
+  class OptimisticLocking
+    def self.retry_lock(subject, &block)
+      loop do
         begin
-          return yield subject
-        rescue StaleObjectError
+          subject.transaction do
+            return block.call(subject)
+          end
+        rescue ActiveRecord::StaleObjectError
           subject.reload
         end
       end
diff --git a/spec/lib/gitlab/optimistic_locking_spec.rb b/spec/lib/gitlab/optimistic_locking_spec.rb
new file mode 100644
index 00000000000..75c78cf077a
--- /dev/null
+++ b/spec/lib/gitlab/optimistic_locking_spec.rb
@@ -0,0 +1,28 @@
+require 'spec_helper'
+
+describe Gitlab::OptimisticLocking, lib: true do
+  describe '#retry_lock' do
+    let!(:pipeline) { create(:ci_pipeline) }
+    let!(:pipeline2) { Ci::Pipeline.find(pipeline.id) }
+
+    it 'does not reload object if state changes' do
+      expect(pipeline).not_to receive(:reload)
+      expect(pipeline).to receive(:succeed).and_call_original
+
+      described_class.retry_lock(pipeline) do |subject|
+        subject.succeed
+      end
+    end
+
+    it 'retries action if exception is raised' do
+      pipeline.succeed
+
+      expect(pipeline2).to receive(:reload).and_call_original
+      expect(pipeline2).to receive(:drop).twice.and_call_original
+
+      described_class.retry_lock(pipeline2) do |subject|
+        subject.drop
+      end
+    end
+  end
+end