vaultik/internal/vaultik/restore_sweeper.go

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:
//
//  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
//     byte count.
//  3. 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.
//
// 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.
type restoreSweeper struct {
	ctx        context.Context //nolint:containedctx // ctx bound at construction by design
	repos      *database.Repositories
	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,
	cache *blobDiskCache,
	threshold int64,
) *restoreSweeper {
	if threshold <= 0 {
		threshold = 1
	}

	return &restoreSweeper{
		ctx:       ctx,
		repos:     repos,
		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) {
	s.bytesAccum += n
	if s.bytesAccum < s.threshold {
		return
	}

	s.bytesAccum = 0
	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.
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 {
			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
}