vaultik/internal/vaultik/restore_sweeper_integration_test.go

package vaultik_test

import (
	"context"
	"fmt"
	"io"
	"maps"
	"math/rand"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"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/database"
	"sneak.berlin/go/vaultik/internal/log"
	"sneak.berlin/go/vaultik/internal/snapshot"
	"sneak.berlin/go/vaultik/internal/storage"
	"sneak.berlin/go/vaultik/internal/ui"
	"sneak.berlin/go/vaultik/internal/vaultik"
)

// TestRestoreSweeperEvictsBlobs exercises the reference-counted blob
// disk cache eviction during restore.
//
// The scenario: 30 unique 1 MB random files plus 10 duplicates of those
// (40 files total, 30 MB of unique content) get backed up with a 10 MB
// blob_size_limit. After backup the snapshot's encrypted blobs are
// restored through Vaultik.Restore, and per-key Get counts on the
// storage layer are recorded. Each blob in the snapshot MUST be
// downloaded exactly once — re-downloads would mean the sweeper either
// evicted a blob that was still needed (LRU regression) or that the
// cache held nothing at all (broken cache).
//
// The duplicates ensure deduplicated files share blobs with their
// originals; the sweeper must keep each blob alive until BOTH the
// original AND every duplicate referencing its chunks have been
// restored.
// buildSweeperFixture writes the source layout for the sweeper test: 30
// unique 1 MB random files plus 10 duplicates of a random subset. The
// PRNG seed is fixed so failures are reproducible; the entropy is what
// matters here — the FastCDC chunker needs realistic-looking data to
// pick chunk boundaries naturally. Returns the expected content by
// path.
func buildSweeperFixture(
	t *testing.T, fs afero.Fs, dataDir string, uniqueFiles, duplicateFiles int,
) map[string][]byte {
	t.Helper()

	const fileSize = 1 * 1024 * 1024

	rng := rand.New(rand.NewSource(42)) //nolint:gosec // G404: fixture data only

	type sourceFile struct {
		path string
		data []byte
	}

	uniques := make([]sourceFile, 0, uniqueFiles)
	expected := make(map[string][]byte, uniqueFiles+duplicateFiles)

	for i := range uniqueFiles {
		data := make([]byte, fileSize)
		_, err := rng.Read(data)
		require.NoError(t, err)

		path := filepath.Join(dataDir, fmt.Sprintf("unique-%02d.bin", i))
		require.NoError(t, afero.WriteFile(fs, path, data, 0o644))
		uniques = append(uniques, sourceFile{path: path, data: data})
		expected[path] = data
	}

	// Copy a subset of the originals to fresh paths so the chunker
	// dedups them against the originals' blobs.
	for i, idx := range rng.Perm(uniqueFiles)[:duplicateFiles] {
		src := uniques[idx]
		dstPath := filepath.Join(dataDir, fmt.Sprintf("dup-%02d.bin", i))
		require.NoError(t, afero.WriteFile(fs, dstPath, src.data, 0o644))

		expected[dstPath] = src.data
	}

	return expected
}

// verifySingleDownloadPerBlob asserts each blob on disk was fetched
// exactly once. >1 means the sweeper evicted a still-needed blob; 0
// means the cache silently stopped being consulted.
func verifySingleDownloadPerBlob(
	t *testing.T, counter *countingStorer, storeDir string,
) {
	t.Helper()

	blobDownloads := 0

	for key, count := range counter.snapshot() {
		if !strings.HasPrefix(key, "blobs/") {
			continue
		}

		assert.Equalf(t, 1, count,
			"blob %s should have been downloaded exactly once during "+
				"restore, got %d", key, count)

		blobDownloads++
	}

	blobCount := countBlobsOnDisk(t, storeDir)
	assert.Equal(t, blobCount, blobDownloads,
		"every blob on disk should have been fetched exactly once during restore")
	t.Logf("restore downloaded %d blobs, each exactly once", blobDownloads)
}

