localrepo: Add Size() for repo

Add a Size() method for repo that calculates disk usage of objects. The
way we currently calculate repository size is to use du(1), which is a
quick and dirty way to get the total storage of a repository. However,
this ends up varying depending on the __how__ the objects are stored by
Git, which is affected by repository maintenance strategies and can vary.

A more consistent way to calculate storage is to get a total of number
of reachable objects with git rev-list --all --objects --disk-usage

2 files changed, 165 insertions, 0 deletions

diff --git a/internal/git/localrepo/repo.go b/internal/git/localrepo/repo.go
index 806b61553..7d7f3e409 100644
--- a/internal/git/localrepo/repo.go
+++ b/internal/git/localrepo/repo.go
@@ -1,8 +1,11 @@
 package localrepo
 
 import (
+	"bytes"
 	"context"
 	"fmt"
+	"strconv"
+	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -90,3 +93,29 @@ func errorWithStderr(err error, stderr []byte) error {
 	}
 	return fmt.Errorf("%w, stderr: %q", err, stderr)
 }
+
+// Size calculates the size of all reachable objects in bytes
+func (repo *Repo) Size(ctx context.Context) (int64, error) {
+	var stdout bytes.Buffer
+	if err := repo.ExecAndWait(ctx,
+		git.SubCmd{
+			Name: "rev-list",
+			Flags: []git.Option{
+				git.Flag{Name: "--all"},
+				git.Flag{Name: "--objects"},
+				git.Flag{Name: "--use-bitmap-index"},
+				git.Flag{Name: "--disk-usage"},
+			},
+		},
+		git.WithStdout(&stdout),
+	); err != nil {
+		return -1, err
+	}
+
+	size, err := strconv.ParseInt(strings.TrimSuffix(stdout.String(), "\n"), 10, 64)
+	if err != nil {
+		return -1, err
+	}
+
+	return size, nil
+}
diff --git a/internal/git/localrepo/repo_test.go b/internal/git/localrepo/repo_test.go
index 356ce2ce5..c305519c8 100644
--- a/internal/git/localrepo/repo_test.go
+++ b/internal/git/localrepo/repo_test.go
@@ -1,13 +1,19 @@
 package localrepo
 
 import (
+	"bytes"
 	"context"
+	"os"
+	"path/filepath"
 	"testing"
 
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/git"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/git/catfile"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/git/gittest"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/gitaly/config"
+	"gitlab.com/gitlab-org/gitaly/v14/internal/testhelper"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/testhelper/testcfg"
 	"gitlab.com/gitlab-org/gitaly/v14/proto/go/gitalypb"
 )
@@ -35,3 +41,133 @@ func TestRepo(t *testing.T) {
 		return New(config.NewLocator(cfg), gitCmdFactory, catfileCache, pbRepo), repoPath
 	})
 }
