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

2 files changed, 74 insertions, 1 deletions

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