// runSweeperBackup performs the backup half of the sweeper test and
// closes the index database so the restore half runs from remote bytes
// only. Returns the snapshot ID.
func runSweeperBackup(
	ctx context.Context, t *testing.T, fs afero.Fs,
	storer *storage.FileStorer, cfg *config.Config,
	dataDir, storeDir, dbPath string,
	chunkSize, maxBlobSize int64,
	uniqueFiles, duplicateFiles int,
) string {
	t.Helper()

	db, err := database.New(ctx, dbPath)
	require.NoError(t, err)

	defer func() { _ = db.Close() }()

	repos := database.NewRepositories(db)

	sm := snapshot.NewSnapshotManager(snapshot.SnapshotManagerParams{
		Repos:   repos,
		Storage: storer,
		Config:  cfg,
	})
	sm.SetFilesystem(fs)

	scanner := snapshot.NewScanner(snapshot.ScannerConfig{
		FS:               fs,
		Storage:          storer,
		ChunkSize:        chunkSize,
		MaxBlobSize:      maxBlobSize,
		CompressionLevel: cfg.CompressionLevel,
		AgeRecipients:    cfg.AgeRecipients,
		Repositories:     repos,
	})

	snapshotID, err := sm.CreateSnapshotWithName(
		ctx, cfg.Hostname, "sweeper", "test-version", "test-git")
	require.NoError(t, err)

	scanResult, err := scanner.Scan(ctx, dataDir, snapshotID)
	require.NoError(t, err)
	require.Equal(t, uniqueFiles+duplicateFiles, scanResult.FilesScanned)
	require.Greater(t, scanResult.BlobsCreated, 1,
		"30 MB of unique data at 10 MB blob size should yield multiple blobs")

	require.NoError(t, sm.CompleteSnapshot(ctx, snapshotID))
	require.NoError(t, sm.ExportSnapshotMetadata(ctx, dbPath, snapshotID))

	// Count blobs actually present on disk; this is the ground-truth
	// figure each blob's GET count must equal exactly once.
	blobCount := countBlobsOnDisk(t, storeDir)
	require.Greater(t, blobCount, 1, "expected more than one blob")
	t.Logf("backup produced %d blobs from %d files (%d unique + %d duplicates)",
		blobCount, uniqueFiles+duplicateFiles, uniqueFiles, duplicateFiles)

	// Force restore to operate without the source-side index, exactly
	// as a real restore on a fresh machine would.
	require.NoError(t, db.Close())

	return snapshotID
}

func TestRestoreSweeperEvictsBlobs(t *testing.T) {
	log.Initialize(log.Config{})
	t.Parallel()

	fs := afero.NewOsFs()
	tempDir := t.TempDir()

	dataDir := filepath.Join(tempDir, "source")
	storeDir := filepath.Join(tempDir, "remote")
	restoreDir := filepath.Join(tempDir, "restored")
	dbPath := filepath.Join(tempDir, "index.sqlite")

	require.NoError(t, fs.MkdirAll(dataDir, 0o755))

	const (
		uniqueFiles    = 30
		duplicateFiles = 10
	)

	expected := buildSweeperFixture(t, fs, dataDir, uniqueFiles, duplicateFiles)

	chunkSize := int64(64 * 1024)
	maxBlobSize := int64(10 * 1024 * 1024)

	storer, err := storage.NewFileStorer(storeDir)
	require.NoError(t, err)

	cfg := &config.Config{
		AgeRecipients:    []string{testAgePublicKey},
		AgeSecretKey:     testAgeSecretKey,
		CompressionLevel: 3,
		Hostname:         testHostname,
		BlobSizeLimit:    config.Size(maxBlobSize),
	}

	ctx := context.Background()

	snapshotID := runSweeperBackup(ctx, t, fs, storer, cfg,
		dataDir, storeDir, dbPath, chunkSize, maxBlobSize,
		uniqueFiles, duplicateFiles)

	counter := newCountingStorer(storer)

	restoreVaultik := &vaultik.Vaultik{
		Config:  cfg,
		Storage: counter,
		Fs:      fs,
		Stdout:  io.Discard,
		Stderr:  io.Discard,
		UI:      ui.NewWithColor(io.Discard, false),
	}
	restoreVaultik.SetContext(ctx)

	require.NoError(t, restoreVaultik.Restore(&vaultik.RestoreOptions{
		SnapshotID: snapshotID,
		TargetDir:  restoreDir,
	}))

	// Verify every restored file byte-matches its source.
	for origPath, want := range expected {
		restoredPath := filepath.Join(restoreDir, origPath)
		got, err := afero.ReadFile(fs, restoredPath)
		require.NoErrorf(t, err, "restored file missing: %s", restoredPath)
		require.Equalf(t, want, got, "byte mismatch for %s", origPath)
	}

	verifySingleDownloadPerBlob(t, counter, storeDir)
}

// countingStorer wraps a Storer and records the number of Get calls per
// key. Used to verify that the restore-side blob cache + sweeper avoid
// re-downloading blobs that are evicted while still needed.
type countingStorer struct {
	storage.Storer

	mu     sync.Mutex
	counts map[string]int
}

func newCountingStorer(inner storage.Storer) *countingStorer {
	return &countingStorer{Storer: inner, counts: make(map[string]int)}
}

func (c *countingStorer) Get(
	ctx context.Context, key string,
) (io.ReadCloser, error) {
	c.mu.Lock()
	c.counts[key]++
	c.mu.Unlock()

	return c.Storer.Get(ctx, key)
}

func (c *countingStorer) snapshot() map[string]int {
	c.mu.Lock()
	defer c.mu.Unlock()

	out := make(map[string]int, len(c.counts))
	maps.Copy(out, c.counts)

	return out
}

// countBlobsOnDisk walks the blobs/ tree of a FileStorer-backed store
// and returns the total number of blob files. Used to ground-truth the
// expected number of restore-time downloads.
func countBlobsOnDisk(t *testing.T, storeDir string) int {
	t.Helper()

	count := 0
	root := filepath.Join(storeDir, "blobs")
	err := filepath.Walk(root, func(_ string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		if !info.IsDir() {
			count++
		}

		return nil
	})
	require.NoError(t, err)

	return count
}