Allow asynchronous rebase operations to be monitored

This MR introduces tracking of the `rebase_jid` for merge requests. As
with `merge_ongoing?`, `rebase_in_progress?` will now return true if a
rebase is proceeding in sidekiq.

After one release, we should remove the Gitaly-based lookup of rebases.
It is much better to track this kind of thing via the database.

14 files changed, 188 insertions, 38 deletions

diff --git a/app/controllers/projects/merge_requests_controller.rb b/app/controllers/projects/merge_requests_controller.rb
index 7ee8e0ea8f8..2aa2508be16 100644
--- a/app/controllers/projects/merge_requests_controller.rb
+++ b/app/controllers/projects/merge_requests_controller.rb
@@ -201,7 +201,7 @@ class Projects::MergeRequestsController < Projects::MergeRequests::ApplicationCo
   end
 
   def rebase
-    RebaseWorker.perform_async(@merge_request.id, current_user.id)
+    @merge_request.rebase_async(current_user.id)
 
     head :ok
   end
diff --git a/app/models/merge_request.rb b/app/models/merge_request.rb
index 8391d526d18..e96e26cc773 100644
--- a/app/models/merge_request.rb
+++ b/app/models/merge_request.rb
@@ -223,7 +223,13 @@ class MergeRequest < ApplicationRecord
   end
 
   def rebase_in_progress?
-    strong_memoize(:rebase_in_progress) do
+    (rebase_jid.present? && Gitlab::SidekiqStatus.running?(rebase_jid)) ||
+      gitaly_rebase_in_progress?
+  end
+
+  # TODO: remove the Gitaly lookup after v12.1, when rebase_jid will be reliable
+  def gitaly_rebase_in_progress?
+    strong_memoize(:gitaly_rebase_in_progress) do
       # The source project can be deleted
       next false unless source_project
 
@@ -389,6 +395,26 @@ class MergeRequest < ApplicationRecord
     update_column(:merge_jid, jid)
   end
 
+  # Set off a rebase asynchronously, atomically updating the `rebase_jid` of
+  # the MR so that the status of the operation can be tracked.
+  def rebase_async(user_id)
+    transaction do
+      lock!
+
+      raise ActiveRecord::StaleObjectError if !open? || rebase_in_progress?
+
+      # Although there is a race between setting rebase_jid here and clearing it
+      # in the RebaseWorker, it can't do any harm since we check both that the
+      # attribute is set *and* that the sidekiq job is still running. So a JID
+      # for a completed RebaseWorker is equivalent to a nil JID.
+      jid = Sidekiq::Worker.skipping_transaction_check do
+        RebaseWorker.perform_async(id, user_id)
+      end
+
+      update_column(:rebase_jid, jid)
+    end
+  end
+
   def merge_participants
     participants = [author]
 
diff --git a/app/services/merge_requests/rebase_service.rb b/app/services/merge_requests/rebase_service.rb
index 4b9921c28ba..8d3b9b05819 100644
--- a/app/services/merge_requests/rebase_service.rb
+++ b/app/services/merge_requests/rebase_service.rb
@@ -15,7 +15,7 @@ module MergeRequests
     end
 
     def rebase
-      if merge_request.rebase_in_progress?
+      if merge_request.gitaly_rebase_in_progress?
         log_error('Rebase task canceled: Another rebase is already in progress', save_message_on_model: true)
         return false
       end
@@ -27,6 +27,8 @@ module MergeRequests
       log_error(REBASE_ERROR, save_message_on_model: true)
       log_error(e.message)
       false
+    ensure
+      merge_request.update_column(:rebase_jid, nil)
     end
   end
 end
diff --git a/app/workers/rebase_worker.rb b/app/workers/rebase_worker.rb
index a6baebc1443..8d06adcd993 100644
--- a/app/workers/rebase_worker.rb
+++ b/app/workers/rebase_worker.rb
@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+# The RebaseWorker must be wrapped in important concurrency code, so should only
+# be scheduled via MergeRequest#rebase_async
 class RebaseWorker
   include ApplicationWorker
 
