Add support for two-step Gitaly Rebase RPC

The new two-step Gitaly `Rebase` RPC yields the rebase commit SHA to the
client before proceeding with the rebase.

This avoids an issue where the rebase commit SHA was returned when the
RPC had fully completed, and in some cases this would be after the Rails
`post_receive` worker services had already run. In these situations,
the merge request did not yet have its rebase_commit_sha attribute set
introducing the possibility for bugs (such as previous approvals being
reset).

https://gitlab.com/gitlab-org/gitlab-ee/issues/5966

9 files changed, 219 insertions, 20 deletions

diff --git a/Gemfile b/Gemfile
index f3c21c720e0..ae9a29370c7 100644
--- a/Gemfile
+++ b/Gemfile
@@ -417,7 +417,7 @@ group :ed25519 do
 end
 
 # Gitaly GRPC client
-gem 'gitaly-proto', '~> 1.26.0', require: 'gitaly'
+gem 'gitaly-proto', '~> 1.27.0', require: 'gitaly'
 
 gem 'grpc', '~> 1.19.0'
 
diff --git a/Gemfile.lock b/Gemfile.lock
index 4db00ba4e18..053dd322b01 100644
--- a/Gemfile.lock
+++ b/Gemfile.lock
@@ -283,7 +283,7 @@ GEM
       gettext_i18n_rails (>= 0.7.1)
       po_to_json (>= 1.0.0)
       rails (>= 3.2.0)
-    gitaly-proto (1.26.0)
+    gitaly-proto (1.27.0)
       grpc (~> 1.0)
     github-markup (1.7.0)
     gitlab-default_value_for (3.1.1)
@@ -1056,7 +1056,7 @@ DEPENDENCIES
   gettext (~> 3.2.2)
   gettext_i18n_rails (~> 1.8.0)
   gettext_i18n_rails_js (~> 1.3)
-  gitaly-proto (~> 1.26.0)
+  gitaly-proto (~> 1.27.0)
   github-markup (~> 1.7.0)
   gitlab-default_value_for (~> 3.1.1)
   gitlab-labkit (~> 0.2.0)
diff --git a/app/models/repository.rb b/app/models/repository.rb
index f495a03ad8e..be17b54ff12 100644
--- a/app/models/repository.rb
+++ b/app/models/repository.rb
@@ -1037,11 +1037,41 @@ class Repository
     raw_repository.fetch_ref(source_repository.raw_repository, source_ref: source_ref, target_ref: target_ref)
   end
 
+  # DEPRECATED: https://gitlab.com/gitlab-org/gitaly/issues/1628
+  def rebase_deprecated(user, merge_request)
+    rebase_sha = raw.rebase_deprecated(
+      user,
+      merge_request.id,
+      branch: merge_request.source_branch,
+      branch_sha: merge_request.source_branch_sha,
+      remote_repository: merge_request.target_project.repository.raw,
+      remote_branch: merge_request.target_branch
+    )
+
+    # To support the full deprecated behaviour, set the
+    # `rebase_commit_sha` for the merge_request here and return the value
+    merge_request.update(rebase_commit_sha: rebase_sha)
+
+    rebase_sha
+  end
+
   def rebase(user, merge_request)
-    raw.rebase(user, merge_request.id, branch: merge_request.source_branch,
-                                       branch_sha: merge_request.source_branch_sha,
-                                       remote_repository: merge_request.target_project.repository.raw,
-                                       remote_branch: merge_request.target_branch)
+    if Feature.disabled?(:two_step_rebase, default_enabled: true)
+      return rebase_deprecated(user, merge_request)
+    end
+
+    MergeRequest.transaction do
+      raw.rebase(
+        user,
+        merge_request.id,
+        branch: merge_request.source_branch,
+        branch_sha: merge_request.source_branch_sha,
+        remote_repository: merge_request.target_project.repository.raw,
+        remote_branch: merge_request.target_branch
+      ) do |commit_id|
+        merge_request.update!(rebase_commit_sha: commit_id)
+      end
+    end
   end
 
   def squash(user, merge_request, message)
diff --git a/app/services/merge_requests/rebase_service.rb b/app/services/merge_requests/rebase_service.rb
index 31b3ebf311e..4b9921c28ba 100644
--- a/app/services/merge_requests/rebase_service.rb
+++ b/app/services/merge_requests/rebase_service.rb
@@ -20,17 +20,7 @@ module MergeRequests
         return false
       end
 
-      log_prefix = "#{self.class.name} info (#{merge_request.to_reference(full: true)}):"
-
-      Gitlab::GitLogger.info("#{log_prefix} rebase started")
-
-      rebase_sha = repository.rebase(current_user, merge_request)
-
-      Gitlab::GitLogger.info("#{log_prefix} rebased to #{rebase_sha}")
-
-      merge_request.update(rebase_commit_sha: rebase_sha)
-
-      Gitlab::GitLogger.info("#{log_prefix} rebase SHA saved: #{rebase_sha}")
+      repository.rebase(current_user, merge_request)
 
       true
     rescue => e
