localrepo: Modify WriteTree to be recursive, and to operate on TreeEntryjc/restore-tree-improvements-refactor

Since WriteTree() operates on trees, change it to Write() and make it a
function for TreeEntry. It will do a depth first walk of all the Entries,
writing new trees as necessary. If a TreeEntry's OID is empty, then
Write() will attempt to write a new tree with all of its Entries.

4 files changed, 295 insertions, 36 deletions

diff --git a/internal/git/localrepo/commit_test.go b/internal/git/localrepo/commit_test.go
index 9238ae7a8..9fe1a19dc 100644
--- a/internal/git/localrepo/commit_test.go
+++ b/internal/git/localrepo/commit_test.go
@@ -32,13 +32,23 @@ func TestWriteCommit(t *testing.T) {
 	require.NoError(t, err)
 
 	treeEntryA := TreeEntry{Path: "file", Mode: "100644", OID: blobID}
-	treeA, err := repo.WriteTree(ctx, []*TreeEntry{&treeEntryA})
-	require.NoError(t, err)
+	treeA := &TreeEntry{
+		Type:    Tree,
+		Mode:    "040000",
+		Entries: []*TreeEntry{&treeEntryA},
+	}
+	require.NoError(t, treeA.Write(
+		ctx,
+		repo))
 
-	treeB, err := repo.WriteTree(ctx, []*TreeEntry{
-		{Path: "file", Mode: "100644", OID: changedBlobID},
-	})
-	require.NoError(t, err)
+	treeB := &TreeEntry{
+		Type: Tree,
+		Mode: "040000",
+		Entries: []*TreeEntry{
+			{Path: "file", Mode: "100644", OID: changedBlobID},
+		},
+	}
+	require.NoError(t, treeB.Write(ctx, repo))
 	commitA, err := repo.WriteCommit(
 		ctx,
 		WriteCommitConfig{
@@ -47,7 +57,7 @@ func TestWriteCommit(t *testing.T) {
 			CommitterName:  "Tazmanian Devil",
 			CommitterEmail: "taz@devils.org",
 			Message:        "I ❤️  Tazmania",
-			TreeID:         treeA,
+			TreeID:         treeA.OID,
 		},
 	)
 	require.NoError(t, err)
@@ -59,7 +69,7 @@ func TestWriteCommit(t *testing.T) {
 			CommitterName:  "Daffy Duck",
 			CommitterEmail: "daffy@ducks.org",
 			Message:        "Big beak",
-			TreeID:         treeB,
+			TreeID:         treeB.OID,
 		},
 	)
 	require.NoError(t, err)
@@ -80,7 +90,7 @@ func TestWriteCommit(t *testing.T) {
 		{
 			desc: "with commit message",
 			cfg: WriteCommitConfig{
-				TreeID:         treeA,
+				TreeID:         treeA.OID,
 				AuthorName:     "Scrooge Mcduck",
 				AuthorEmail:    "chief@ducks.org",
 				CommitterName:  "Mickey Mouse",
@@ -90,7 +100,7 @@ func TestWriteCommit(t *testing.T) {
 				Message:        "my custom message\n\ntrailer\n",
 			},
 			expectedCommit: strings.Join([]string{
-				"tree " + string(treeA),
+				"tree " + string(treeA.OID),
 				fmt.Sprintf(
 					"author Scrooge Mcduck <chief@ducks.org> %d %s",
 					commitDate.Unix(),
@@ -110,7 +120,7 @@ func TestWriteCommit(t *testing.T) {
 		{
 			desc: "with multiple parents",
 			cfg: WriteCommitConfig{
-				TreeID:         treeA,
+				TreeID:         treeA.OID,
 				Parents:        []git.ObjectID{commitA, commitB},
 				AuthorName:     "Scrooge Mcduck",
 				AuthorEmail:    "chief@ducks.org",
@@ -121,7 +131,7 @@ func TestWriteCommit(t *testing.T) {
 				Message:        "my custom message",
 			},
 			expectedCommit: strings.Join([]string{
-				"tree " + treeA.String(),
+				"tree " + treeA.OID.String(),
 				"parent " + commitA.String(),
 				"parent " + commitB.String(),
 				fmt.Sprintf(
@@ -141,7 +151,7 @@ func TestWriteCommit(t *testing.T) {
 		{
 			desc: "with reference",
 			cfg: WriteCommitConfig{
-				TreeID:         treeA,
+				TreeID:         treeA.OID,
 				Parents:        []git.ObjectID{commitA, commitB},
 				AuthorName:     "Scrooge Mcduck",
 				AuthorEmail:    "chief@ducks.org",
@@ -153,7 +163,7 @@ func TestWriteCommit(t *testing.T) {
 				Reference:      "refs/heads/foo",
 			},
 			expectedCommit: strings.Join([]string{
-				"tree " + treeA.String(),
+				"tree " + treeA.OID.String(),
 				"parent " + commitA.String(),
 				"parent " + commitB.String(),
 				fmt.Sprintf(
@@ -203,14 +213,19 @@ func TestWriteCommit_validation(t *testing.T) {
 
 	blobID, err := repo.WriteBlob(ctx, "", strings.NewReader("foo"))
 	require.NoError(t, err)
-	treeID, err := repo.WriteTree(ctx, []*TreeEntry{
-		{
-			OID:  blobID,
-			Mode: "100644",
-			Path: "file1",
+	tree := &TreeEntry{
+		Type: Tree,
+		Mode: "040000",
+		Entries: []*TreeEntry{
+			{
+				OID:  blobID,
+				Mode: "100644",
+				Path: "file1",
+			},
 		},
-	})
-	require.NoError(t, err)
+	}
+	require.NoError(t, tree.Write(ctx, repo))
+	treeID := tree.OID
 
 	testCases := []struct {
 		desc        string
diff --git a/internal/git/localrepo/tree.go b/internal/git/localrepo/tree.go
index 8ece22088..49ffa910a 100644
--- a/internal/git/localrepo/tree.go
+++ b/internal/git/localrepo/tree.go
@@ -295,9 +295,38 @@ func (t *TreeEntry) Walk(callback func(path string, entry *TreeEntry) error) err
 	return nil
 }
 
-// WriteTree writes a new tree object from a slice of entries. This function does not verify whether OIDs
+// Write does a depth first walk, writing new trees when the OID of the
+// TreeEntry is empty.
+func (t *TreeEntry) Write(
+	ctx context.Context,
+	repo *Repo,
+) error {
+	var err error
+
+	if t.OID != "" {
+		return nil
+	}
+
+	for _, e := range t.Entries {
+		if e.Type == Tree && e.OID == "" {
+			if err := e.Write(ctx, repo); err != nil {
+				return err
+			}
+		}
+	}
+
+	treeOID, err := repo.writeEntries(ctx, t.Entries)
+	if err != nil {
+		return fmt.Errorf("writing tree: %w", err)
+	}
+
+	t.OID = treeOID
+	return nil
+}
+
+// writeTreeEntries writes a new tree object from a slice of entries. This function does not verify whether OIDs
 // referred to by tree entries actually exist in the repository.
-func (repo *Repo) WriteTree(ctx context.Context, entries []*TreeEntry) (git.ObjectID, error) {
+func (repo *Repo) writeEntries(ctx context.Context, entries []*TreeEntry) (git.ObjectID, error) {
 	var tree bytes.Buffer
 
 	sort.Stable(TreeEntriesByPath(entries))
diff --git a/internal/git/localrepo/tree_test.go b/internal/git/localrepo/tree_test.go
index dda012fd1..e4794f996 100644
--- a/internal/git/localrepo/tree_test.go
+++ b/internal/git/localrepo/tree_test.go
@@ -2,6 +2,7 @@ package localrepo
 
 import (
 	"bytes"
+	"context"
 	"os"
 	"path/filepath"
 	"sort"
@@ -17,7 +18,7 @@ import (
 	"gitlab.com/gitlab-org/gitaly/v16/internal/testhelper/testcfg"
 )
 
-func TestWriteTree(t *testing.T) {
+func TestTreeEntry_Write(t *testing.T) {
 	cfg := testcfg.Build(t)
 	ctx := testhelper.Context(t)
 
@@ -35,14 +36,19 @@ func TestWriteTree(t *testing.T) {
 	blobID, err := repo.WriteBlob(ctx, "file", bytes.NewBufferString("foobar\n"))
 	require.NoError(t, err)
 
-	treeID, err := repo.WriteTree(ctx, []*TreeEntry{
-		{
-			OID:  blobID,
-			Mode: "100644",
-			Path: "file",
+	tree := &TreeEntry{
+		Type: Tree,
+		Mode: "040000",
+		Entries: []*TreeEntry{
+			{
+				OID:  blobID,
+				Mode: "100644",
+				Path: "file",
+			},
 		},
-	})
-	require.NoError(t, err)
+	}
+	require.NoError(t, tree.Write(ctx, repo))
+	treeID := tree.OID
 
 	nonExistentBlobID, err := repo.WriteBlob(ctx, "file1", bytes.NewBufferString("content"))
 	require.NoError(t, err)
@@ -186,8 +192,14 @@ func TestWriteTree(t *testing.T) {
 		tc := tc
 		t.Run(tc.desc, func(t *testing.T) {
 			t.Parallel()
+			tree := &TreeEntry{
+				Type:    Tree,
+				Mode:    "040000",
+				Entries: tc.entries,
+			}
 
-			oid, err := repo.WriteTree(ctx, tc.entries)
+			err := tree.Write(ctx, repo)
+			oid := tree.OID
 			if tc.expectedErrString != "" {
 				if gitVersion.HashObjectFsck() {
 					switch e := err.(type) {
@@ -356,7 +368,6 @@ func TestTreeEntryByPath(t *testing.T) {
 			},
 		},
 	}
-
 	for _, tc := range testCases {
 		t.Run(tc.desc, func(t *testing.T) {
 			sort.Stable(TreeEntriesByPath(tc.input))
@@ -597,6 +608,206 @@ func TestReadTree(t *testing.T) {
 	}
 }
 
+func TestWriteTreeRecursively(t *testing.T) {
+	cfg := testcfg.Build(t)
+
+	testCases := []struct {
+		desc      string
+		setupTree func(*testing.T, string) *TreeEntry
+	}{
+		{
+			desc: "every level has a new tree",
+			setupTree: func(t *testing.T, repoPath string) *TreeEntry {
+				blobA := gittest.WriteBlob(t, cfg, repoPath, []byte("a"))
+
+				return &TreeEntry{
+					OID:  "",
+					Type: Tree,
+					Mode: "040000",
+					Entries: []*TreeEntry{
+						{
+							OID:  "",
+							Type: Tree,
+							Mode: "040000",
+							Path: "dirA",
+							Entries: []*TreeEntry{
+								{
+									OID:  "",
+									Mode: "040000",
+									Type: Tree,
+									Path: "dirB",
+									Entries: []*TreeEntry{
+										{
+											OID:  blobA,
+											Type: Blob,
+											Mode: "100644",
+											Path: "file1",
+										},
+									},
+								},
+							},
+						},
+					},
+				}
+			},
+		},
+		{
+			desc: "only some new trees",
+			setupTree: func(t *testing.T, repoPath string) *TreeEntry {
+				blobA := gittest.WriteBlob(t, cfg, repoPath, []byte("a"))
+				blobB := gittest.WriteBlob(t, cfg, repoPath, []byte("b"))
+				blobC := gittest.WriteBlob(t, cfg, repoPath, []byte("c"))
+				blobD := gittest.WriteBlob(t, cfg, repoPath, []byte("d"))
+				dirDTree := gittest.WriteTree(
+					t,
+					cfg,
+					repoPath,
+					[]gittest.TreeEntry{
+						{
+							OID:  blobC,
+							Mode: "100644",
+							Path: "file3",
+						},
+						{
+							OID:  blobD,
+							Mode: "100644",
+							Path: "file4",
+						},
+					})
+				dirCTree := gittest.WriteTree(
+					t,
+					cfg,
+					repoPath,
+					[]gittest.TreeEntry{
+						{
+							OID:  dirDTree,
+							Mode: "040000",
+							Path: "dirD",
+						},
+					},
+				)
+
+				return &TreeEntry{
+					OID:  "",
+					Type: Tree,
+					Mode: "040000",
+					Entries: []*TreeEntry{
+						{
+							OID:  "",
+							Type: Tree,
+							Mode: "040000",
+							Path: "dirA",
+							Entries: []*TreeEntry{
+								{
+									OID:  "",
+									Mode: "040000",
+									Type: Tree,
+									Path: "dirB",
+									Entries: []*TreeEntry{
+										{
+											OID:  blobA,
+											Type: Blob,
+											Mode: "100644",
+											Path: "file1",
+										},
+										{
+											OID:  blobB,
+											Type: Blob,
+											Mode: "100644",
+											Path: "file2",
+										},
+									},
+								},
+							},
+						},
+						{
+							OID:  dirCTree,
+							Type: Tree,
+							Mode: "040000",
+							Path: "dirC",
+							Entries: []*TreeEntry{
+								{
+									OID:  dirDTree,
+									Mode: "040000",
+									Type: Tree,
+									Path: "dirD",
+									Entries: []*TreeEntry{
+										{
+											OID:  blobC,
+											Type: Blob,
+											Mode: "100644",
+											Path: "file3",
+										},
+										{
+											OID:  blobD,
+											Type: Blob,
+											Mode: "100644",
+											Path: "file4",
+										},
+									},
+								},
+							},
+						},
+					},
+				}
+			},
+		},
+	}
+
+	for _, tc := range testCases {
+		tc := tc
+		t.Run(tc.desc, func(t *testing.T) {
+			t.Parallel()
+
+			ctx := testhelper.Context(t)
+			repoProto, repoPath := gittest.CreateRepository(t, ctx, cfg, gittest.CreateRepositoryConfig{
+				SkipCreationViaService: true,
+			})
+			repo := NewTestRepo(t, cfg, repoProto)
+
+			treeEntry := tc.setupTree(t, repoPath)
+			require.NoError(t, treeEntry.Write(
+				ctx,
+				repo,
+			))
+
+			requireTreeEquals(t, ctx, repo, treeEntry.OID, treeEntry)
+		})
+	}
+}
+
+func requireTreeEquals(
+	t *testing.T,
+	ctx context.Context,
+	repo *Repo,
+	treeOID git.ObjectID,
+	expect *TreeEntry,
+) {
+	tree, err := repo.ReadTree(ctx, git.Revision(treeOID))
+	require.NoError(t, err)
+
+	for i, entry := range tree.Entries {
+		if tree.Entries[i].OID != "" {
+			require.Equal(t, expect.Entries[i].Mode, entry.Mode)
+			require.Equal(t, expect.Entries[i].Type, entry.Type)
+			require.Equal(t, expect.Entries[i].OID, entry.OID)
+			require.Equal(t, expect.Entries[i].Path, entry.Path)
+		} else {
+			objectInfo, err := repo.ReadObjectInfo(
+				ctx,
+				git.Revision(entry.OID),
+			)
+			require.NoError(t, err)
+
+			require.Equal(t, expect.Entries[i].Type, ToEnum(objectInfo.Type))
+		}
+
+		if entry.Type == Tree {
+			requireTreeEquals(t, ctx, repo, entry.OID, expect.Entries[i])
+		}
+	}
+}
+
 func TestReadTree_WithRecursive(t *testing.T) {
 	t.Parallel()
 
diff --git a/internal/gitaly/service/operations/submodules.go b/internal/gitaly/service/operations/submodules.go
index 90b96ddf2..b40ea1b4c 100644
--- a/internal/gitaly/service/operations/submodules.go
+++ b/internal/gitaly/service/operations/submodules.go
@@ -129,11 +129,15 @@ func (s *Server) updateSubmodule(ctx context.Context, quarantineRepo *localrepo.
 			return "", err
 		}
 
-		newTreeID, err := quarantineRepo.WriteTree(ctx, newEntries)
-		if err != nil {
+		newTree := &localrepo.TreeEntry{
+			Type:    localrepo.Tree,
+			Mode:    "040000",
+			Entries: newEntries,
+		}
+		if err := newTree.Write(ctx, quarantineRepo); err != nil {
 			return "", fmt.Errorf("write tree: %w", err)
 		}
-		replaceWith = newTreeID
+		replaceWith = newTree.OID
 
 		if path == "." {
 			break