Sketch TreeEntryFinderjv-tree-entry-finder

4 files changed, 22 insertions, 5 deletions

diff --git a/internal/service/blob/get_blobs.go b/internal/service/blob/get_blobs.go
index 4bb70ee5f..342966129 100644
--- a/internal/service/blob/get_blobs.go
+++ b/internal/service/blob/get_blobs.go
@@ -13,11 +13,13 @@ import (
 )
 
 func sendGetBlobsResponse(req *gitalypb.GetBlobsRequest, stream gitalypb.BlobService_GetBlobsServer, c *catfile.Batch) error {
+	tef := commit.NewTreeEntryFinder(c)
+
 	for _, revisionPath := range req.RevisionPaths {
 		revision := revisionPath.Revision
 		path := revisionPath.Path
 
-		treeEntry, err := commit.TreeEntryForRevisionAndPath(c, revision, string(path))
+		treeEntry, err := tef.TreeEntryForRevisionAndPath(revision, string(path))
 		if err != nil {
 			return err
 		}
diff --git a/internal/service/commit/tree_entries_helper.go b/internal/service/commit/tree_entries_helper.go
index 369b1b2a1..4323efe44 100644
--- a/internal/service/commit/tree_entries_helper.go
+++ b/internal/service/commit/tree_entries_helper.go
@@ -50,6 +50,20 @@ func extractEntryInfoFromTreeData(treeData *bytes.Buffer, commitOid, rootOid, ro
 	return entries, nil
 }
 
+type revisionPath struct{ revision, path string }
+
+type TreeEntryFinder struct {
+	c         *catfile.Batch
+	treeCache map[revisionPath][]*gitalypb.TreeEntry
+}
+
+func NewTreeEntryFinder(c *catfile.Batch) *TreeEntryFinder {
+	return &TreeEntryFinder{
+		c:         c,
+		treeCache: make(map[revisionPath][]*gitalypb.TreeEntry),
+	}
+}
+
 func treeEntries(c *catfile.Batch, revision, path string, rootOid string, recursive bool) ([]*gitalypb.TreeEntry, error) {
 	if path == "." {
 		path = ""
@@ -124,8 +138,8 @@ func treeEntries(c *catfile.Batch, revision, path string, rootOid string, recurs
 }
 
 // TreeEntryForRevisionAndPath returns a TreeEntry struct for the object present at the revision/path pair.
-func TreeEntryForRevisionAndPath(c *catfile.Batch, revision, path string) (*gitalypb.TreeEntry, error) {
-	entries, err := treeEntries(c, revision, pathPkg.Dir(path), "", false)
+func (tef *TreeEntryFinder) TreeEntryForRevisionAndPath(revision, path string) (*gitalypb.TreeEntry, error) {
+	entries, err := treeEntries(tef.c, revision, pathPkg.Dir(path), "", false)
 	if err != nil {
 		return nil, err
 	}
diff --git a/internal/service/commit/tree_entry.go b/internal/service/commit/tree_entry.go
index 0b054803d..e8d73bc7f 100644
--- a/internal/service/commit/tree_entry.go
+++ b/internal/service/commit/tree_entry.go
@@ -14,7 +14,7 @@ import (
 )
 
 func sendTreeEntry(stream gitalypb.CommitService_TreeEntryServer, c *catfile.Batch, revision, path string, limit int64) error {
-	treeEntry, err := TreeEntryForRevisionAndPath(c, revision, path)
+	treeEntry, err := NewTreeEntryFinder(c).TreeEntryForRevisionAndPath(revision, path)
 	if err != nil {
 		return err
 	}
diff --git a/internal/service/repository/archive.go b/internal/service/repository/archive.go
index 2848910f7..a41ff090d 100644
--- a/internal/service/repository/archive.go
+++ b/internal/service/repository/archive.go
@@ -84,7 +84,8 @@ func validateGetArchivePrecondition(ctx context.Context, in *gitalypb.GetArchive
 		return err
 	}
 
-	treeEntry, err := commit.TreeEntryForRevisionAndPath(c, in.GetCommitId(), strings.TrimRight(path, "/"))
+	tef := commit.NewTreeEntryFinder(c)
+	treeEntry, err := tef.TreeEntryForRevisionAndPath(in.GetCommitId(), strings.TrimRight(path, "/"))
 	if err != nil {
 		return err
 	}