vaultik/internal/vaultik/restore_plan.go

package vaultik

import (
	"context"
	"errors"
	"fmt"
	"math"
	"os"

	"sneak.berlin/go/vaultik/internal/database"
	"sneak.berlin/go/vaultik/internal/types"
)

// Sentinel errors for restore planning index lookups.
var (
	errPlanChunkMissing  = errors.New("chunk missing from blob map")
	errPlanBlobIDMissing = errors.New("blob id missing from id-to-hash map")
)

// restorePlan orders restore-time file processing by blob locality. The
// goal is to keep the blob disk cache occupancy as small as possible:
// download one blob, drain every file referencing only that blob, let
// the sweeper free the blob, then move on. Files that span multiple
// blobs are processed when their full blob set is on disk.
//
// The plan keeps two indexes:
//
//   - fileBlobs: for each pending file, the set of blob hashes it
//     still needs that are NOT yet in the cache. Files with an empty
//     set are "ready" — they can be restored from the current cache
//     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.
type restorePlan struct {
	fileBlobs map[types.FileID]map[string]struct{}
	blobFiles map[string]map[types.FileID]struct{}
	ready     []types.FileID
	cached    map[string]struct{}
}

// newRestorePlan builds the file→blob index for the given files. Files
// whose chunks reference no blobs (symlinks, directories) start in the
// ready queue immediately.
func newRestorePlan(
	ctx context.Context,
	repos *database.Repositories,
	files []*database.File,
	chunkToBlobMap map[string]*database.BlobChunk,
	blobIDToHash map[string]string,
) (*restorePlan, error) {
	p := &restorePlan{
		fileBlobs: make(map[types.FileID]map[string]struct{}, len(files)),
		blobFiles: make(map[string]map[types.FileID]struct{}),
		ready:     make([]types.FileID, 0, len(files)),
		cached:    make(map[string]struct{}),
	}
	for _, f := range files {
		if f.IsSymlink() || f.Mode&uint32(os.ModeDir) != 0 {
			// No chunks to fetch — restore can run immediately.
			p.fileBlobs[f.ID] = nil
			p.ready = append(p.ready, f.ID)

			continue
		}

		fileChunks, err := repos.FileChunks.GetByFileID(ctx, f.ID)
		if err != nil {
			return nil, fmt.Errorf("planning %s: %w", f.Path, err)
		}

		blobs := make(map[string]struct{})

		for _, fc := range fileChunks {
			bc, ok := chunkToBlobMap[fc.ChunkHash.String()]
			if !ok {
				return nil, fmt.Errorf("planning %s: %w: %s",
					f.Path, errPlanChunkMissing, shortHash(fc.ChunkHash.String()))
			}

			hash, ok := blobIDToHash[bc.BlobID.String()]
			if !ok {
				return nil, fmt.Errorf("planning %s: %w: %s",
					f.Path, errPlanBlobIDMissing, bc.BlobID)
			}

			blobs[hash] = struct{}{}
		}

		p.fileBlobs[f.ID] = blobs
		for hash := range blobs {
			set, ok := p.blobFiles[hash]
			if !ok {
				set = make(map[types.FileID]struct{})
				p.blobFiles[hash] = set
			}

			set[f.ID] = struct{}{}
		}

		if len(blobs) == 0 {
			p.ready = append(p.ready, f.ID)
		}
	}

	return p, nil
}

// markBlobCached records that the named blob is now resident in the
// disk cache and moves any pending file whose remaining-uncached-blobs
// set just dropped to empty onto the ready queue.
func (p *restorePlan) markBlobCached(blobHash string) {
	if _, already := p.cached[blobHash]; already {
		return
	}

	p.cached[blobHash] = struct{}{}
	for fileID := range p.blobFiles[blobHash] {
		blobs := p.fileBlobs[fileID]
		delete(blobs, blobHash)

		if len(blobs) == 0 {
			p.ready = append(p.ready, fileID)
		}
	}
}

// popReady returns the next ready file, removing it from the queue. If
// no file is ready, the second return value is false.
func (p *restorePlan) popReady() (types.FileID, bool) {
	if len(p.ready) == 0 {
		return types.FileID{}, false
	}

	id := p.ready[0]
	p.ready = p.ready[1:]

	return id, true
}

// finishFile drops a restored file from both indexes so subsequent
// planning calls don't reconsider it.
func (p *restorePlan) finishFile(fileID types.FileID) {
	for hash := range p.fileBlobs[fileID] {
		if set, ok := p.blobFiles[hash]; ok {
			delete(set, fileID)

			if len(set) == 0 {
				delete(p.blobFiles, hash)
			}
		}
	}

	delete(p.fileBlobs, fileID)
	// Also scrub the file from any blobFiles entries where it might
	// still appear even after its uncached-blob set was emptied.
	for hash, set := range p.blobFiles {
		if _, ok := set[fileID]; ok {
			delete(set, fileID)

			if len(set) == 0 {
				delete(p.blobFiles, hash)
			}
		}
	}
}

// pickNextDownload returns the pending file whose remaining-uncached
// blob set is smallest (with ties broken by FileID string compare so
// the choice is deterministic across runs). This file's blobs are
// downloaded next, after which it — together with any other pending
// files whose blob sets become empty — moves to the ready queue.
//
// The second return value is false when no file needs a download, so a
// genuine file carrying the nil UUID is picked rather than mistaken for
// "nothing left".
func (p *restorePlan) pickNextDownload() (types.FileID, bool) {
	var best types.FileID

	found := false
	bestCount := math.MaxInt

	var bestID string

	for id, blobs := range p.fileBlobs {
		n := len(blobs)
		if n == 0 {
			// Already-ready files should have been popped via
			// popReady; ignore here just in case.
			continue
		}

		idStr := id.String()
		if !found || n < bestCount || (n == bestCount && idStr < bestID) {
			best = id
			found = true
			bestCount = n
			bestID = idStr
		}
	}

	return best, found
}

// blobsNeeded returns the uncached blob hashes for fileID in any order.
func (p *restorePlan) blobsNeeded(fileID types.FileID) []string {
	blobs := p.fileBlobs[fileID]

	out := make([]string, 0, len(blobs))
	for h := range blobs {
		out = append(out, h)
	}

	return out
}

// filesReferencingBlob returns the pending files that reference the
// named blob, in any order. The result is a copy, so the caller may
// finishFile each of them while ranging over it.
func (p *restorePlan) filesReferencingBlob(blobHash string) []types.FileID {
	files := p.blobFiles[blobHash]

	out := make([]types.FileID, 0, len(files))
	for id := range files {
		out = append(out, id)
	}

	return out
}

// hasPending reports whether any unfinished files remain.
func (p *restorePlan) hasPending() bool {
	return len(p.fileBlobs) > 0
}