diff --git a/changelogs/unreleased/5966-rebase-with-block.yml b/changelogs/unreleased/5966-rebase-with-block.yml
new file mode 100644
index 00000000000..9272a02977f
--- /dev/null
+++ b/changelogs/unreleased/5966-rebase-with-block.yml
@@ -0,0 +1,5 @@
+---
+title: Fix approvals sometimes being reset after a merge request is rebased
+merge_request: 27446
+author:
+type: fixed
diff --git a/lib/gitlab/git/repository.rb b/lib/gitlab/git/repository.rb
index 55bd77f6c4a..508499f227c 100644
--- a/lib/gitlab/git/repository.rb
+++ b/lib/gitlab/git/repository.rb
@@ -834,7 +834,8 @@ module Gitlab
         gitaly_repository_client.create_from_snapshot(url, auth)
       end
 
-      def rebase(user, rebase_id, branch:, branch_sha:, remote_repository:, remote_branch:)
+      # DEPRECATED: https://gitlab.com/gitlab-org/gitaly/issues/1628
+      def rebase_deprecated(user, rebase_id, branch:, branch_sha:, remote_repository:, remote_branch:)
         wrapped_gitaly_errors do
           gitaly_operation_client.user_rebase(user, rebase_id,
                                             branch: branch,
@@ -844,6 +845,20 @@ module Gitlab
         end
       end
 
+      def rebase(user, rebase_id, branch:, branch_sha:, remote_repository:, remote_branch:, &block)
+        wrapped_gitaly_errors do
+          gitaly_operation_client.rebase(
+            user,
+            rebase_id,
+            branch: branch,
+            branch_sha: branch_sha,
+            remote_repository: remote_repository,
+            remote_branch: remote_branch,
+            &block
+          )
+        end
+      end
+
       def rebase_in_progress?(rebase_id)
         wrapped_gitaly_errors do
           gitaly_repository_client.rebase_in_progress?(rebase_id)
diff --git a/lib/gitlab/gitaly_client/operation_service.rb b/lib/gitlab/gitaly_client/operation_service.rb
index b0f328ce3d4..e4a59ee3f9b 100644
--- a/lib/gitlab/gitaly_client/operation_service.rb
+++ b/lib/gitlab/gitaly_client/operation_service.rb
@@ -197,6 +197,7 @@ module Gitlab
                                    start_repository: start_repository)
       end
 
