diff --git a/internal/vaultik/blobcache.go b/internal/vaultik/blobcache.go new file mode 100644 index 0000000..e7d60f4 --- /dev/null +++ b/internal/vaultik/blobcache.go @@ -0,0 +1,83 @@ +package vaultik + +import ( + "container/list" + "sync" +) + +// defaultMaxBlobCacheBytes is the default maximum size of the blob cache (1 GB). +const defaultMaxBlobCacheBytes = 1 << 30 // 1 GiB + +// blobCacheEntry is an entry in the LRU blob cache. +type blobCacheEntry struct { + key string + data []byte +} + +// blobLRUCache is a simple LRU cache for downloaded blob data, bounded by +// total byte size to prevent memory exhaustion during large restores. +type blobLRUCache struct { + mu sync.Mutex + maxBytes int64 + curBytes int64 + ll *list.List + items map[string]*list.Element +} + +// newBlobLRUCache creates a new LRU blob cache with the given maximum size in bytes. +func newBlobLRUCache(maxBytes int64) *blobLRUCache { + return &blobLRUCache{ + maxBytes: maxBytes, + ll: list.New(), + items: make(map[string]*list.Element), + } +} + +// Get returns the blob data for the given key, or nil if not cached. +func (c *blobLRUCache) Get(key string) ([]byte, bool) { + c.mu.Lock() + defer c.mu.Unlock() + + if ele, ok := c.items[key]; ok { + c.ll.MoveToFront(ele) + return ele.Value.(*blobCacheEntry).data, true + } + return nil, false +} + +// Put adds a blob to the cache, evicting least-recently-used entries if needed. +func (c *blobLRUCache) Put(key string, data []byte) { + c.mu.Lock() + defer c.mu.Unlock() + + entrySize := int64(len(data)) + + // If this single entry exceeds max, don't cache it + if entrySize > c.maxBytes { + return + } + + // If already present, update in place + if ele, ok := c.items[key]; ok { + c.ll.MoveToFront(ele) + old := ele.Value.(*blobCacheEntry) + c.curBytes += entrySize - int64(len(old.data)) + old.data = data + } else { + ele := c.ll.PushFront(&blobCacheEntry{key: key, data: data}) + c.items[key] = ele + c.curBytes += entrySize + } + + // Evict from back until under limit + for c.curBytes > c.maxBytes && c.ll.Len() > 0 { + oldest := c.ll.Back() + if oldest == nil { + break + } + entry := oldest.Value.(*blobCacheEntry) + c.ll.Remove(oldest) + delete(c.items, entry.key) + c.curBytes -= int64(len(entry.data)) + } +} diff --git a/internal/vaultik/restore.go b/internal/vaultik/restore.go index 015c533..d2b408b 100644 --- a/internal/vaultik/restore.go +++ b/internal/vaultik/restore.go @@ -109,7 +109,7 @@ func (v *Vaultik) Restore(opts *RestoreOptions) error { // Step 5: Restore files result := &RestoreResult{} - blobCache := make(map[string][]byte) // Cache downloaded and decrypted blobs + blobCache := newBlobLRUCache(defaultMaxBlobCacheBytes) // LRU cache bounded to ~1 GiB for i, file := range files { if v.ctx.Err() != nil { @@ -299,7 +299,7 @@ func (v *Vaultik) restoreFile( targetDir string, identity age.Identity, chunkToBlobMap map[string]*database.BlobChunk, - blobCache map[string][]byte, + blobCache *blobLRUCache, result *RestoreResult, ) error { // Calculate target path - use full original path under target directory @@ -383,7 +383,7 @@ func (v *Vaultik) restoreRegularFile( targetPath string, identity age.Identity, chunkToBlobMap map[string]*database.BlobChunk, - blobCache map[string][]byte, + blobCache *blobLRUCache, result *RestoreResult, ) error { // Get file chunks in order @@ -417,13 +417,13 @@ func (v *Vaultik) restoreRegularFile( // Download and decrypt blob if not cached blobHashStr := blob.Hash.String() - blobData, ok := blobCache[blobHashStr] + blobData, ok := blobCache.Get(blobHashStr) if !ok { blobData, err = v.downloadBlob(ctx, blobHashStr, blob.CompressedSize, identity) if err != nil { return fmt.Errorf("downloading blob %s: %w", blobHashStr[:16], err) } - blobCache[blobHashStr] = blobData + blobCache.Put(blobHashStr, blobData) result.BlobsDownloaded++ result.BytesDownloaded += blob.CompressedSize }