Backport Squash/Rebase refactor from EEce-gitaly-worktree-refactor

1 files changed, 140 insertions, 8 deletions

diff --git a/lib/gitlab/git/repository.rb b/lib/gitlab/git/repository.rb
index 867fc2a42f6..d50de2fb6c1 100644
--- a/lib/gitlab/git/repository.rb
+++ b/lib/gitlab/git/repository.rb
@@ -18,6 +18,8 @@ module Gitlab
         GIT_ALTERNATE_OBJECT_DIRECTORIES_RELATIVE
       ].freeze
       SEARCH_CONTEXT_LINES = 3
+      REBASE_WORKTREE_PREFIX = 'rebase'.freeze
+      SQUASH_WORKTREE_PREFIX = 'squash'.freeze
 
       NoRepository = Class.new(StandardError)
       InvalidBlobName = Class.new(StandardError)
@@ -1090,13 +1092,8 @@ module Gitlab
         raise ArgumentError, "invalid ref_path #{ref_path.inspect}" if ref_path.include?(' ')
         raise ArgumentError, "invalid ref #{ref.inspect}" if ref.include?("\x00")
 
-        command = [Gitlab.config.git.bin_path] + %w[update-ref --stdin -z]
         input = "update #{ref_path}\x00#{ref}\x00\x00"
-        output, status = circuit_breaker.perform do
-          popen(command, path) { |stdin| stdin.write(input) }
-        end
-
-        raise GitError, output unless status.zero?
+        run_git!(%w[update-ref --stdin -z]) { |stdin| stdin.write(input) }
       end
 
       def fetch_ref(source_repository, source_ref:, target_ref:)
@@ -1118,14 +1115,22 @@ module Gitlab
       end
 
       # Refactoring aid; allows us to copy code from app/models/repository.rb
-      def run_git(args, env: {}, nice: false)
+      def run_git(args, chdir: path, env: {}, nice: false, &block)
         cmd = [Gitlab.config.git.bin_path, *args]
         cmd.unshift("nice") if nice
         circuit_breaker.perform do
-          popen(cmd, path, env)
+          popen(cmd, chdir, env, &block)
         end
       end
 
+      def run_git!(args, chdir: path, env: {}, nice: false, &block)
+        output, status = run_git(args, chdir: chdir, env: env, nice: nice, &block)
+
+        raise GitError, output unless status.zero?
+
+        output
+      end
+
       # Refactoring aid; allows us to copy code from app/models/repository.rb
       def run_git_with_timeout(args, timeout, env: {})
         circuit_breaker.perform do
@@ -1195,6 +1200,64 @@ module Gitlab
         raise GitError.new("Could not fsck repository:\n#{output}") unless status.zero?
       end
 