+      # DEPRECATED: https://gitlab.com/gitlab-org/gitaly/issues/1628
       def user_rebase(user, rebase_id, branch:, branch_sha:, remote_repository:, remote_branch:)
         request = Gitaly::UserRebaseRequest.new(
           repository: @gitaly_repo,
@@ -225,6 +226,49 @@ module Gitlab
         end
       end
 
+      def rebase(user, rebase_id, branch:, branch_sha:, remote_repository:, remote_branch:)
+        request_enum = QueueEnumerator.new
+        rebase_sha = nil
+
+        response_enum = GitalyClient.call(
+          @repository.storage,
+          :operation_service,
+          :user_rebase_confirmable,
+          request_enum.each,
+          remote_storage: remote_repository.storage
+        )
+
+        # First request
+        request_enum.push(
+          Gitaly::UserRebaseConfirmableRequest.new(
+            header: Gitaly::UserRebaseConfirmableRequest::Header.new(
+              repository: @gitaly_repo,
+              user: Gitlab::Git::User.from_gitlab(user).to_gitaly,
+              rebase_id: rebase_id.to_s,
+              branch: encode_binary(branch),
+              branch_sha: branch_sha,
+              remote_repository: remote_repository.gitaly_repository,
+              remote_branch: encode_binary(remote_branch)
+            )
+          )
+        )
+
+        perform_next_gitaly_rebase_request(response_enum) do |response|
+          rebase_sha = response.rebase_sha
+        end
+
+        yield rebase_sha
+
+        # Second request confirms with gitaly to finalize the rebase
+        request_enum.push(Gitaly::UserRebaseConfirmableRequest.new(apply: true))
+
+        perform_next_gitaly_rebase_request(response_enum)
+
+        rebase_sha
+      ensure
+        request_enum.close
+      end
+
       def user_squash(user, squash_id, branch, start_sha, end_sha, author, message)
         request = Gitaly::UserSquashRequest.new(
           repository: @gitaly_repo,
@@ -346,6 +390,20 @@ module Gitlab
 
       private
 
+      def perform_next_gitaly_rebase_request(response_enum)
+        response = response_enum.next
+
+        if response.pre_receive_error.present?
+          raise Gitlab::Git::PreReceiveError, response.pre_receive_error
+        elsif response.git_error.present?
+          raise Gitlab::Git::Repository::GitError, response.git_error
+        end
+
+        yield response if block_given?
+
+        response
+      end
+
       def call_cherry_pick_or_revert(rpc, user:, commit:, branch_name:, message:, start_branch_name:, start_repository:)
         request_class = "Gitaly::User#{rpc.to_s.camelcase}Request".constantize
 
diff --git a/spec/models/repository_spec.rb b/spec/models/repository_spec.rb
index a6c3d5756aa..9ff0f355fd4 100644
--- a/spec/models/repository_spec.rb
+++ b/spec/models/repository_spec.rb
@@ -1451,6 +1451,91 @@ describe Repository do
     end
   end
 
+  describe '#rebase' do
+    let(:merge_request) { create(:merge_request, source_branch: 'feature', target_branch: 'master', source_project: project) }
+
+    shared_examples_for 'a method that can rebase successfully' do
+      it 'returns the rebase commit sha' do
+        rebase_commit_sha = repository.rebase(user, merge_request)
+        head_sha = merge_request.source_project.repository.commit(merge_request.source_branch).sha
+
+        expect(rebase_commit_sha).to eq(head_sha)
+      end
+
+      it 'sets the `rebase_commit_sha` for the given merge request' do
+        rebase_commit_sha = repository.rebase(user, merge_request)
+
+        expect(rebase_commit_sha).not_to be_nil
+        expect(merge_request.rebase_commit_sha).to eq(rebase_commit_sha)
+      end
+    end
+
+    context 'when two_step_rebase feature is enabled' do
+      before do
+        stub_feature_flags(two_step_rebase: true)
+      end
+
+      it_behaves_like 'a method that can rebase successfully'
+
+      it 'executes the new Gitaly RPC' do
+        expect_any_instance_of(Gitlab::GitalyClient::OperationService).to receive(:rebase)
+        expect_any_instance_of(Gitlab::GitalyClient::OperationService).not_to receive(:user_rebase)
+
+        repository.rebase(user, merge_request)
+      end
+
+      describe 'rolling back the `rebase_commit_sha`' do
+        let(:new_sha) { Digest::SHA1.hexdigest('foo') }
+
+        it 'does not rollback when there are no errors' do
+          second_response = double(pre_receive_error: nil, git_error: nil)
+          mock_gitaly(second_response)
+
+          repository.rebase(user, merge_request)
+
+          expect(merge_request.reload.rebase_commit_sha).to eq(new_sha)
+        end
+
+        it 'does rollback when an error is encountered in the second step' do
+          second_response = double(pre_receive_error: 'my_error', git_error: nil)
+          mock_gitaly(second_response)
+
+          expect do
+            repository.rebase(user, merge_request)
+          end.to raise_error(Gitlab::Git::PreReceiveError)
+
+          expect(merge_request.reload.rebase_commit_sha).to be_nil
+        end
+
+        def mock_gitaly(second_response)
+          responses = [
+            double(rebase_sha: new_sha).as_null_object,
+            second_response
+          ]
+
+          expect_any_instance_of(
+            Gitaly::OperationService::Stub
+          ).to receive(:user_rebase_confirmable).and_return(responses.each)
+        end
+      end
+    end
+
+    context 'when two_step_rebase feature is disabled' do
+      before do
+        stub_feature_flags(two_step_rebase: false)
+      end
+
+      it_behaves_like 'a method that can rebase successfully'
+
+      it 'executes the deprecated Gitaly RPC' do
+        expect_any_instance_of(Gitlab::GitalyClient::OperationService).to receive(:user_rebase)
+        expect_any_instance_of(Gitlab::GitalyClient::OperationService).not_to receive(:rebase)
+
+        repository.rebase(user, merge_request)
+      end
+    end
+  end
+
   describe '#revert' do
     let(:new_image_commit) { repository.commit('33f3729a45c02fc67d00adb1b8bca394b0e761d9') }
     let(:update_image_commit) { repository.commit('2f63565e7aac07bcdadb654e253078b727143ec4') }
diff --git a/spec/services/merge_requests/rebase_service_spec.rb b/spec/services/merge_requests/rebase_service_spec.rb
index 0d9523a4c2c..a443e4588d9 100644
--- a/spec/services/merge_requests/rebase_service_spec.rb
+++ b/spec/services/merge_requests/rebase_service_spec.rb
@@ -73,7 +73,7 @@ describe MergeRequests::RebaseService do
     end
 
     context 'valid params' do
-      describe 'successful rebase' do
+      shared_examples_for 'a service that can execute a successful rebase' do
         before do
           service.execute(merge_request)
         end
@@ -99,6 +99,22 @@ describe MergeRequests::RebaseService do
         end
       end
 
+      context 'when the two_step_rebase feature is enabled' do
+        before do
+          stub_feature_flags(two_step_rebase: true)
+        end
+
+        it_behaves_like 'a service that can execute a successful rebase'
+      end
+
+      context 'when the two_step_rebase feature is disabled' do
+        before do
+          stub_feature_flags(two_step_rebase: false)
+        end
+
+        it_behaves_like 'a service that can execute a successful rebase'
+      end
+
       context 'fork' do
         describe 'successful fork rebase' do
           let(:forked_project) do