Stream blobs to disk and restore files in blob-locality order
Three coordinated changes drop restore wall-clock by orders of
magnitude on real-world snapshots and bring memory use back under
control:
* Streaming download into the disk cache. New
blobDiskCache.PutFromReader takes an io.Reader and io.Copy's it
straight into the cache file. The old downloadBlob path did
io.ReadAll on the decrypted plaintext stream — for a 50 GB blob
that meant 50 GB in RAM before the cache write. The whole chain
(Storage.Get → age.Decrypt → zstd.NewReader → io.Copy) is now
fully streaming; peak RAM per blob is bounded by ~64 KB of
internal age/zstd buffers plus the io.Copy buffer.
* Chunk extraction via ReadAt. After a blob is on disk, restore
reads chunks via blobDiskCache.ReadAt(hash, offset, length) so
only the chunk's bytes ever touch RAM. The previous code path
called blobCache.Get for every cache-hit chunk, which read the
entire blob (e.g. 10 GB) from disk into a []byte just to slice
out a few KB — single-handedly ~900 ms per cache hit on the
user's photo snapshot.
* Locality-aware restore ordering. New restorePlan indexes
file→blob_set and blob→file_set at restore start, then drives
the loop so that every file whose full blob set is on disk is
drained before any new blob downloads. After the drain queue
empties, the planner picks the pending file with the smallest
uncached-blob count, downloads those blobs, and drains again.
A sweep is forced right before each download so the just-
completed blob is evicted before the next one is Put, keeping
peak disk-cache occupancy at 1 for path-order-adversarial
snapshots.
The restore hot path also moves onto a restoreSession struct so
restoreFile/restoreRegularFile/etc. take only the per-call file
argument instead of threading 9+ parameters through every signature.
The new BlobRepository.GetAll method lets the session build a single
blob-id → blob-hash map at start instead of doing one DB query per
chunk.
TestRestoreLocalityAndReadAt passes: peak_len ≤ 1, get_calls = 0,
readat_calls > 0, every blob fetched exactly once.
This commit is contained in:
@@ -0,0 +1,185 @@
|
||||
package vaultik
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"math"
|
||||
"os"
|
||||
|
||||
"sneak.berlin/go/vaultik/internal/database"
|
||||
"sneak.berlin/go/vaultik/internal/types"
|
||||
)
|
||||
|
||||
// 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: chunk %s missing from blob map",
|
||||
f.Path, fc.ChunkHash.String()[:16])
|
||||
}
|
||||
hash, ok := blobIDToHash[bc.BlobID.String()]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("planning %s: blob id %s missing from id-to-hash map",
|
||||
f.Path, 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 zero FileID return means nothing is pending.
|
||||
func (p *restorePlan) pickNextDownload() types.FileID {
|
||||
var best types.FileID
|
||||
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 n < bestCount || (n == bestCount && (best.IsZero() || idStr < bestID)) {
|
||||
best = id
|
||||
bestCount = n
|
||||
bestID = idStr
|
||||
}
|
||||
}
|
||||
return best
|
||||
}
|
||||
|
||||
// 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
|
||||
}
|
||||
|
||||
// hasPending reports whether any unfinished files remain.
|
||||
func (p *restorePlan) hasPending() bool {
|
||||
return len(p.fileBlobs) > 0
|
||||
}
|
||||
Reference in New Issue
Block a user