Merge branch 'sh-worktree-cleanup' into 'master'

Properly clean up worktrees after commit operations

Closes #1593

See merge request gitlab-org/gitaly!1383

6 files changed, 92 insertions, 33 deletions

diff --git a/changelogs/unreleased/sh-worktree-cleanup.yml b/changelogs/unreleased/sh-worktree-cleanup.yml
new file mode 100644
index 000000000..8b239dcc7
--- /dev/null
+++ b/changelogs/unreleased/sh-worktree-cleanup.yml
@@ -0,0 +1,5 @@
+---
+title: Properly clean up worktrees after commit operations
+merge_request: 1383
+author:
+type: fixed
diff --git a/internal/service/operations/squash_test.go b/internal/service/operations/squash_test.go
index 05bdb8bbe..18b6afd39 100644
--- a/internal/service/operations/squash_test.go
+++ b/internal/service/operations/squash_test.go
@@ -1,6 +1,7 @@
 package operations
 
 import (
+	"fmt"
 	"io/ioutil"
 	"os"
 	"path/filepath"
@@ -145,6 +146,16 @@ index 5b812ff..ff85394 100644
 	require.Equal(t, user.Name, commit.Committer.Name)
 	require.Equal(t, user.Email, commit.Committer.Email)
 	require.Equal(t, commitMessage, commit.Subject)
+
+	// Ensure Git metadata is cleaned up
+	worktreeList := text.ChompBytes(testhelper.MustRunCommand(t, nil, "git", "-C", testRepoPath, "worktree", "list", "--porcelain"))
+	expectedOut := fmt.Sprintf("worktree %s\nHEAD %s\nbranch refs/heads/3way-test\n", testRepoPath, baseSha)
+	require.Equal(t, expectedOut, worktreeList)
+
+	// Ensure actual worktree is removed
+	files, err := ioutil.ReadDir(filepath.Join(testRepoPath, ".git", "gitlab-worktree"))
+	require.NoError(t, err)
+	require.Equal(t, 0, len(files))
 }
 
 func TestSplitIndex(t *testing.T) {
diff --git a/ruby/lib/gitlab/git/repository.rb b/ruby/lib/gitlab/git/repository.rb
index ad14b35bb..1d9aee50d 100644
--- a/ruby/lib/gitlab/git/repository.rb
+++ b/ruby/lib/gitlab/git/repository.rb
@@ -350,7 +350,7 @@ module Gitlab
       end
 
       def rebase(user, rebase_id, branch:, branch_sha:, remote_repository:, remote_branch:)
-        rebase_path = worktree_path(REBASE_WORKTREE_PREFIX, rebase_id)
+        worktree = Gitlab::Git::Worktree.new(path, REBASE_WORKTREE_PREFIX, rebase_id)
         env = git_env.merge(user.git_env)
 
         if remote_repository.is_a?(RemoteRepository)
@@ -360,13 +360,13 @@ module Gitlab
           remote_repo_path = remote_repository.path
         end
 
-        with_worktree(rebase_path, branch, env: env) do
+        with_worktree(worktree, branch, env: env) do
           run_git!(
             %W[pull --rebase #{remote_repo_path} #{remote_branch}],
-            chdir: rebase_path, env: env, include_stderr: true
+            chdir: worktree.path, env: env, include_stderr: true
           )
 
-          rebase_sha = run_git!(%w[rev-parse HEAD], chdir: rebase_path, env: env).strip
+          rebase_sha = run_git!(%w[rev-parse HEAD], chdir: worktree.path, env: env).strip
 
           yield rebase_sha if block_given?
 
@@ -377,7 +377,7 @@ module Gitlab
       end
 
       def squash(user, squash_id, branch:, start_sha:, end_sha:, author:, message:)
-        squash_path = worktree_path(SQUASH_WORKTREE_PREFIX, squash_id)
+        worktree = Gitlab::Git::Worktree.new(path, SQUASH_WORKTREE_PREFIX, squash_id)
         env = git_env.merge(user.git_env).merge(
           'GIT_AUTHOR_NAME' => author.name,
           'GIT_AUTHOR_EMAIL' => author.email
@@ -387,40 +387,40 @@ module Gitlab
           %W[diff --name-only --diff-filter=ar --binary #{diff_range}]
         ).chomp
 
-        with_worktree(squash_path, branch, sparse_checkout_files: diff_files, env: env) do
+        with_worktree(worktree, 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 --3way --whitespace=nowarn], chdir: squash_path, env: env, include_stderr: true) do |stdin|
+          run_git!(%w[apply --index --3way --whitespace=nowarn], chdir: worktree.path, env: env, include_stderr: true) do |stdin|
             stdin.binmode
             stdin.write(diff)
           end
 
           # Commit the `diff_range` diff
-          run_git!(%W[commit --no-verify --message #{message}], chdir: squash_path, env: env, include_stderr: true)
+          run_git!(%W[commit --no-verify --message #{message}], chdir: worktree.path, env: env, include_stderr: true)
 
           # 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
+            %w[rev-parse --quiet --verify HEAD], chdir: worktree.path, env: env
           ).chomp
         end
       end
 
       def commit_patches(start_point, patches, extra_env: {})
-        worktree_path = worktree_path(AM_WORKTREE_PREFIX, SecureRandom.hex)
+        worktree = Gitlab::Git::Worktree.new(path, AM_WORKTREE_PREFIX, SecureRandom.hex)
         env = git_env.merge(extra_env)
 
-        with_worktree(worktree_path, start_point, env: env) do
-          result, status = run_git(%w[am --quiet --3way], chdir: worktree_path, env: env) do |stdin|
+        with_worktree(worktree, start_point, env: env) do
+          result, status = run_git(%w[am --quiet --3way], chdir: worktree.path, env: env) do |stdin|
             loop { stdin.write(patches.next) }
           end
 
           raise Gitlab::Git::PatchError, result unless status == 0
 
           run_git!(
-            %w[rev-parse --quiet --verify HEAD], chdir: worktree_path, env: env
+            %w[rev-parse --quiet --verify HEAD], chdir: worktree.path, env: env
           ).chomp
         end
       end
@@ -859,7 +859,7 @@ module Gitlab
         nil
       end
 
-      def with_worktree(worktree_path, branch, sparse_checkout_files: nil, env:)
+      def with_worktree(worktree, branch, sparse_checkout_files: nil, env:)
         base_args = %w[worktree add --detach]
 
         run_git!(%w[config core.splitIndex true])
@@ -870,22 +870,21 @@ module Gitlab
         # 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).chomp
+          run_git!(base_args + ['--no-checkout', worktree.path], env: env)
+          worktree_git_path = run_git!(%w[rev-parse --git-dir], chdir: worktree.path).chomp
 
           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)
+          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)
+          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)
+        run_git(%W[worktree remove -f #{worktree.name}])
       end
 
       # Adding a worktree means checking out the repository. For large repos,
@@ -901,14 +900,6 @@ module Gitlab
         File.write(File.join(worktree_info_path, 'sparse-checkout'), files)
       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 rugged_fetch_source_branch(source_repository, source_branch, local_ref)
         with_repo_branch_commit(source_repository, source_branch) do |commit|
           if commit
diff --git a/ruby/lib/gitlab/git/worktree.rb b/ruby/lib/gitlab/git/worktree.rb
new file mode 100644
index 000000000..59b62e5e3
--- /dev/null
+++ b/ruby/lib/gitlab/git/worktree.rb
@@ -0,0 +1,26 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Git
+    class Worktree
+      attr_reader :path, :name
+
+      def initialize(repo_path, prefix, id)
+        @repo_path = repo_path
+        @prefix = prefix
+        @id = id.to_s
+        @name = "#{prefix}-#{id}"
+        @path = worktree_path
+      end
+
+      private
+
+      def worktree_path
+        raise ArgumentError, "worktree id can't be empty" unless @id.present?
+        raise ArgumentError, "worktree id can't contain slashes " if @id.include?("/")
+
+        File.join(@repo_path, 'gitlab-worktree', @name)
+      end
+    end
+  end
+end
diff --git a/ruby/spec/lib/gitlab/git/repository_spec.rb b/ruby/spec/lib/gitlab/git/repository_spec.rb
index c7cedbdb7..595111814 100644
--- a/ruby/spec/lib/gitlab/git/repository_spec.rb
+++ b/ruby/spec/lib/gitlab/git/repository_spec.rb
@@ -661,10 +661,10 @@ describe Gitlab::Git::Repository do # rubocop:disable Metrics/BlockLength
       it 'checks out only the files in the diff' do
         allow(repository).to receive(:with_worktree).and_wrap_original do |m, *args|
           m.call(*args) do
-            worktree_path = args[0]
-            files_pattern = File.join(worktree_path, '**', '*')
+            worktree = args[0]
+            files_pattern = File.join(worktree.path, '**', '*')
             expected = expected_files.map do |path|
-              File.expand_path(path, worktree_path)
+              File.expand_path(path, worktree.path)
             end
 
             expect(Dir[files_pattern]).to eq(expected)
@@ -685,8 +685,8 @@ describe Gitlab::Git::Repository do # rubocop:disable Metrics/BlockLength
         it 'does not include the renamed file in the sparse checkout' do
           allow(repository).to receive(:with_worktree).and_wrap_original do |m, *args|
             m.call(*args) do
-              worktree_path = args[0]
-              files_pattern = File.join(worktree_path, '**', '*')
+              worktree = args[0]
+              files_pattern = File.join(worktree.path, '**', '*')
 
               expect(Dir[files_pattern]).not_to include('CHANGELOG')
               expect(Dir[files_pattern]).not_to include('encoding/CHANGELOG')
@@ -818,6 +818,11 @@ describe Gitlab::Git::Repository do # rubocop:disable Metrics/BlockLength
 
         expect(new_rev).not_to be_nil
         expect(commit.message).to eq("A commit from a patch\n")
+
+        # Ensure worktree cleanup occurs
+        result, status = repository.send(:run_git, %w[worktree list --porcelain])
+        expect(status).to eq(0)
+        expect(result).to eq("worktree #{repository_path}\nbare\n\n")
       end
     end
 
diff --git a/ruby/spec/lib/gitlab/git/worktree_spec.rb b/ruby/spec/lib/gitlab/git/worktree_spec.rb
new file mode 100644
index 000000000..ac4834d63
--- /dev/null
+++ b/ruby/spec/lib/gitlab/git/worktree_spec.rb
@@ -0,0 +1,21 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe Gitlab::Git::Worktree do
+  context '#initialize' do
+    let(:repo_path) { '/tmp/test' }
+    let(:prefix) { 'rebase' }
+
+    it 'generates valid path' do
+      worktree = described_class.new(repo_path, prefix, 12345)
+
+      expect(worktree.path).to eq('/tmp/test/gitlab-worktree/rebase-12345')
+    end
+
+    it 'rejects bad IDs' do
+      expect { described_class.new(repo_path, prefix, '') }.to raise_error(ArgumentError)
+      expect { described_class.new(repo_path, prefix, '/test/me') }.to raise_error(ArgumentError)
+    end
+  end
+end