backup: Try resetting refs for an optimal restorejliu-backup-reset-refs

Instead of recreating the repository from scratch during a restore, try
resetting the refs in the repository to the same state as they are in
the backup. If this doesn't work (for example if objects get pruned), we
proceed with recreating the repository from the bundle.

4 files changed, 160 insertions, 53 deletions

diff --git a/internal/backup/backup.go b/internal/backup/backup.go
index 02af48455..882327517 100644
--- a/internal/backup/backup.go
+++ b/internal/backup/backup.go
@@ -359,9 +359,18 @@ func (mgr *Manager) Restore(ctx context.Context, req *RestoreRequest) error {
 		return fmt.Errorf("manager: no backup steps")
 	}
 
-	// If we can't reset the refs, perform a full restore by recreating the repo and cloning from the bundle.
-	if err := mgr.restoreFromBundle(ctx, repo, backup, req.AlwaysCreate); err != nil {
-		return fmt.Errorf("manager: restore from bundle: %w", err)
+	// Restore Git objects, potentially from increments.
+	if err := mgr.restoreFromRefs(ctx, repo, backup); err != nil {
+		mgr.logger.WithFields(log.Fields{
+			"storage":       req.Repository.GetStorageName(),
+			"relative_path": req.Repository.GetRelativePath(),
+			"backup_id":     backup.ID,
+		}).Warn("unable to reset refs. Proceeding with a normal restore")
+
+		// If we can't reset the refs, perform a full restore by recreating the repo and cloning from the bundle.
+		if err := mgr.restoreFromBundle(ctx, repo, backup, req.AlwaysCreate); err != nil {
+			return fmt.Errorf("manager: restore from bundle: %w", err)
+		}
 	}
 
 	// Restore custom hooks. The path is the same regardless of increment.
@@ -369,6 +378,22 @@ func (mgr *Manager) Restore(ctx context.Context, req *RestoreRequest) error {
 	return mgr.restoreCustomHooks(ctx, repo, latestStep.CustomHooksPath)
 }
 
+func (mgr *Manager) restoreFromRefs(ctx context.Context, repo Repository, backup *Backup) error {
+	latestStep := backup.Steps[len(backup.Steps)-1]
+	refs, err := mgr.readRefs(ctx, latestStep.RefPath)
+	if err != nil {
+		return fmt.Errorf("read refs from backup: %w", err)
+	}
+
+	// Explicitly reset HEAD to the default branch tracked by the manifest if available. In a
+	// bundle restore, this would've been done during repository creation.
+	if defaultBranch, ok := git.ReferenceName(backup.HeadReference).Branch(); ok {
+		refs = append(refs, git.NewSymbolicReference("HEAD", git.ReferenceName(defaultBranch)))
+	}
+
+	return repo.ResetRefs(ctx, refs)
+}
+
 func (mgr *Manager) restoreFromBundle(ctx context.Context, repo Repository, backup *Backup, alwaysCreate bool) error {
 	hash, err := git.ObjectHashByFormat(backup.ObjectFormat)
 	if err != nil {
diff --git a/internal/backup/backup_test.go b/internal/backup/backup_test.go
index d5027d0be..f7cddf10c 100644
--- a/internal/backup/backup_test.go
+++ b/internal/backup/backup_test.go
@@ -1,6 +1,7 @@
 package backup_test
 
 import (
+	"bytes"
 	"context"
 	"fmt"
 	"os"
@@ -21,6 +22,7 @@ import (
 	"gitlab.com/gitlab-org/gitaly/v16/internal/gitaly/transaction"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/grpc/client"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/helper/perm"
+	"gitlab.com/gitlab-org/gitaly/v16/internal/log"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/testhelper"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/testhelper/testcfg"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/testhelper/testserver"
@@ -567,22 +569,22 @@ func TestManager_Restore_latest(t *testing.T) {
 
 	for _, managerTC := range []struct {
 		desc  string
-		setup func(t testing.TB, sink backup.Sink, locator backup.Locator) *backup.Manager
+		setup func(t testing.TB, sink backup.Sink, locator backup.Locator, logger log.LogrusLogger) *backup.Manager
 	}{
 		{
 			desc: "RPC manager",
-			setup: func(tb testing.TB, sink backup.Sink, locator backup.Locator) *backup.Manager {
+			setup: func(tb testing.TB, sink backup.Sink, locator backup.Locator, logger log.LogrusLogger) *backup.Manager {
 				pool := client.NewPool()
 				tb.Cleanup(func() {
 					testhelper.MustClose(tb, pool)
 				})
 
-				return backup.NewManager(sink, testhelper.SharedLogger(t), locator, pool)
+				return backup.NewManager(sink, logger, locator, pool)
 			},
 		},
 		{
 			desc: "Local manager",
-			setup: func(tb testing.TB, sink backup.Sink, locator backup.Locator) *backup.Manager {
+			setup: func(tb testing.TB, sink backup.Sink, locator backup.Locator, logger log.LogrusLogger) *backup.Manager {
 				if testhelper.IsPraefectEnabled() {
 					tb.Skip("local backup manager expects to operate on the local filesystem so cannot operate through praefect")
 				}
@@ -593,7 +595,7 @@ func TestManager_Restore_latest(t *testing.T) {
 				tb.Cleanup(catfileCache.Stop)
 				txManager := transaction.NewTrackingManager()
 
-				return backup.NewManagerLocal(sink, testhelper.SharedLogger(t), locator, storageLocator, gitCmdFactory, catfileCache, txManager, repoCounter)
+				return backup.NewManagerLocal(sink, logger, locator, storageLocator, gitCmdFactory, catfileCache, txManager, repoCounter)
 			},
 		},
 	} {
@@ -610,45 +612,73 @@ func TestManager_Restore_latest(t *testing.T) {
 
 			repoClient := gitalypb.NewRepositoryServiceClient(cc)
 
-			_, repoPath := gittest.CreateRepository(t, ctx, cfg)
-			commitID := gittest.WriteCommit(t, cfg, repoPath, gittest.WithBranch("main"))
-			gittest.WriteTag(t, cfg, repoPath, "v1.0.0", commitID.Revision())
-			repoChecksum := gittest.ChecksumRepo(t, cfg, repoPath)
-			repoBundle := gittest.BundleRepo(t, cfg, repoPath, "-")
-			repoRefs := gittest.Exec(t, cfg, "-C", repoPath, "show-ref", "--head")
-
 			backupRoot := testhelper.TempDir(t)
 
 			for _, tc := range []struct {
-				desc          string
-				locators      []string
-				setup         func(tb testing.TB) (*gitalypb.Repository, *git.Checksum)
-				alwaysCreate  bool
-				expectExists  bool
-				expectedPaths []string
-				expectedErrAs error
+				desc                           string
+				locators                       []string
+				setup                          func(tb testing.TB) (*gitalypb.Repository, *git.Checksum)
+				alwaysCreate                   bool
+				expectExists                   bool
+				expectedPaths                  []string
+				expectedErrAs                  error
+				shouldUseResetRefsOptimisation bool
 			}{
 				{
-					desc:     "existing repo, without hooks",
+					desc:     "non-optimised",
 					locators: []string{"legacy", "pointer"},
 					setup: func(tb testing.TB) (*gitalypb.Repository, *git.Checksum) {
-						repo, _ := gittest.CreateRepository(t, ctx, cfg)
+						_, repoPath, _ := createAndSeedRepository(t, ctx, cfg)
+						repoChecksum, repoBundle, repoRefs := createBackupArtifacts(t, cfg, repoPath)
 
+						// Restoring into an empty repo, so ref reset won't work.
+						repo, _ := gittest.CreateRepository(t, ctx, cfg)
 						relativePath := stripRelativePath(tb, repo)
 						testhelper.WriteFiles(tb, backupRoot, map[string]any{
 							relativePath + ".bundle": repoBundle,
 							relativePath + ".refs":   repoRefs,
 						})
 
+						customHooksPath := filepath.Join(backupRoot, relativePath, "custom_hooks.tar")
+						require.NoError(tb, os.MkdirAll(filepath.Join(backupRoot, relativePath), perm.PublicDir))
+						testhelper.CopyFile(tb, mustCreateCustomHooksArchive(t, ctx), customHooksPath)
+
 						return repo, repoChecksum
 					},
+					expectedPaths: []string{
+						"custom_hooks/pre-commit.sample",
+						"custom_hooks/prepare-commit-msg.sample",
+						"custom_hooks/pre-push.sample",
+					},
 					expectExists: true,
 				},
 				{
+					desc:     "existing repo, without hooks",
+					locators: []string{"legacy", "pointer"},
+					setup: func(tb testing.TB) (*gitalypb.Repository, *git.Checksum) {
+						repo, repoPath, head := createAndSeedRepository(t, ctx, cfg)
+						repoChecksum, repoBundle, repoRefs := createBackupArtifacts(t, cfg, repoPath)
+						gittest.WriteCommit(t, cfg, repoPath, gittest.WithParents(head), gittest.WithBranch("main"))
+
+						relativePath := stripRelativePath(tb, repo)
+						testhelper.WriteFiles(tb, backupRoot, map[string]any{
+							relativePath + ".bundle": repoBundle,
+							relativePath + ".refs":   repoRefs,
+						})
+
+						return repo, repoChecksum
+					},
+					expectExists:                   true,
+					shouldUseResetRefsOptimisation: true,
+				},
+
+				{
 					desc:     "existing repo, with hooks",
 					locators: []string{"legacy", "pointer"},
 					setup: func(tb testing.TB) (*gitalypb.Repository, *git.Checksum) {
-						repo, _ := gittest.CreateRepository(t, ctx, cfg)
+						repo, repoPath, head := createAndSeedRepository(t, ctx, cfg)
+						repoChecksum, repoBundle, repoRefs := createBackupArtifacts(t, cfg, repoPath)
+						gittest.WriteCommit(t, cfg, repoPath, gittest.WithParents(head), gittest.WithBranch("main"))
 
 						relativePath := stripRelativePath(tb, repo)
 						customHooksPath := filepath.Join(backupRoot, relativePath, "custom_hooks.tar")
@@ -667,7 +697,8 @@ func TestManager_Restore_latest(t *testing.T) {
 						"custom_hooks/prepare-commit-msg.sample",
 						"custom_hooks/pre-push.sample",
 					},
-					expectExists: true,
+					expectExists:                   true,
+					shouldUseResetRefsOptimisation: true,
 				},
 				{
 					desc:     "missing backup",
@@ -723,27 +754,33 @@ func TestManager_Restore_latest(t *testing.T) {
 					desc:     "nonexistent repo",
 					locators: []string{"legacy", "pointer"},
 					setup: func(tb testing.TB) (*gitalypb.Repository, *git.Checksum) {
-						repo := &gitalypb.Repository{
+						_, repoPath, _ := createAndSeedRepository(t, ctx, cfg)
+						repoChecksum, repoBundle, repoRefs := createBackupArtifacts(t, cfg, repoPath)
+
+						nonexistentRepo := &gitalypb.Repository{
 							StorageName:  "default",
 							RelativePath: gittest.NewRepositoryName(tb),
 						}
 
-						relativePath := stripRelativePath(tb, repo)
+						relativePath := stripRelativePath(tb, nonexistentRepo)
 						testhelper.WriteFiles(tb, backupRoot, map[string]any{
 							relativePath + ".bundle": repoBundle,
 							relativePath + ".refs":   repoRefs,
 						})
 
-						return repo, repoChecksum
+						return nonexistentRepo, repoChecksum
 					},
-					expectExists: true,
+					expectExists:                   true,
+					shouldUseResetRefsOptimisation: true,
 				},
 				{
 					desc:     "single incremental",
 					locators: []string{"pointer"},
 					setup: func(tb testing.TB) (*gitalypb.Repository, *git.Checksum) {
 						const backupID = "abc123"
-						repo, _ := gittest.CreateRepository(t, ctx, cfg)
+						repo, repoPath, head := createAndSeedRepository(t, ctx, cfg)
+						repoChecksum, repoBundle, repoRefs := createBackupArtifacts(t, cfg, repoPath)
+						gittest.WriteCommit(t, cfg, repoPath, gittest.WithParents(head), gittest.WithBranch("main"))
 
 						relativePath := stripRelativePath(tb, repo)
 						testhelper.WriteFiles(tb, backupRoot, map[string]any{
@@ -755,14 +792,15 @@ func TestManager_Restore_latest(t *testing.T) {
 
 						return repo, repoChecksum
 					},
-					expectExists: true,
+					expectExists:                   true,
+					shouldUseResetRefsOptimisation: true,
 				},
 				{
 					desc:     "single incremental, empty backup",
 					locators: []string{"pointer"},
 					setup: func(tb testing.TB) (*gitalypb.Repository, *git.Checksum) {
 						const backupID = "abc123"
-						repo, _ := gittest.CreateRepository(t, ctx, cfg)
+						repo, _, _ := createAndSeedRepository(t, ctx, cfg)
 
 						relativePath := stripRelativePath(tb, repo)
 						testhelper.WriteFiles(tb, backupRoot, map[string]any{
@@ -834,7 +872,8 @@ func TestManager_Restore_latest(t *testing.T) {
 
 						return repo, checksum
 					},
-					expectExists: true,
+					expectExists:                   true,
+					shouldUseResetRefsOptimisation: true,
 				},
 			} {
 				t.Run(tc.desc, func(t *testing.T) {
@@ -850,7 +889,12 @@ func TestManager_Restore_latest(t *testing.T) {
 							locator, err := backup.ResolveLocator(locatorName, sink)
 							require.NoError(t, err)
 
-							fsBackup := managerTC.setup(t, sink, locator)
+							loggerOutput := &bytes.Buffer{}
+							logger := testhelper.NewLogger(t,
+								testhelper.WithLoggerName(fmt.Sprintf("%s-%s", locatorName, tc.desc)),
+								testhelper.WithOutput(loggerOutput))
+
+							fsBackup := managerTC.setup(t, sink, locator, logger)
 							err = fsBackup.Restore(ctx, &backup.RestoreRequest{
 								Server:           storage.ServerInfo{Address: cfg.SocketPath, Token: cfg.Auth.Token},
 								Repository:       repo,
@@ -888,6 +932,12 @@ func TestManager_Restore_latest(t *testing.T) {
 									require.FileExists(t, filepath.Join(repoPath, p))
 								}
 							}
+
+							if tc.shouldUseResetRefsOptimisation {
+								require.NotContains(t, loggerOutput.String(), "unable to reset refs. Proceeding with a normal restore")
+							} else {
+								require.Contains(t, loggerOutput.String(), "unable to reset refs. Proceeding with a normal restore")
+							}
 						})
 					}
 				})
@@ -908,22 +958,22 @@ func TestManager_Restore_specific(t *testing.T) {
 
 	for _, managerTC := range []struct {
 		desc  string
-		setup func(t testing.TB, sink backup.Sink, locator backup.Locator) *backup.Manager
+		setup func(t testing.TB, sink backup.Sink, locator backup.Locator, logger log.LogrusLogger) *backup.Manager
 	}{
 		{
 			desc: "RPC manager",
-			setup: func(tb testing.TB, sink backup.Sink, locator backup.Locator) *backup.Manager {
+			setup: func(tb testing.TB, sink backup.Sink, locator backup.Locator, logger log.LogrusLogger) *backup.Manager {
 				pool := client.NewPool()
 				tb.Cleanup(func() {
 					testhelper.MustClose(tb, pool)
 				})
 
-				return backup.NewManager(sink, testhelper.SharedLogger(t), locator, pool)
+				return backup.NewManager(sink, logger, locator, pool)
 			},
 		},
 		{
 			desc: "Local manager",
-			setup: func(tb testing.TB, sink backup.Sink, locator backup.Locator) *backup.Manager {
+			setup: func(tb testing.TB, sink backup.Sink, locator backup.Locator, logger log.LogrusLogger) *backup.Manager {
 				if testhelper.IsPraefectEnabled() {
 					tb.Skip("local backup manager expects to operate on the local filesystem so cannot operate through praefect")
 				}
@@ -934,7 +984,7 @@ func TestManager_Restore_specific(t *testing.T) {
 				tb.Cleanup(catfileCache.Stop)
 				txManager := transaction.NewTrackingManager()
 
-				return backup.NewManagerLocal(sink, testhelper.SharedLogger(t), locator, storageLocator, gitCmdFactory, catfileCache, txManager, repoCounter)
+				return backup.NewManagerLocal(sink, logger, locator, storageLocator, gitCmdFactory, catfileCache, txManager, repoCounter)
 			},
 		},
 	} {
@@ -961,13 +1011,14 @@ func TestManager_Restore_specific(t *testing.T) {
 			backupRoot := testhelper.TempDir(t)
 
 			for _, tc := range []struct {
-				desc                  string
-				setup                 func(tb testing.TB) (*gitalypb.Repository, *git.Checksum)
-				alwaysCreate          bool
-				expectExists          bool
-				expectedPaths         []string
-				expectedErrAs         error
-				expectedHeadReference git.ReferenceName
+				desc                           string
+				setup                          func(tb testing.TB) (*gitalypb.Repository, *git.Checksum)
+				alwaysCreate                   bool
+				expectExists                   bool
+				expectedPaths                  []string
+				expectedErrAs                  error
+				expectedHeadReference          git.ReferenceName
+				shouldUseResetRefsOptimisation bool
 			}{
 				{
 					desc: "missing backup",
@@ -996,7 +1047,8 @@ func TestManager_Restore_specific(t *testing.T) {
 
 						return repo, repoChecksum
 					},
-					expectExists: true,
+					expectExists:                   true,
+					shouldUseResetRefsOptimisation: true,
 				},
 				{
 					desc: "many incrementals",
@@ -1056,7 +1108,8 @@ func TestManager_Restore_specific(t *testing.T) {
 
 						return repo, checksum
 					},
-					expectExists: true,
+					expectExists:                   true,
+					shouldUseResetRefsOptimisation: true,
 				},
 				{
 					desc: "manifest, empty backup",
@@ -1108,8 +1161,9 @@ custom_hooks_path = 'custom_hooks.tar'
 
 						return repo, checksum
 					},
-					expectExists:          true,
-					expectedHeadReference: "refs/heads/banana",
+					expectExists:                   true,
+					expectedHeadReference:          "refs/heads/banana",
+					shouldUseResetRefsOptimisation: true,
 				},
 			} {
 				t.Run(tc.desc, func(t *testing.T) {
@@ -1121,7 +1175,12 @@ custom_hooks_path = 'custom_hooks.tar'
 					locator, err := backup.ResolveLocator("pointer", sink)
 					require.NoError(t, err)
 
-					fsBackup := managerTC.setup(t, sink, locator)
+					loggerOutput := &bytes.Buffer{}
+					logger := testhelper.NewLogger(t,
+						testhelper.WithLoggerName(tc.desc),
+						testhelper.WithOutput(loggerOutput))
+
+					fsBackup := managerTC.setup(t, sink, locator, logger)
 					err = fsBackup.Restore(ctx, &backup.RestoreRequest{
 						Server:           storage.ServerInfo{Address: cfg.SocketPath, Token: cfg.Auth.Token},
 						Repository:       repo,
@@ -1167,6 +1226,12 @@ custom_hooks_path = 'custom_hooks.tar'
 						ref := gittest.GetSymbolicRef(t, cfg, repoPath, "HEAD")
 						require.Equal(t, tc.expectedHeadReference, git.ReferenceName(ref.Target))
 					}
+
+					if tc.shouldUseResetRefsOptimisation {
+						require.NotContains(t, loggerOutput.String(), "unable to reset refs. Proceeding with a normal restore")
+					} else {
+						require.Contains(t, loggerOutput.String(), "unable to reset refs. Proceeding with a normal restore")
+					}
 				})
 			}
 		})
@@ -1275,3 +1340,19 @@ func mustCreateCustomHooksArchive(t *testing.T, ctx context.Context) string {
 
 	return archivePath
 }
+
+func createAndSeedRepository(t *testing.T, ctx context.Context, cfg config.Cfg) (repo *gitalypb.Repository, repoPath string, head git.ObjectID) {
+	repo, repoPath = gittest.CreateRepository(t, ctx, cfg)
+	commitID := gittest.WriteCommit(t, cfg, repoPath, gittest.WithBranch("main"))
+	gittest.WriteTag(t, cfg, repoPath, "v1.0.0", commitID.Revision())
+
+	return repo, repoPath, commitID
+}
+
+func createBackupArtifacts(t *testing.T, cfg config.Cfg, repoPath string) (repoChecksum *git.Checksum, repoBundle []byte, repoRefs []byte) {
+	repoChecksum = gittest.ChecksumRepo(t, cfg, repoPath)
+	repoBundle = gittest.BundleRepo(t, cfg, repoPath, "-")
+	repoRefs = gittest.Exec(t, cfg, "-C", repoPath, "show-ref", "--head")
+
+	return repoChecksum, repoBundle, repoRefs
+}
diff --git a/internal/backup/repository_test.go b/internal/backup/repository_test.go
index 599c3bbd7..e54c1ef91 100644
--- a/internal/backup/repository_test.go
+++ b/internal/backup/repository_test.go
@@ -32,6 +32,7 @@ func removeHeadReference(refs []git.Reference) []git.Reference {
 
 func TestRemoteRepository_ResetRefs(t *testing.T) {
 	cfg := testcfg.Build(t)
+	testcfg.BuildGitalyHooks(t, cfg)
 	cfg.SocketPath = testserver.RunGitalyServer(t, cfg, setup.RegisterAll)
 	ctx := testhelper.Context(t)
 
diff --git a/internal/gitaly/service/repository/restore_repository_test.go b/internal/gitaly/service/repository/restore_repository_test.go
index ce56ec877..d3a3fdd7d 100644
--- a/internal/gitaly/service/repository/restore_repository_test.go
+++ b/internal/gitaly/service/repository/restore_repository_test.go
@@ -140,7 +140,7 @@ func TestRestoreRepository(t *testing.T) {
 					backupID: "",
 				}
 			},
-			expectedErr: structerr.NewFailedPrecondition("restore repository: manager: repository skipped: read refs: doesn't exist").WithDetail(
+			expectedErr: structerr.NewFailedPrecondition("restore repository: manager: restore from bundle: repository skipped: read refs: doesn't exist").WithDetail(
 				&gitalypb.RestoreRepositoryResponse_SkippedError{},
 			),
 		},