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

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'
 

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)

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,
+    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)
+        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)

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

changelogs/unreleased/5966-rebase-with-block.yml

+---
+title: Fix approvals sometimes being reset after a merge request is rebased
+merge_request: 27446
+author:
+type: fixed

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)

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
 

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') }

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