diff --git a/changelogs/unreleased/54117-transactional-rebase.yml b/changelogs/unreleased/54117-transactional-rebase.yml
new file mode 100644
index 00000000000..d0c93114c49
--- /dev/null
+++ b/changelogs/unreleased/54117-transactional-rebase.yml
@@ -0,0 +1,5 @@
+---
+title: Allow asynchronous rebase operations to be monitored
+merge_request: 29940
+author:
+type: fixed
diff --git a/db/migrate/20190621151636_add_merge_request_rebase_jid.rb b/db/migrate/20190621151636_add_merge_request_rebase_jid.rb
new file mode 100644
index 00000000000..1fed5690ead
--- /dev/null
+++ b/db/migrate/20190621151636_add_merge_request_rebase_jid.rb
@@ -0,0 +1,9 @@
+# frozen_string_literal: true
+
+class AddMergeRequestRebaseJid < ActiveRecord::Migration[5.1]
+  DOWNTIME = false
+
+  def change
+    add_column :merge_requests, :rebase_jid, :string
+  end
+end
diff --git a/db/schema.rb b/db/schema.rb
index 4bcc8b5f1d7..9cc45bb1e47 100644
--- a/db/schema.rb
+++ b/db/schema.rb
@@ -1989,6 +1989,7 @@ ActiveRecord::Schema.define(version: 20190628185004) do
     t.boolean "allow_maintainer_to_push"
     t.integer "state_id", limit: 2
     t.integer "approvals_before_merge"
+    t.string "rebase_jid"
     t.index ["assignee_id"], name: "index_merge_requests_on_assignee_id", using: :btree
     t.index ["author_id"], name: "index_merge_requests_on_author_id", using: :btree
     t.index ["created_at"], name: "index_merge_requests_on_created_at", using: :btree