+      def rebase(user, rebase_id, branch:, branch_sha:, remote_repository:, remote_branch:)
+        rebase_path = worktree_path(REBASE_WORKTREE_PREFIX, rebase_id)
+        env = git_env_for_user(user)
+
+        with_worktree(rebase_path, branch, env: env) do
+          run_git!(
+            %W(pull --rebase #{remote_repository.path} #{remote_branch}),
+            chdir: rebase_path, env: env
+          )
+
+          rebase_sha = run_git!(%w(rev-parse HEAD), chdir: rebase_path, env: env).strip
+
+          Gitlab::Git::OperationService.new(user, self)
+            .update_branch(branch, rebase_sha, branch_sha)
+
+          rebase_sha
+        end
+      end
+
+      def rebase_in_progress?(rebase_id)
+        fresh_worktree?(worktree_path(REBASE_WORKTREE_PREFIX, rebase_id))
+      end
+
+      def squash(user, squash_id, branch:, start_sha:, end_sha:, author:, message:)
+        squash_path = worktree_path(SQUASH_WORKTREE_PREFIX, squash_id)
+        env = git_env_for_user(user).merge(
+          'GIT_AUTHOR_NAME' => author.name,
+          'GIT_AUTHOR_EMAIL' => author.email
+        )
+        diff_range = "#{start_sha}...#{end_sha}"
+        diff_files = run_git!(
+          %W(diff --name-only --diff-filter=a --binary #{diff_range})
+        ).chomp
+
+        with_worktree(squash_path, branch, sparse_checkout_files: diff_files, env: env) do
+          # Apply diff of the `diff_range` to the worktree
+          diff = run_git!(%W(diff --binary #{diff_range}))
+          run_git!(%w(apply --index), chdir: squash_path, env: env) do |stdin|
+            stdin.write(diff)
+          end
+
+          # Commit the `diff_range` diff
+          run_git!(%W(commit --no-verify --message #{message}), chdir: squash_path, env: env)
+
+          # Return the squash sha. May print a warning for ambiguous refs, but
+          # we can ignore that with `--quiet` and just take the SHA, if present.
+          # HEAD here always refers to the current HEAD commit, even if there is
+          # another ref called HEAD.
+          run_git!(
+            %w(rev-parse --quiet --verify HEAD), chdir: squash_path, env: env
+          ).chomp
+        end
+      end
+
+      def squash_in_progress?(squash_id)
+        fresh_worktree?(worktree_path(SQUASH_WORKTREE_PREFIX, squash_id))
+      end
+
       def gitaly_repository
         Gitlab::GitalyClient::Util.repository(@storage, @relative_path, @gl_repository)
       end
@@ -1231,6 +1294,57 @@ module Gitlab
 
       private
 
+      def fresh_worktree?(path)
+        File.exist?(path) && !clean_stuck_worktree(path)
+      end
+
+      def with_worktree(worktree_path, branch, sparse_checkout_files: nil, env:)
+        base_args = %w(worktree add --detach)
+
+        # Note that we _don't_ want to test for `.present?` here: If the caller
+        # passes an non nil empty value it means it still wants sparse checkout
+        # but just isn't interested in any file, perhaps because it wants to
+        # checkout files in by a changeset but that changeset only adds files.
+        if sparse_checkout_files
+          # Create worktree without checking out
+          run_git!(base_args + ['--no-checkout', worktree_path], env: env)
+          worktree_git_path = run_git!(%w(rev-parse --git-dir), chdir: worktree_path)
+
+          configure_sparse_checkout(worktree_git_path, sparse_checkout_files)
+
+          # After sparse checkout configuration, checkout `branch` in worktree
+          run_git!(%W(checkout --detach #{branch}), chdir: worktree_path, env: env)
+        else
+          # Create worktree and checkout `branch` in it
+          run_git!(base_args + [worktree_path, branch], env: env)
+        end
+
+        yield
+      ensure
+        FileUtils.rm_rf(worktree_path) if File.exist?(worktree_path)
+        FileUtils.rm_rf(worktree_git_path) if worktree_git_path && File.exist?(worktree_git_path)
+      end
+
+      def clean_stuck_worktree(path)
+        return false unless File.mtime(path) < 15.minutes.ago
+
+        FileUtils.rm_rf(path)
+        true
+      end
+
+      # Adding a worktree means checking out the repository. For large repos,
+      # this can be very expensive, so set up sparse checkout for the worktree
+      # to only check out the files we're interested in.
+      def configure_sparse_checkout(worktree_git_path, files)
+        run_git!(%w(config core.sparseCheckout true))
+
+        return if files.empty?
+
+        worktree_info_path = File.join(worktree_git_path, 'info')
+        FileUtils.mkdir_p(worktree_info_path)
+        File.write(File.join(worktree_info_path, 'sparse-checkout'), files)
+      end
+
       def rugged_fetch_source_branch(source_repository, source_branch, local_ref)
         with_repo_branch_commit(source_repository, source_branch) do |commit|
           if commit
@@ -1242,6 +1356,24 @@ module Gitlab
         end
       end
 
+      def worktree_path(prefix, id)
+        id = id.to_s
+        raise ArgumentError, "worktree id can't be empty" unless id.present?
+        raise ArgumentError, "worktree id can't contain slashes " if id.include?("/")
+
+        File.join(path, 'gitlab-worktree', "#{prefix}-#{id}")
+      end
+
+      def git_env_for_user(user)
+        {
+          'GIT_COMMITTER_NAME' => user.name,
+          'GIT_COMMITTER_EMAIL' => user.email,
+          'GL_ID' => Gitlab::GlId.gl_id(user),
+          'GL_PROTOCOL' => Gitlab::Git::Hook::GL_PROTOCOL,
+          'GL_REPOSITORY' => gl_repository
+        }
+      end
+
       # Gitaly note: JV: Trying to get rid of the 'filter' option so we can implement this with 'git'.
       def branches_filter(filter: nil, sort_by: nil)
         # n+1: https://gitlab.com/gitlab-org/gitlab-ce/issues/37464