Merge branch 'benchmark_backup_repository' into 'master'

Add benchmark for BackupRepository RPC

See merge request https://gitlab.com/gitlab-org/gitaly/-/merge_requests/6305

Merged-by: Justin Tobler <jtobler@gitlab.com>
Approved-by: Justin Tobler <jtobler@gitlab.com>
Reviewed-by: Will Chandler <wchandler@gitlab.com>
Co-authored-by: Will Chandler <wchandler@gitlab.com>
Co-authored-by: James Fargher <jfargher@gitlab.com>

4 files changed, 114 insertions, 36 deletions

diff --git a/Makefile b/Makefile
index faa049bee..c7b7ec5d9 100644
--- a/Makefile
+++ b/Makefile
@@ -652,4 +652,4 @@ ${PROTOC_GEN_DOC}:    TOOL_PACKAGE = github.com/pseudomuto/protoc-gen-doc/cmd/pr
 ${DELVE}:             TOOL_PACKAGE = github.com/go-delve/delve/cmd/dlv
 
 ${BENCHMARK_REPO}:
-	${GIT} clone --bare ${GIT_QUIET} https://gitlab.com/gitlab-org/gitlab.git $@
+	${GIT} clone --mirror ${GIT_QUIET} https://gitlab.com/gitlab-org/gitlab.git $@
diff --git a/internal/backup/backup.go b/internal/backup/backup.go
index 267ec5f99..a11aa4953 100644
--- a/internal/backup/backup.go
+++ b/internal/backup/backup.go
@@ -1,6 +1,7 @@
 package backup
 
 import (
+	"bufio"
 	"bytes"
 	"context"
 	"errors"
@@ -82,7 +83,8 @@ type Repository interface {
 	ListRefs(ctx context.Context) ([]git.Reference, error)
 	// GetCustomHooks fetches the custom hooks archive.
 	GetCustomHooks(ctx context.Context, out io.Writer) error
-	// CreateBundle fetches a bundle that contains refs matching patterns.
+	// CreateBundle fetches a bundle that contains refs matching patterns. When
+	// patterns is nil all refs are bundled.
 	CreateBundle(ctx context.Context, out io.Writer, patterns io.Reader) error
 	// Remove removes the repository. Does not return an error if the
 	// repository cannot be found.
@@ -326,33 +328,39 @@ func (mgr *Manager) writeBundle(ctx context.Context, repo Repository, step *Step
 		return nil
 	}
 
-	negatedRefs, err := mgr.negatedKnownRefs(ctx, step)
-	if err != nil {
-		return fmt.Errorf("write bundle: %w", err)
-	}
-	defer func() {
-		if err := negatedRefs.Close(); err != nil && returnErr == nil {
-			returnErr = fmt.Errorf("write bundle: %w", err)
-		}
-	}()
-
-	patternReader, patternWriter := io.Pipe()
-	defer func() {
-		if err := patternReader.Close(); err != nil && returnErr == nil {
-			returnErr = fmt.Errorf("write bundle: %w", err)
+	var patterns io.Reader
+	// Full backup, no need to check for known refs.
+	if len(step.PreviousRefPath) > 0 {
+		negatedRefs, err := mgr.negatedKnownRefs(ctx, step)
+		if err != nil {
+			return fmt.Errorf("write bundle: %w", err)
 		}
-	}()
-	go func() {
-		defer patternWriter.Close()
+		defer func() {
+			if err := negatedRefs.Close(); err != nil && returnErr == nil {
+				returnErr = fmt.Errorf("write bundle: %w", err)
+			}
+		}()
 
-		for _, ref := range refs {
-			_, err := fmt.Fprintln(patternWriter, ref.Name)
-			if err != nil {
-				_ = patternWriter.CloseWithError(err)
-				return
+		patternReader, patternWriter := io.Pipe()
+		defer func() {
+			if err := patternReader.Close(); err != nil && returnErr == nil {
+				returnErr = fmt.Errorf("write bundle: %w", err)
 			}
-		}
-	}()
+		}()
+		go func() {
+			defer patternWriter.Close()
+
+			for _, ref := range refs {
+				_, err := fmt.Fprintln(patternWriter, ref.Name)
+				if err != nil {
+					_ = patternWriter.CloseWithError(err)
+					return
+				}
+			}
+		}()
+
+		patterns = io.MultiReader(negatedRefs, patternReader)
+	}
 
 	w := NewLazyWriter(func() (io.WriteCloser, error) {
 		return mgr.sink.GetWriter(ctx, step.BundlePath)
@@ -363,7 +371,7 @@ func (mgr *Manager) writeBundle(ctx context.Context, repo Repository, step *Step
 		}
 	}()
 
-	if err := repo.CreateBundle(ctx, w, io.MultiReader(negatedRefs, patternReader)); err != nil {
+	if err := repo.CreateBundle(ctx, w, patterns); err != nil {
 		if errors.Is(err, localrepo.ErrEmptyBundle) {
 			return fmt.Errorf("write bundle: %w: no changes to bundle", ErrSkipped)
 		}
@@ -493,13 +501,18 @@ func (mgr *Manager) writeRefs(ctx context.Context, path string, refs []git.Refer
 		}
 	}()
 
+	buf := bufio.NewWriter(w)
 	for _, ref := range refs {
-		_, err = fmt.Fprintf(w, "%s %s\n", ref.Target, ref.Name)
+		_, err = fmt.Fprintf(buf, "%s %s\n", ref.Target, ref.Name)
 		if err != nil {
 			return fmt.Errorf("write refs: %w", err)
 		}
 	}
 
+	if err := buf.Flush(); err != nil {
+		return fmt.Errorf("write refs: %w", err)
+	}
+
 	return nil
 }
 
diff --git a/internal/backup/repository.go b/internal/backup/repository.go
index 65207cf78..9e67e124d 100644
--- a/internal/backup/repository.go
+++ b/internal/backup/repository.go
@@ -85,13 +85,46 @@ func (rr *remoteRepository) GetCustomHooks(ctx context.Context, out io.Writer) e
 	return nil
 }
 
-// CreateBundle fetches a bundle that contains refs matching patterns.
+// CreateBundle fetches a bundle that contains refs matching patterns. When
+// patterns is nil all refs are bundled.
 func (rr *remoteRepository) CreateBundle(ctx context.Context, out io.Writer, patterns io.Reader) error {
+	if patterns != nil {
+		return rr.createBundlePatterns(ctx, out, patterns)
+	}
+
+	return rr.createBundle(ctx, out)
+}
+
+func (rr *remoteRepository) createBundle(ctx context.Context, out io.Writer) error {
 	repoClient := rr.newRepoClient()
-	stream, err := repoClient.CreateBundleFromRefList(ctx)
+	stream, err := repoClient.CreateBundle(ctx, &gitalypb.CreateBundleRequest{
+		Repository: rr.repo,
+	})
 	if err != nil {
 		return fmt.Errorf("remote repository: create bundle: %w", err)
 	}
+
+	bundle := streamio.NewReader(func() ([]byte, error) {
+		resp, err := stream.Recv()
+		if structerr.GRPCCode(err) == codes.FailedPrecondition {
+			err = localrepo.ErrEmptyBundle
+		}
+		return resp.GetData(), err
+	})
+
+	if _, err := io.Copy(out, bundle); err != nil {
+		return fmt.Errorf("remote repository: create bundle: %w", err)
+	}
+
+	return nil
+}
+
+func (rr *remoteRepository) createBundlePatterns(ctx context.Context, out io.Writer, patterns io.Reader) error {
+	repoClient := rr.newRepoClient()
+	stream, err := repoClient.CreateBundleFromRefList(ctx)
+	if err != nil {
+		return fmt.Errorf("remote repository: create bundle patterns: %w", err)
+	}
 	c := chunk.New(&createBundleFromRefListSender{
 		stream: stream,
 	})
@@ -102,7 +135,7 @@ func (rr *remoteRepository) CreateBundle(ctx context.Context, out io.Writer, pat
 		if errors.Is(err, io.EOF) {
 			break
 		} else if err != nil {
-			return fmt.Errorf("remote repository: create bundle: %w", err)
+			return fmt.Errorf("remote repository: create bundle patterns: %w", err)
 		}
 
 		line = bytes.TrimSuffix(line, []byte("\n"))
@@ -111,14 +144,14 @@ func (rr *remoteRepository) CreateBundle(ctx context.Context, out io.Writer, pat
 			Repository: rr.repo,
 			Patterns:   [][]byte{line},
 		}); err != nil {
-			return fmt.Errorf("remote repository: create bundle: %w", err)
+			return fmt.Errorf("remote repository: create bundle patterns: %w", err)
 		}
 	}
 	if err := c.Flush(); err != nil {
-		return fmt.Errorf("remote repository: create bundle: %w", err)
+		return fmt.Errorf("remote repository: create bundle patterns: %w", err)
 	}
 	if err := stream.CloseSend(); err != nil {
-		return fmt.Errorf("remote repository: create bundle: %w", err)
+		return fmt.Errorf("remote repository: create bundle patterns: %w", err)
 	}
 
 	bundle := streamio.NewReader(func() ([]byte, error) {
@@ -130,7 +163,7 @@ func (rr *remoteRepository) CreateBundle(ctx context.Context, out io.Writer, pat
 	})
 
 	if _, err := io.Copy(out, bundle); err != nil {
-		return fmt.Errorf("remote repository: create bundle: %w", err)
+		return fmt.Errorf("remote repository: create bundle patterns: %w", err)
 	}
 
 	return nil
@@ -310,7 +343,8 @@ func (r *localRepository) GetCustomHooks(ctx context.Context, out io.Writer) err
 	return nil
 }
 
-// CreateBundle fetches a bundle that contains refs matching patterns.
+// CreateBundle fetches a bundle that contains refs matching patterns. When
+// patterns is nil all refs are bundled.
 func (r *localRepository) CreateBundle(ctx context.Context, out io.Writer, patterns io.Reader) error {
 	err := r.repo.CreateBundle(ctx, out, &localrepo.CreateBundleOpts{
 		Patterns: patterns,
diff --git a/internal/gitaly/service/repository/backup_repository_test.go b/internal/gitaly/service/repository/backup_repository_test.go
index 5191d79d1..2a6b0dbcd 100644
--- a/internal/gitaly/service/repository/backup_repository_test.go
+++ b/internal/gitaly/service/repository/backup_repository_test.go
@@ -3,6 +3,7 @@ package repository
 import (
 	"context"
 	"path/filepath"
+	"strconv"
 	"strings"
 	"testing"
 
@@ -186,3 +187,33 @@ func TestServerBackupRepository(t *testing.T) {
 		})
 	}
 }
+
+func BenchmarkBackupRepository(b *testing.B) {
+	ctx := testhelper.Context(b)
+
+	backupRoot := testhelper.TempDir(b)
+	backupSink, err := backup.ResolveSink(ctx, backupRoot)
+	require.NoError(b, err)
+
+	backupLocator, err := backup.ResolveLocator("pointer", backupSink)
+	require.NoError(b, err)
+
+	cfg, client := setupRepositoryService(b,
+		testserver.WithBackupSink(backupSink),
+		testserver.WithBackupLocator(backupLocator),
+	)
+
+	repo, _ := gittest.CreateRepository(b, ctx, cfg, gittest.CreateRepositoryConfig{
+		Seed: "benchmark.git",
+	})
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		_, err := client.BackupRepository(ctx, &gitalypb.BackupRepositoryRequest{
+			Repository:       repo,
+			VanityRepository: repo,
+			BackupId:         strconv.Itoa(i),
+		})
+		require.NoError(b, err)
+	}
+}