diff --git a/doc/api/merge_requests.md b/doc/api/merge_requests.md
index 69db9f97a35..5a20db4b647 100644
--- a/doc/api/merge_requests.md
+++ b/doc/api/merge_requests.md
@@ -1485,8 +1485,14 @@ PUT /projects/:id/merge_requests/:merge_request_iid/rebase
 curl --header "PRIVATE-TOKEN: <your_access_token>" https://gitlab.example.com/api/v4/projects/76/merge_requests/1/rebase
 ```
 
-This is an asynchronous request. The API will return an empty `202 Accepted`
-response if the request is enqueued successfully.
+This is an asynchronous request. The API will return a `202 Accepted` response
+if the request is enqueued successfully, with a response containing:
+
+```json
+{
+  "rebase_in_progress": true
+}
+```
 
 You can poll the [Get single MR](#get-single-mr) endpoint with the
 `include_rebase_in_progress` parameter to check the status of the
diff --git a/lib/api/api.rb b/lib/api/api.rb
index 20f8c637274..42499c5b41e 100644
--- a/lib/api/api.rb
+++ b/lib/api/api.rb
@@ -52,7 +52,10 @@ module API
       rack_response({ 'message' => '404 Not found' }.to_json, 404)
     end
 
-    rescue_from ::Gitlab::ExclusiveLeaseHelpers::FailedToObtainLockError do
+    rescue_from(
+      ::ActiveRecord::StaleObjectError,
+      ::Gitlab::ExclusiveLeaseHelpers::FailedToObtainLockError
+    ) do
       rack_response({ 'message' => '409 Conflict: Resource lock' }.to_json, 409)
     end
 
diff --git a/lib/api/merge_requests.rb b/lib/api/merge_requests.rb
index 6b8c1a2c0e8..64ee82cd775 100644
--- a/lib/api/merge_requests.rb
+++ b/lib/api/merge_requests.rb
@@ -429,9 +429,10 @@ module API
 
         authorize_push_to_merge_request!(merge_request)
 
-        RebaseWorker.perform_async(merge_request.id, current_user.id)
+        merge_request.rebase_async(current_user.id)
 
         status :accepted
+        present rebase_in_progress: merge_request.rebase_in_progress?
       end
 
       desc 'List issues that will be closed on merge' do
diff --git a/lib/gitlab/import_export/import_export.yml b/lib/gitlab/import_export/import_export.yml
index a0fb051e806..01437c67fa9 100644
--- a/lib/gitlab/import_export/import_export.yml
+++ b/lib/gitlab/import_export/import_export.yml
@@ -160,6 +160,7 @@ excluded_attributes:
     - :milestone_id
     - :ref_fetched
     - :merge_jid
+    - :rebase_jid
     - :latest_merge_request_diff_id
   award_emoji:
     - :awardable_id
diff --git a/spec/models/merge_request_spec.rb b/spec/models/merge_request_spec.rb
index a2547755510..fe6d68aff3f 100644
--- a/spec/models/merge_request_spec.rb
+++ b/spec/models/merge_request_spec.rb
@@ -7,6 +7,8 @@ describe MergeRequest do
   include ProjectForksHelper
   include ReactiveCachingHelpers
 
+  using RSpec::Parameterized::TableSyntax
+
   subject { create(:merge_request) }
 
   describe 'associations' do
@@ -1996,6 +1998,47 @@ describe MergeRequest do
     end
   end
 
+  describe '#rebase_async' do
+    let(:merge_request) { create(:merge_request) }
+    let(:user_id) { double(:user_id) }
+    let(:rebase_jid) { 'rebase-jid' }
+
+    subject(:execute) { merge_request.rebase_async(user_id) }
+
+    it 'atomically enqueues a RebaseWorker job and updates rebase_jid' do
+      expect(RebaseWorker)
+        .to receive(:perform_async)
+        .with(merge_request.id, user_id)
+        .and_return(rebase_jid)
+
+      expect(merge_request).to receive(:lock!).and_call_original
+
+      execute
+
+      expect(merge_request.rebase_jid).to eq(rebase_jid)
+    end
+
+    it 'refuses to enqueue a job if a rebase is in progress' do
+      merge_request.update_column(:rebase_jid, rebase_jid)
+
+      expect(RebaseWorker).not_to receive(:perform_async)
+      expect(Gitlab::SidekiqStatus)
+        .to receive(:running?)
+        .with(rebase_jid)
+        .and_return(true)
+
+      expect { execute }.to raise_error(ActiveRecord::StaleObjectError)
+    end
+
+    it 'refuses to enqueue a job if the MR is not open' do
+      merge_request.update_column(:state, 'foo')
+
+      expect(RebaseWorker).not_to receive(:perform_async)
+
+      expect { execute }.to raise_error(ActiveRecord::StaleObjectError)
+    end
+  end
+
   describe '#mergeable?' do
     let(:project) { create(:project) }
 
@@ -2946,40 +2989,64 @@ describe MergeRequest do
   end
 
   describe '#rebase_in_progress?' do
-    shared_examples 'checking whether a rebase is in progress' do
-      let(:repo_path) do
-        Gitlab::GitalyClient::StorageSettings.allow_disk_access do
-          subject.source_project.repository.path
-        end
-      end
-      let(:rebase_path) { File.join(repo_path, "gitlab-worktree", "rebase-#{subject.id}") }
+    where(:rebase_jid, :jid_valid, :result) do
+      'foo' | true  | true
+      'foo' | false | false
+      ''    | true  | false
+      nil   | true  | false
+    end
 
-      before do
-        system(*%W(#{Gitlab.config.git.bin_path} -C #{repo_path} worktree add --detach #{rebase_path} master))
+    with_them do
+      let(:merge_request) { create(:merge_request) }
+
+      subject { merge_request.rebase_in_progress? }
+
+      it do
+        # Stub out the legacy gitaly implementation
+        allow(merge_request).to receive(:gitaly_rebase_in_progress?) { false }
+
+        allow(Gitlab::SidekiqStatus).to receive(:running?).with(rebase_jid) { jid_valid }
+
+        merge_request.rebase_jid = rebase_jid
+
+        is_expected.to eq(result)
       end
+    end
+  end
 
-      it 'returns true when there is a current rebase directory' do
-        expect(subject.rebase_in_progress?).to be_truthy
+  describe '#gitaly_rebase_in_progress?' do
+    let(:repo_path) do
+      Gitlab::GitalyClient::StorageSettings.allow_disk_access do
+        subject.source_project.repository.path
       end
+    end
+    let(:rebase_path) { File.join(repo_path, "gitlab-worktree", "rebase-#{subject.id}") }
+
+    before do
+      system(*%W(#{Gitlab.config.git.bin_path} -C #{repo_path} worktree add --detach #{rebase_path} master))
+    end
 
-      it 'returns false when there is no rebase directory' do
-        FileUtils.rm_rf(rebase_path)
+    it 'returns true when there is a current rebase directory' do
+      expect(subject.rebase_in_progress?).to be_truthy
+    end
 
-        expect(subject.rebase_in_progress?).to be_falsey
-      end
+    it 'returns false when there is no rebase directory' do
+      FileUtils.rm_rf(rebase_path)
 
-      it 'returns false when the rebase directory has expired' do
-        time = 20.minutes.ago.to_time
-        File.utime(time, time, rebase_path)
+      expect(subject.rebase_in_progress?).to be_falsey
+    end
 
-        expect(subject.rebase_in_progress?).to be_falsey
-      end
+    it 'returns false when the rebase directory has expired' do
+      time = 20.minutes.ago.to_time
+      File.utime(time, time, rebase_path)
 
-      it 'returns false when the source project has been removed' do
-        allow(subject).to receive(:source_project).and_return(nil)
+      expect(subject.rebase_in_progress?).to be_falsey
+    end
 
-        expect(subject.rebase_in_progress?).to be_falsey
-      end
+    it 'returns false when the source project has been removed' do
+      allow(subject).to receive(:source_project).and_return(nil)
+
+      expect(subject.rebase_in_progress?).to be_falsey
     end
   end
 
diff --git a/spec/requests/api/merge_requests_spec.rb b/spec/requests/api/merge_requests_spec.rb
index a82ecb4fd63..ced853caab4 100644
--- a/spec/requests/api/merge_requests_spec.rb
+++ b/spec/requests/api/merge_requests_spec.rb
@@ -2033,6 +2033,9 @@ describe API::MergeRequests do
 
       expect(response).to have_gitlab_http_status(202)
       expect(RebaseWorker.jobs.size).to eq(1)
+
+      expect(merge_request.reload).to be_rebase_in_progress
+      expect(json_response['rebase_in_progress']).to be(true)
     end
 
     it 'returns 403 if the user cannot push to the branch' do
@@ -2043,6 +2046,16 @@ describe API::MergeRequests do
 
       expect(response).to have_gitlab_http_status(403)
     end
+
+    it 'returns 409 if a rebase is already in progress' do
+      Sidekiq::Testing.fake! do
+        merge_request.rebase_async(user.id)
+
+        put api("/projects/#{project.id}/merge_requests/#{merge_request.iid}/rebase", user)
+      end
+
+      expect(response).to have_gitlab_http_status(409)
+    end
   end
 
   describe 'Time tracking' do
diff --git a/spec/services/merge_requests/rebase_service_spec.rb b/spec/services/merge_requests/rebase_service_spec.rb
index 7e2f03d1097..ee9caaf2f47 100644
--- a/spec/services/merge_requests/rebase_service_spec.rb
+++ b/spec/services/merge_requests/rebase_service_spec.rb
@@ -6,10 +6,12 @@ describe MergeRequests::RebaseService do
   include ProjectForksHelper
 
   let(:user) { create(:user) }
+  let(:rebase_jid) { 'fake-rebase-jid' }
   let(:merge_request) do
-    create(:merge_request,
+    create :merge_request,
            source_branch: 'feature_conflict',
-           target_branch: 'master')
+           target_branch: 'master',
+           rebase_jid: rebase_jid
   end
   let(:project) { merge_request.project }
   let(:repository) { project.repository.raw }
@@ -23,11 +25,11 @@ describe MergeRequests::RebaseService do
   describe '#execute' do
     context 'when another rebase is already in progress' do
       before do
-        allow(merge_request).to receive(:rebase_in_progress?).and_return(true)
+        allow(merge_request).to receive(:gitaly_rebase_in_progress?).and_return(true)
       end
 
       it 'saves the error message' do
-        subject.execute(merge_request)
+        service.execute(merge_request)
 
         expect(merge_request.reload.merge_error).to eq 'Rebase task canceled: Another rebase is already in progress'
       end
@@ -36,6 +38,13 @@ describe MergeRequests::RebaseService do
         expect(service.execute(merge_request)).to match(status: :error,
                                                         message: described_class::REBASE_ERROR)
       end
+
+      it 'clears rebase_jid' do
+        expect { service.execute(merge_request) }
+          .to change { merge_request.rebase_jid }
+          .from(rebase_jid)
+          .to(nil)
+      end
     end
 
     shared_examples 'sequence of failure and success' do
@@ -43,14 +52,19 @@ describe MergeRequests::RebaseService do
         allow(repository).to receive(:gitaly_operation_client).and_raise('Something went wrong')
 
         service.execute(merge_request)
+        merge_request.reload
 
-        expect(merge_request.reload.merge_error).to eq described_class::REBASE_ERROR
+        expect(merge_request.reload.merge_error).to eq(described_class::REBASE_ERROR)
+        expect(merge_request.rebase_jid).to eq(nil)
 
         allow(repository).to receive(:gitaly_operation_client).and_call_original
+        merge_request.update!(rebase_jid: rebase_jid)
 
         service.execute(merge_request)
+        merge_request.reload
 
-        expect(merge_request.reload.merge_error).to eq nil
+        expect(merge_request.merge_error).to eq(nil)
+        expect(merge_request.rebase_jid).to eq(nil)
       end
     end
 
@@ -72,7 +86,7 @@ describe MergeRequests::RebaseService do
       it 'saves a generic error message' do
         subject.execute(merge_request)
 
-        expect(merge_request.reload.merge_error).to eq described_class::REBASE_ERROR
+        expect(merge_request.reload.merge_error).to eq(described_class::REBASE_ERROR)
       end
 
       it 'returns an error' do