repository: Use git-cat-file to calculate repository size

Use git-cat-file to list all objects and accumulate the sizes in order
to get a number of the total size of a repository.

Changelog: changed

3 files changed, 124 insertions, 18 deletions

diff --git a/internal/gitaly/service/repository/size.go b/internal/gitaly/service/repository/size.go
index 1879aead9..6b3002c15 100644
--- a/internal/gitaly/service/repository/size.go
+++ b/internal/gitaly/service/repository/size.go
@@ -10,7 +10,13 @@ import (
 	"github.com/grpc-ecosystem/go-grpc-middleware/logging/logrus/ctxlogrus"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/command"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/git"
+	"gitlab.com/gitlab-org/gitaly/v15/internal/git/catfile"
+	"gitlab.com/gitlab-org/gitaly/v15/internal/git/gitpipe"
+	"gitlab.com/gitlab-org/gitaly/v15/internal/git/housekeeping"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/git/localrepo"
+	"gitlab.com/gitlab-org/gitaly/v15/internal/git/objectpool"
+	"gitlab.com/gitlab-org/gitaly/v15/internal/gitaly/storage"
+	"gitlab.com/gitlab-org/gitaly/v15/internal/gitaly/transaction"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/metadata/featureflag"
 	"gitlab.com/gitlab-org/gitaly/v15/proto/go/gitalypb"
 )
