diff --git a/internal/vaultik/blobcache.go b/internal/vaultik/blobcache.go index 313c3f9..505f796 100644 --- a/internal/vaultik/blobcache.go +++ b/internal/vaultik/blobcache.go @@ -39,9 +39,9 @@ type blobDiskCacheEntry struct { // internally and will fall back to LRU eviction if curBytes exceeds // maxBytes. Restore passes math.MaxInt64 as maxBytes and drives // eviction itself via Delete() through restoreSweeper, which deletes -// each blob the moment every file that references its chunks has been -// written. LRU never fires under that configuration; it is kept as a -// safety net for callers that don't manage eviction themselves. +// each blob once no file still to be restored in this run needs it. +// LRU never fires under that configuration; it is kept as a safety net +// for callers that don't manage eviction themselves. // // Get/ReadAt/peak-Len counters are debugging instrumentation used by // tests to assert that the restore code path uses ReadAt rather than diff --git a/internal/vaultik/restore.go b/internal/vaultik/restore.go index 9b114ee..a3c8a3f 100644 --- a/internal/vaultik/restore.go +++ b/internal/vaultik/restore.go @@ -323,12 +323,6 @@ func (v *Vaultik) restoreAllFiles( _ = blobCache.Close() }() - // Per-restore sweep state: every blob_size_limit/100 bytes written, - // scan the cache and delete any blob whose remaining file references - // are all already restored. - sweeper := newRestoreSweeper(v.ctx, repos, blobCache, - v.Config.BlobSizeLimit.Int64()/sweepIntervalDivisor) - blobByHash, blobIDToHash, err := v.buildBlobIndexes(repos) if err != nil { return nil, err @@ -340,6 +334,12 @@ func (v *Vaultik) restoreAllFiles( return nil, fmt.Errorf("building restore plan: %w", err) } + // Per-restore sweep state: every blob_size_limit/100 bytes written, + // scan the cache and delete any blob that no pending file in the + // plan references. + sweeper := newRestoreSweeper(plan, blobCache, + v.Config.BlobSizeLimit.Int64()/sweepIntervalDivisor) + // Ranging over files rather than the map reports them in a stable order. for _, file := range files { if planErr, ok := unplanned[file.ID]; ok { @@ -427,10 +427,8 @@ func (v *Vaultik) runRestoreLoop( continue } - // Record the file as restored so the sweeper can free blobs - // once all referencing files are done, and drop it from the - // plan's indexes so future picks ignore it. - session.sweeper.fileRestored(fileID.String()) + // Drop the file from the plan's indexes so future picks ignore + // it and the sweeper can free blobs no pending file needs. plan.finishFile(fileID) processed++ diff --git a/internal/vaultik/restore_interrupt_test.go b/internal/vaultik/restore_interrupt_test.go index 9d96577..28e09f0 100644 --- a/internal/vaultik/restore_interrupt_test.go +++ b/internal/vaultik/restore_interrupt_test.go @@ -74,7 +74,8 @@ func TestRestoreCleansTempDirOnInterrupt(t *testing.T) { buildLocalityFixture(t, fs, dataDir) cfg, storer, snapshotID := setupLocalityBackup( - context.Background(), t, fs, dataDir, storeDir, dbPath) + context.Background(), t, fs, dataDir, storeDir, dbPath, + localitySourceBytes) // Point the "" temp paths (the blob cache directory and the // snapshot-database directory) at a private directory so the test can diff --git a/internal/vaultik/restore_locality_test.go b/internal/vaultik/restore_locality_test.go index 67b1310..39f15a1 100644 --- a/internal/vaultik/restore_locality_test.go +++ b/internal/vaultik/restore_locality_test.go @@ -63,6 +63,10 @@ type localityCopy struct { data []byte } +// localitySourceBytes is the size of each source file of the locality +// fixture. +const localitySourceBytes = 1024 * 1024 + // buildLocalityFixture writes the adversarial source layout described // in TestRestoreLocalityAndReadAt: 15 sources of 1 MiB each (which the // backup packs into 3 blobs of 5 chunks) plus 9 interleaved-name copies @@ -73,7 +77,6 @@ func buildLocalityFixture( t.Helper() const ( - srcBytes = 1024 * 1024 srcCount = 15 blobsCount = 3 perBlob = srcCount / blobsCount @@ -83,7 +86,7 @@ func buildLocalityFixture( for i := range srcCount { s := &localitySource{ path: fmt.Sprintf("src-%03d.bin", i+1), - data: randomBytes(t, srcBytes), + data: randomBytes(t, localitySourceBytes), } sources[i] = s require.NoError(t, afero.WriteFile( @@ -115,19 +118,20 @@ func buildLocalityFixture( // setupLocalityBackup runs the backup half of the locality test: a // snapshot of dataDir into a file storer with sizes tuned so 15 one-chunk -// files pack into 3 blobs. The index database is closed before returning -// so the restore half runs from remote bytes only. +// files of srcBytes pack into 3 blobs. The index database is closed +// before returning so the restore half runs from remote bytes only. func setupLocalityBackup( ctx context.Context, t *testing.T, fs afero.Fs, - dataDir, storeDir, dbPath string, + dataDir, storeDir, dbPath string, srcBytes int64, ) (*config.Config, *storage.FileStorer, string) { t.Helper() - // chunkSize avg = 4 MiB makes minSize = 1 MiB, so a 1 MiB file - // becomes one chunk. maxBlobSize = 5 MiB packs exactly 5 chunks - // per blob, yielding 3 blobs from 15 source files. - chunkSize := int64(4 * 1024 * 1024) - maxBlobSize := int64(5 * 1024 * 1024) + // An average chunk size of four times srcBytes makes the minimum + // chunk srcBytes, so each source file becomes one chunk. A blob limit + // of five times srcBytes packs exactly 5 chunks per blob, yielding 3 + // blobs from 15 source files. + chunkSize := 4 * srcBytes + maxBlobSize := 5 * srcBytes storer, err := storage.NewFileStorer(storeDir) require.NoError(t, err) @@ -208,7 +212,7 @@ func TestRestoreLocalityAndReadAt(t *testing.T) { ctx := context.Background() cfg, storer, snapshotID := setupLocalityBackup( - ctx, t, fs, dataDir, storeDir, dbPath) + ctx, t, fs, dataDir, storeDir, dbPath, localitySourceBytes) // Wrap the storer so we can count downloads per blob key. counter := newCountingStorer(storer) diff --git a/internal/vaultik/restore_plan.go b/internal/vaultik/restore_plan.go index 2615c52..24b52d5 100644 --- a/internal/vaultik/restore_plan.go +++ b/internal/vaultik/restore_plan.go @@ -31,7 +31,8 @@ var ( // with no further downloads. // - blobFiles: for each blob, the set of pending files referencing // it. Used to short-circuit "when this blob lands, which files -// become ready" without a global scan. +// become ready" without a global scan, and by the sweeper to tell +// whether a cached blob is still needed. type restorePlan struct { fileBlobs map[types.FileID]map[string]struct{} blobFiles map[string]map[types.FileID]struct{} @@ -237,6 +238,12 @@ func (p *restorePlan) filesReferencingBlob(blobHash string) []types.FileID { return out } +// pendingFileUsesBlob reports whether any pending file references the +// named blob, whether or not the blob is cached. +func (p *restorePlan) pendingFileUsesBlob(blobHash string) bool { + return len(p.blobFiles[blobHash]) > 0 +} + // hasPending reports whether any unfinished files remain. func (p *restorePlan) hasPending() bool { return len(p.fileBlobs) > 0 diff --git a/internal/vaultik/restore_plan_test.go b/internal/vaultik/restore_plan_test.go index ce23200..0306717 100644 --- a/internal/vaultik/restore_plan_test.go +++ b/internal/vaultik/restore_plan_test.go @@ -63,15 +63,6 @@ func TestRunRestoreLoopFailsOnAbandonedFiles(t *testing.T) { t.Cleanup(func() { _ = cache.Close() }) - v := &Vaultik{ctx: ctx} - session := &restoreSession{ - v: v, - ctx: ctx, - repos: repos, - sweeper: newRestoreSweeper(ctx, repos, cache, 1), - result: &RestoreResult{}, - } - // A file that is still pending but whose uncached-blob set is empty // and which was never queued as ready: the loop can neither restore // nor download it. This is the abandonment the guard must catch. @@ -83,6 +74,15 @@ func TestRunRestoreLoopFailsOnAbandonedFiles(t *testing.T) { cached: map[string]struct{}{}, } + v := &Vaultik{ctx: ctx} + session := &restoreSession{ + v: v, + ctx: ctx, + repos: repos, + sweeper: newRestoreSweeper(plan, cache, 1), + result: &RestoreResult{}, + } + err = v.runRestoreLoop(session, plan, map[types.FileID]*database.File{}, 0) require.ErrorIs(t, err, errRestoreIncomplete) } diff --git a/internal/vaultik/restore_sweeper.go b/internal/vaultik/restore_sweeper.go index 5659f6b..004a557 100644 --- a/internal/vaultik/restore_sweeper.go +++ b/internal/vaultik/restore_sweeper.go @@ -1,46 +1,30 @@ package vaultik -import ( - "context" - "fmt" - - "sneak.berlin/go/vaultik/internal/database" - "sneak.berlin/go/vaultik/internal/log" -) - -// restoreSweeper frees cached blobs once all files that reference any of -// their chunks have been restored. It works as follows: +// restoreSweeper frees cached blobs once no file still to be restored in +// this run needs them. It works as follows: // -// 1. Callers add a file's ID to an in-memory restored set via -// fileRestored once the file is fully written to disk. -// 2. After each chunk is restored, chunkRestored accumulates a running +// 1. After each chunk is restored, chunkRestored accumulates a running // byte count. -// 3. When the accumulator crosses a threshold (one hundredth of the +// 2. When the accumulator crosses a threshold (one hundredth of the // configured blob size — so a sweep runs about a hundred times per // blob's worth of restored bytes), the sweeper iterates every key in -// the cache. For each cached blob it asks the DB which files -// reference any chunk in that blob, then compares that list against -// the in-memory restored set. If any referencing file is missing -// from the set the blob is kept; otherwise the cache entry is -// deleted. +// the cache and deletes each blob that the restore plan no longer +// needs. // -// All DB reads happen against the snapshot's temporary metadata DB, -// which is local, indexed, and not under contention — the queries are -// cheap and run at most once per blob per sweep interval. +// The plan drops a file once it is restored or has failed, and holds +// only the files this run restores, so files outside the requested paths +// never keep a blob. type restoreSweeper struct { - ctx context.Context //nolint:containedctx // ctx bound at construction by design - repos *database.Repositories + plan *restorePlan cache *blobDiskCache threshold int64 bytesAccum int64 - restored map[string]struct{} } // newRestoreSweeper returns a sweeper that triggers eviction every // `threshold` bytes restored. Callers should pass blob_size_limit/100. func newRestoreSweeper( - ctx context.Context, - repos *database.Repositories, + plan *restorePlan, cache *blobDiskCache, threshold int64, ) *restoreSweeper { @@ -49,21 +33,12 @@ func newRestoreSweeper( } return &restoreSweeper{ - ctx: ctx, - repos: repos, + plan: plan, cache: cache, threshold: threshold, - restored: make(map[string]struct{}), } } -// fileRestored records a file as fully restored. After this call, any -// blob whose only remaining references come from files in the restored -// set may be evicted on the next sweep. -func (s *restoreSweeper) fileRestored(fileID string) { - s.restored[fileID] = struct{}{} -} - // chunkRestored accounts n bytes against the sweep threshold and runs a // sweep if the threshold has been crossed since the last sweep. func (s *restoreSweeper) chunkRestored(n int64) { @@ -76,60 +51,11 @@ func (s *restoreSweeper) chunkRestored(n int64) { s.sweep() } -// sweep deletes any cached blob whose chunks are no longer referenced -// by an unrestored file. Per-blob DB failures are logged and the blob -// is kept — we'd rather hold a blob longer than risk a re-download. +// sweep deletes every cached blob that no pending file references. func (s *restoreSweeper) sweep() { for _, blobHash := range s.cache.Keys() { - needed, err := s.blobStillNeeded(blobHash) - if err != nil { - log.Debug("sweeper referencing-files query failed", - "blob_hash", blobHash[:16], "error", err) - - continue - } - - if !needed { + if !s.plan.pendingFileUsesBlob(blobHash) { s.cache.Delete(blobHash) } } } - -// blobStillNeeded returns true if any file that references a chunk in -// this blob has not yet been restored. On any error the function -// returns true — keeping the blob is always the safe answer because we -// can't prove we're done with it. -func (s *restoreSweeper) blobStillNeeded(blobHash string) (bool, error) { - rows, err := s.repos.DB().Conn().QueryContext(s.ctx, ` - SELECT DISTINCT fc.file_id - FROM file_chunks fc - JOIN blob_chunks bc ON bc.chunk_hash = fc.chunk_hash - JOIN blobs b ON b.id = bc.blob_id - WHERE b.blob_hash = ? - `, blobHash) - if err != nil { - return true, fmt.Errorf("querying referencing files: %w", err) - } - - defer func() { _ = rows.Close() }() - - for rows.Next() { - var fileID string - - err := rows.Scan(&fileID) - if err != nil { - return true, fmt.Errorf("scanning file_id: %w", err) - } - - if _, ok := s.restored[fileID]; !ok { - return true, nil - } - } - - err = rows.Err() - if err != nil { - return true, err - } - - return false, nil -} diff --git a/internal/vaultik/restore_sweeper_test.go b/internal/vaultik/restore_sweeper_test.go new file mode 100644 index 0000000..82729c9 --- /dev/null +++ b/internal/vaultik/restore_sweeper_test.go @@ -0,0 +1,163 @@ +package vaultik //nolint:testpackage // reads the restore blob cache's peak size + +import ( + "context" + "fmt" + "io" + "path/filepath" + "testing" + + "github.com/spf13/afero" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "sneak.berlin/go/vaultik/internal/config" + "sneak.berlin/go/vaultik/internal/log" + "sneak.berlin/go/vaultik/internal/storage" + "sneak.berlin/go/vaultik/internal/ui" +) + +// TestRestoreFreesBlobsNoPendingFileNeeds asserts that the restore cache +// frees a blob once no file still to be restored in this run needs it. +// Every original file has a copy with the same content, so every blob +// holds chunks of copies: a restore that leaves the copies out of its +// paths, or fails on them, must not keep the blobs for them. +func TestRestoreFreesBlobsNoPendingFileNeeds(t *testing.T) { + log.Initialize(log.Config{}) + t.Parallel() + + // Small files keep the backup and both restores quick. + const fileBytes = 16 * 1024 + + fs := afero.NewOsFs() + tempDir := t.TempDir() + + dataDir := filepath.Join(tempDir, "source") + originalsDir := filepath.Join(dataDir, "originals") + copiesDir := filepath.Join(dataDir, "copies") + storeDir := filepath.Join(tempDir, "remote") + dbPath := filepath.Join(tempDir, "index.sqlite") + + originals := writeFilesWithCopies(t, fs, originalsDir, copiesDir, fileBytes) + + ctx := context.Background() + + cfg, storer, snapshotID := setupLocalityBackup( + ctx, t, fs, dataDir, storeDir, dbPath, fileBytes) + + t.Run("restore by path", func(t *testing.T) { + t.Parallel() + + restoreDir := filepath.Join(tempDir, "by-path") + + peak, err := restorePeakCachedBlobs(ctx, t, fs, cfg, storer, + &RestoreOptions{ + SnapshotID: snapshotID, + TargetDir: restoreDir, + Paths: []string{originalsDir}, + }) + require.NoError(t, err) + + requireRestoredFiles(t, fs, restoreDir, originals) + assert.LessOrEqualf(t, peak, 1, + "peak cached blobs was %d; blobs were kept for unselected files", + peak) + }) + + t.Run("restore with failed files", func(t *testing.T) { + t.Parallel() + + restoreDir := filepath.Join(tempDir, "with-failures") + + // A regular file where the copies directory belongs makes every + // copy fail to restore. + blocker := filepath.Join(restoreDir, copiesDir) + require.NoError(t, fs.MkdirAll(filepath.Dir(blocker), 0o755)) + require.NoError(t, afero.WriteFile(fs, blocker, nil, 0o644)) + + peak, err := restorePeakCachedBlobs(ctx, t, fs, cfg, storer, + &RestoreOptions{ + SnapshotID: snapshotID, + TargetDir: restoreDir, + SkipErrors: true, + }) + require.ErrorIs(t, err, errFilesFailedRestore) + + requireRestoredFiles(t, fs, restoreDir, originals) + assert.LessOrEqualf(t, peak, 1, + "peak cached blobs was %d; blobs were kept for failed files", peak) + }) +} + +// writeFilesWithCopies writes 15 random files of fileBytes into +// originalsDir and a copy of each into copiesDir. With +// setupLocalityBackup given the same size, each file is one chunk and the +// 15 chunks fill 3 blobs. Returns the originals' content by path. +func writeFilesWithCopies( + t *testing.T, fs afero.Fs, originalsDir, copiesDir string, fileBytes int, +) map[string][]byte { + t.Helper() + + const fileCount = 15 + + require.NoError(t, fs.MkdirAll(originalsDir, 0o755)) + require.NoError(t, fs.MkdirAll(copiesDir, 0o755)) + + originals := make(map[string][]byte, fileCount) + + for i := range fileCount { + name := fmt.Sprintf("file-%02d.bin", i) + data := randomBytes(t, fileBytes) + original := filepath.Join(originalsDir, name) + originals[original] = data + + require.NoError(t, afero.WriteFile(fs, original, data, 0o644)) + require.NoError(t, afero.WriteFile( + fs, filepath.Join(copiesDir, name), data, 0o644)) + } + + return originals +} + +// requireRestoredFiles requires every file in want to have been restored +// under restoreDir with its content intact. +func requireRestoredFiles( + t *testing.T, fs afero.Fs, restoreDir string, want map[string][]byte, +) { + t.Helper() + + for path, data := range want { + got, err := afero.ReadFile(fs, filepath.Join(restoreDir, path)) + require.NoErrorf(t, err, "file missing after restore: %s", path) + require.Equalf(t, data, got, "byte mismatch for %s", path) + } +} + +// restorePeakCachedBlobs runs a restore and returns the most blobs its +// cache held at once, together with the restore's error. +func restorePeakCachedBlobs( + ctx context.Context, t *testing.T, fs afero.Fs, + cfg *config.Config, storer storage.Storer, opts *RestoreOptions, +) (int, error) { + t.Helper() + + var cache *blobDiskCache + + v := &Vaultik{ + Config: cfg, + Storage: storer, + Fs: fs, + Stdout: io.Discard, + Stderr: io.Discard, + UI: ui.NewWithColor(io.Discard, false), + restoreCacheObserver: func(c *blobDiskCache) { + cache = c + }, + } + v.SetContext(ctx) + + err := v.Restore(opts) + + require.NotNil(t, cache, "restoreCacheObserver must fire during restore") + + return cache.PeakLen(), err +}