GetTreeEntries: Limit number of flattened treessh-limit-tree-recursions

When there are many subdirectories in a folder, GetTreeEntries can get
bogged down trying to flatten all of them. Limit the number of flattened
calls to 100 to avoid this.

Relates to https://gitlab.com/gitlab-org/gitaly/-/issues/3445

Changelog: fixed

3 files changed, 49 insertions, 0 deletions

diff --git a/internal/gitaly/service/commit/tree_entries.go b/internal/gitaly/service/commit/tree_entries.go
index f7c31af84..8e9f5dfd3 100644
--- a/internal/gitaly/service/commit/tree_entries.go
+++ b/internal/gitaly/service/commit/tree_entries.go
@@ -28,6 +28,8 @@ func validateGetTreeEntriesRequest(in *gitalypb.GetTreeEntriesRequest) error {
 }
 
 func populateFlatPath(ctx context.Context, c catfile.Batch, entries []*gitalypb.TreeEntry) error {
+	flattenedTrees := 0
+
 	for _, entry := range entries {
 		entry.FlatPath = entry.Path
 
@@ -35,6 +37,11 @@ func populateFlatPath(ctx context.Context, c catfile.Batch, entries []*gitalypb.
 			continue
 		}
 
+		flattenedTrees++
+		if flattenedTrees > maxFlattenedTrees {
+			break
+		}
+
 		for i := 1; i < defaultFlatTreeRecursion; i++ {
 			subEntries, err := treeEntries(ctx, c, entry.CommitOid, string(entry.FlatPath), "", false)
 
diff --git a/internal/gitaly/service/commit/tree_entries_helper.go b/internal/gitaly/service/commit/tree_entries_helper.go
index 4654a8df2..83e4799e7 100644
--- a/internal/gitaly/service/commit/tree_entries_helper.go
+++ b/internal/gitaly/service/commit/tree_entries_helper.go
@@ -60,6 +60,7 @@ func (tef *TreeEntryFinder) FindByRevisionAndPath(ctx context.Context, revision,
 const (
 	oidSize                  = sha1.Size
 	defaultFlatTreeRecursion = 10
+	maxFlattenedTrees        = 100
 )
 
 func extractEntryInfoFromTreeData(treeData io.Reader, commitOid, rootOid, rootPath, oid string) ([]*gitalypb.TreeEntry, error) {
diff --git a/internal/gitaly/service/commit/tree_entries_test.go b/internal/gitaly/service/commit/tree_entries_test.go
index 80fb6796d..204366857 100644
--- a/internal/gitaly/service/commit/tree_entries_test.go
+++ b/internal/gitaly/service/commit/tree_entries_test.go
@@ -443,6 +443,47 @@ func TestSuccessfulGetTreeEntries_FlatPathMaxDeep_SingleFoldersStructure(t *test
 	}}, fetchedEntries)
 }
 
+func TestSuccessfulGetTreeEntries_FlatPathMaxSubdirectories_SingleFoldersStructure(t *testing.T) {
+	cfg, repo, repoPath, client := setupCommitServiceWithRepo(t, false)
+
+	for i := 0; i < maxFlattenedTrees+10; i++ {
+		folderName := filepath.Join(repoPath, fmt.Sprintf("%d", i), "next")
+		require.NoError(t, os.MkdirAll(folderName, 0755))
+		testhelper.MustRunCommand(t, nil, "touch", filepath.Join(folderName, ".gitkeep"))
+	}
+
+	gittest.Exec(t, cfg, "-C", repoPath, "add", "--all")
+	gittest.Exec(t, cfg, "-C", repoPath, "commit", "-m", "Many folders")
+
+	commitID := text.ChompBytes(gittest.Exec(t, cfg, "-C", repoPath, "rev-parse", "HEAD"))
+
+	request := &gitalypb.GetTreeEntriesRequest{
+		Repository: repo,
+		Revision:   []byte(commitID),
+		Path:       []byte("."),
+		Recursive:  false,
+	}
+
+	ctx, cancel := testhelper.Context()
+	defer cancel()
+
+	// request entries of the tree with single-folder structure on each level
+	entriesClient, err := client.GetTreeEntries(ctx, request)
+	require.NoError(t, err)
+
+	fetchedEntries := getTreeEntriesFromTreeEntryClient(t, entriesClient)
+
+	count := 0
+	for _, entry := range fetchedEntries {
+		if strings.Contains(string(entry.FlatPath), "next") {
+			count++
+		}
+	}
+
+	require.Equal(t, maxFlattenedTrees, count)
+	require.GreaterOrEqual(t, len(fetchedEntries), maxFlattenedTrees)
+}
+
 func TestFailedGetTreeEntriesRequestDueToValidationError(t *testing.T) {
 	t.Parallel()
 	_, repo, _, client := setupCommitServiceWithRepo(t, true)