@@ -20,40 +26,126 @@ func (s *server) RepositorySize(ctx context.Context, in *gitalypb.RepositorySize
 	var size int64
 	var err error
 
-	var excludes []string
-	for refPrefix := range git.InternalRefPrefixes {
-		excludes = append(excludes, refPrefix+"*")
-	}
-
 	path, err := repo.Path()
 	if err != nil {
 		return nil, err
 	}
 
+	duSize := getPathSize(ctx, path)
+
 	if featureflag.RevlistForRepoSize.IsEnabled(ctx) {
-		size, err = repo.Size(
+		size, err = calculateSizeWithRevlist(ctx, repo)
+		if err != nil {
+			return nil, fmt.Errorf("calculating repository size with git-rev-list: %w,", err)
+		}
+
+		ctxlogrus.Extract(ctx).
+			WithField("repo_size_revlist", size).
+			WithField("repo_size_du", duSize).
+			Info("repository size calculated")
+	} else if featureflag.CatfileRepoSize.IsEnabled(ctx) {
+		size, err = calculateSizeWithCatfile(
 			ctx,
-			localrepo.WithExcludeRefs(excludes...),
-			localrepo.WithoutAlternates(),
+			repo,
+			s.locator,
+			s.gitCmdFactory,
+			s.catfileCache,
+			s.txManager,
+			s.housekeepingManager,
 		)
 		if err != nil {
-			return nil, err
+			return nil, fmt.Errorf("calculating repository size with git-cat-file: %w", err)
 		}
-		// return the size in kb to remain consistent
-		size = size / 1024
-
-		duSize := getPathSize(ctx, path)
 
 		ctxlogrus.Extract(ctx).
-			WithField("repo_size_revlist", size).
-			WithField("repo_size_du", duSize).Info("repository size calculated")
+			WithField("repo_size_catfile", size).
+			WithField("repo_size_du", duSize).
+			Info("repository size calculated")
 	} else {
-		size = getPathSize(ctx, path)
+		size = duSize
 	}
 
 	return &gitalypb.RepositorySizeResponse{Size: size}, nil
 }
 
+// calculateSizeWithCatfile calculates the repository size using git-cat-file.
+// In the case the repository belongs to a pool, it will subract the total
+// size of the pool repository objects from its total size. One limitation of
+// this approach is that we don't distinguish whether an object in the pool
+// repository belongs to the fork repository, so in fact we may end up with a
+// smaller total size and theoretically could go negative.
+func calculateSizeWithCatfile(
+	ctx context.Context,
+	repo *localrepo.Repo,
+	locator storage.Locator,
+	gitCmdFactory git.CommandFactory,
+	catfileCache catfile.Cache,
+	txManager transaction.Manager,
+	housekeepingManager housekeeping.Manager,
+) (int64, error) {
+	var size int64
+
+	catfileInfoIterator := gitpipe.CatfileInfoAllObjects(
+		ctx,
+		repo,
+	)
+
+	for catfileInfoIterator.Next() {
+		size += catfileInfoIterator.Result().ObjectSize()
+	}
+
+	if err := catfileInfoIterator.Err(); err != nil {
+		return 0, err
+	}
+
+	var poolSize int64
+
+	if pool, err := objectpool.FromRepo(
+		locator,
+		gitCmdFactory,
+		catfileCache,
+		txManager,
+		housekeepingManager,
+		repo,
+	); err == nil && pool != nil {
+		catfileInfoIterator = gitpipe.CatfileInfoAllObjects(
+			ctx,
+			pool.Repo,
+		)
+
+		for catfileInfoIterator.Next() {
+			poolSize += catfileInfoIterator.Result().ObjectSize()
+		}
+
+		if err := catfileInfoIterator.Err(); err != nil {
+			return 0, err
+		}
+	}
+
+	size -= poolSize
+	// return the size in kb
+	return size / 1024, nil
+}
+
+func calculateSizeWithRevlist(ctx context.Context, repo *localrepo.Repo) (int64, error) {
+	var excludes []string
+	for refPrefix := range git.InternalRefPrefixes {
+		excludes = append(excludes, refPrefix+"*")
+	}
+
+	size, err := repo.Size(
+		ctx,
+		localrepo.WithExcludeRefs(excludes...),
+		localrepo.WithoutAlternates(),
+	)
+	if err != nil {
+		return 0, err
+	}
+
+	// return the size in kb to remain consistent
+	return (size / 1024), nil
+}
+
 func (s *server) GetObjectDirectorySize(ctx context.Context, in *gitalypb.GetObjectDirectorySizeRequest) (*gitalypb.GetObjectDirectorySizeResponse, error) {
 	repo := s.localrepo(in.GetRepository())
 
diff --git a/internal/gitaly/service/repository/size_test.go b/internal/gitaly/service/repository/size_test.go
index e4bc784d3..98f129bea 100644
--- a/internal/gitaly/service/repository/size_test.go
+++ b/internal/gitaly/service/repository/size_test.go
@@ -30,7 +30,10 @@ const testRepoMinSizeKB = 10000
 func TestRepositorySize_SuccessfulRequest(t *testing.T) {
 	t.Parallel()
 
-	featureSet := testhelper.NewFeatureSets(featureflag.RevlistForRepoSize)
+	featureSet := testhelper.NewFeatureSets(
+		featureflag.RevlistForRepoSize,
+		featureflag.CatfileRepoSize,
+	)
 
 	featureSet.Run(t, testSuccessfulRepositorySizeRequest)
 	featureSet.Run(t, testSuccessfulRepositorySizeRequestPoolMember)
@@ -94,7 +97,8 @@ func testSuccessfulRepositorySizeRequestPoolMember(t *testing.T, ctx context.Con
 	response, err = repoClient.RepositorySize(ctx, sizeRequest)
 	require.NoError(t, err)
 
-	if featureflag.RevlistForRepoSize.IsEnabled(ctx) {
+	if featureflag.RevlistForRepoSize.IsEnabled(ctx) ||
+		featureflag.CatfileRepoSize.IsEnabled(ctx) {
 		assert.Equal(t, int64(0), response.GetSize())
 	} else {
 		assert.Less(t, response.GetSize(), sizeBeforePool)
diff --git a/internal/metadata/featureflag/ff_catfile_for_repo_size.go b/internal/metadata/featureflag/ff_catfile_for_repo_size.go
new file mode 100644
index 000000000..d3f409818
--- /dev/null
+++ b/internal/metadata/featureflag/ff_catfile_for_repo_size.go
@@ -0,0 +1,10 @@
+package featureflag
+
+// CatfileRepoSize will enable the rate limiter to reject requests beyond a configured
+// rate.
+var CatfileRepoSize = NewFeatureFlag(
+	"catfile_repo_size",
+	"v15.3.0",
+	"https://gitlab.com/gitlab-org/gitaly/-/issues/4421",
+	false,
+)