+
+func TestSize(t *testing.T) {
+	cfg := testcfg.Build(t)
+	gitCmdFactory := gittest.NewCommandFactory(t, cfg)
+	catfileCache := catfile.NewCache(cfg)
+	t.Cleanup(catfileCache.Stop)
+
+	testCases := []struct {
+		desc         string
+		setup        func(repoPath string, t *testing.T)
+		expectedSize int64
+	}{
+		{
+			desc:         "empty repository",
+			expectedSize: 0,
+		},
+		{
+			desc: "one committed file",
+			setup: func(repoPath string, t *testing.T) {
+				require.NoError(t, os.WriteFile(
+					filepath.Join(repoPath, "file"),
+					bytes.Repeat([]byte("a"), 1000),
+					0o644,
+				))
+
+				cmd := gittest.NewCommand(t, cfg, "-C", repoPath, "add", "file")
+				require.NoError(t, cmd.Run())
+				cmd = gittest.NewCommand(t, cfg, "-C", repoPath, "commit", "-m", "initial")
+				require.NoError(t, cmd.Run())
+			},
+			expectedSize: 202,
+		},
+		{
+			desc: "one large loose blob",
+			setup: func(repoPath string, t *testing.T) {
+				require.NoError(t, os.WriteFile(
+					filepath.Join(repoPath, "file"),
+					bytes.Repeat([]byte("a"), 1000),
+					0o644,
+				))
+
+				cmd := gittest.NewCommand(t, cfg, "-C", repoPath, "checkout", "-b", "branch-a")
+				require.NoError(t, cmd.Run())
+				cmd = gittest.NewCommand(t, cfg, "-C", repoPath, "add", "file")
+				require.NoError(t, cmd.Run())
+				cmd = gittest.NewCommand(t, cfg, "-C", repoPath, "commit", "-m", "initial")
+				require.NoError(t, cmd.Run())
+				cmd = gittest.NewCommand(t, cfg, "-C", repoPath, "update-ref", "-d", "refs/heads/branch-a")
+				require.NoError(t, cmd.Run())
+			},
+			expectedSize: 0,
+		},
+		{
+			desc: "modification to blob without repack",
+			setup: func(repoPath string, t *testing.T) {
+				require.NoError(t, os.WriteFile(
+					filepath.Join(repoPath, "file"),
+					bytes.Repeat([]byte("a"), 1000),
+					0o644,
+				))
+
+				cmd := gittest.NewCommand(t, cfg, "-C", repoPath, "add", "file")
+				require.NoError(t, cmd.Run())
+				cmd = gittest.NewCommand(t, cfg, "-C", repoPath, "commit", "-m", "initial")
+				require.NoError(t, cmd.Run())
+
+				f, err := os.OpenFile(
+					filepath.Join(repoPath, "file"),
+					os.O_APPEND|os.O_WRONLY,
+					0o644)
+				require.NoError(t, err)
+				defer f.Close()
+				_, err = f.WriteString("a")
+				assert.NoError(t, err)
+
+				cmd = gittest.NewCommand(t, cfg, "-C", repoPath, "commit", "-am", "modification")
+				require.NoError(t, cmd.Run())
+			},
+			expectedSize: 437,
+		},
+		{
+			desc: "modification to blob after repack",
+			setup: func(repoPath string, t *testing.T) {
+				require.NoError(t, os.WriteFile(
+					filepath.Join(repoPath, "file"),
+					bytes.Repeat([]byte("a"), 1000),
+					0o644,
+				))
+
+				cmd := gittest.NewCommand(t, cfg, "-C", repoPath, "add", "file")
+				require.NoError(t, cmd.Run())
+				cmd = gittest.NewCommand(t, cfg, "-C", repoPath, "commit", "-m", "initial")
+				require.NoError(t, cmd.Run())
+
+				f, err := os.OpenFile(
+					filepath.Join(repoPath, "file"),
+					os.O_APPEND|os.O_WRONLY,
+					0o644)
+				require.NoError(t, err)
+				defer f.Close()
+				_, err = f.WriteString("a")
+				assert.NoError(t, err)
+
+				cmd = gittest.NewCommand(t, cfg, "-C", repoPath, "commit", "-am", "modification")
+				require.NoError(t, cmd.Run())
+
+				cmd = gittest.NewCommand(t, cfg, "-C", repoPath, "repack", "-a", "-d")
+				require.NoError(t, cmd.Run())
+			},
+			expectedSize: 391,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.desc, func(t *testing.T) {
+			pbRepo, repoPath := gittest.InitRepo(t, cfg, cfg.Storages[0], gittest.InitRepoOpts{
+				WithWorktree: true,
+			})
+			repo := New(config.NewLocator(cfg), gitCmdFactory, catfileCache, pbRepo)
+			if tc.setup != nil {
+				tc.setup(repoPath, t)
+			}
+
+			ctx := testhelper.Context(t)
+			size, err := repo.Size(ctx)
+			require.NoError(t, err)
+			assert.Equal(t, tc.expectedSize, size)
+		})
+	}
+}