vaultik/internal/vaultik/blobcache.go

package vaultik

import (
	"fmt"
	"os"
	"path/filepath"
	"sync"
)

// blobDiskCacheEntry tracks a cached blob on disk.
type blobDiskCacheEntry struct {
	key  string
	size int64
	prev *blobDiskCacheEntry
	next *blobDiskCacheEntry
}

// blobDiskCache is an LRU cache that stores blobs on disk instead of in memory.
// Blobs are written to a temp directory keyed by their hash. When total size
// exceeds maxBytes, the least-recently-used entries are evicted (deleted from disk).
type blobDiskCache struct {
	mu       sync.Mutex
	dir      string
	maxBytes int64
	curBytes int64
	items    map[string]*blobDiskCacheEntry
	head     *blobDiskCacheEntry // most recent
	tail     *blobDiskCacheEntry // least recent
}

// newBlobDiskCache creates a new disk-based blob cache with the given max size.
func newBlobDiskCache(maxBytes int64) (*blobDiskCache, error) {
	dir, err := os.MkdirTemp("", "vaultik-blobcache-*")
	if err != nil {
		return nil, fmt.Errorf("creating blob cache dir: %w", err)
	}
	return &blobDiskCache{
		dir:      dir,
		maxBytes: maxBytes,
		items:    make(map[string]*blobDiskCacheEntry),
	}, nil
}

func (c *blobDiskCache) path(key string) string {
	return filepath.Join(c.dir, key)
}

func (c *blobDiskCache) unlink(e *blobDiskCacheEntry) {
	if e.prev != nil {
		e.prev.next = e.next
	} else {
		c.head = e.next
	}
	if e.next != nil {
		e.next.prev = e.prev
	} else {
		c.tail = e.prev
	}
	e.prev = nil
	e.next = nil
}

func (c *blobDiskCache) pushFront(e *blobDiskCacheEntry) {
	e.prev = nil
	e.next = c.head
	if c.head != nil {
		c.head.prev = e
	}
	c.head = e
	if c.tail == nil {
		c.tail = e
	}
}

func (c *blobDiskCache) evictLRU() {
	if c.tail == nil {
		return
	}
	victim := c.tail
	c.unlink(victim)
	delete(c.items, victim.key)
	c.curBytes -= victim.size
	_ = os.Remove(c.path(victim.key))
}

// Put writes blob data to disk cache. Entries larger than maxBytes are silently skipped.
func (c *blobDiskCache) Put(key string, data []byte) error {
	entrySize := int64(len(data))

	c.mu.Lock()
	defer c.mu.Unlock()

	if entrySize > c.maxBytes {
		return nil
	}

	// Remove old entry if updating
	if e, ok := c.items[key]; ok {
		c.unlink(e)
		c.curBytes -= e.size
		_ = os.Remove(c.path(key))
		delete(c.items, key)
	}

	if err := os.WriteFile(c.path(key), data, 0600); err != nil {
		return fmt.Errorf("writing blob to cache: %w", err)
	}

	e := &blobDiskCacheEntry{key: key, size: entrySize}
	c.pushFront(e)
	c.items[key] = e
	c.curBytes += entrySize

	for c.curBytes > c.maxBytes && c.tail != nil {
		c.evictLRU()
	}

	return nil
}

// Get reads a cached blob from disk. Returns data and true on hit.
func (c *blobDiskCache) Get(key string) ([]byte, bool) {
	c.mu.Lock()
	e, ok := c.items[key]
	if !ok {
		c.mu.Unlock()
		return nil, false
	}
	c.unlink(e)
	c.pushFront(e)
	c.mu.Unlock()

	data, err := os.ReadFile(c.path(key))
	if err != nil {
		c.mu.Lock()
		if e2, ok2 := c.items[key]; ok2 && e2 == e {
			c.unlink(e)
			delete(c.items, key)
			c.curBytes -= e.size
		}
		c.mu.Unlock()
		return nil, false
	}
	return data, true
}

// ReadAt reads a slice of a cached blob without loading the entire blob into memory.
func (c *blobDiskCache) ReadAt(key string, offset, length int64) ([]byte, error) {
	c.mu.Lock()
	e, ok := c.items[key]
	if !ok {
		c.mu.Unlock()
		return nil, fmt.Errorf("key %q not in cache", key)
	}
	if offset+length > e.size {
		c.mu.Unlock()
		return nil, fmt.Errorf("read beyond blob size: offset=%d length=%d size=%d", offset, length, e.size)
	}
	c.unlink(e)
	c.pushFront(e)
	c.mu.Unlock()

	f, err := os.Open(c.path(key))
	if err != nil {
		return nil, err
	}
	defer func() { _ = f.Close() }()

	buf := make([]byte, length)
	if _, err := f.ReadAt(buf, offset); err != nil {
		return nil, err
	}
	return buf, nil
}

// Has returns whether a key exists in the cache.
func (c *blobDiskCache) Has(key string) bool {
	c.mu.Lock()
	defer c.mu.Unlock()
	_, ok := c.items[key]
	return ok
}

// Delete removes a blob from the cache and its disk file. No-op if absent.
// Used by restore's sweep logic to free blobs whose chunks have all been
// restored (so they will never be needed again during this restore).
func (c *blobDiskCache) Delete(key string) {
	c.mu.Lock()
	defer c.mu.Unlock()
	e, ok := c.items[key]
	if !ok {
		return
	}
	c.unlink(e)
	delete(c.items, key)
	c.curBytes -= e.size
	_ = os.Remove(c.path(key))
}

// Keys returns a snapshot of all cached keys. Safe for iteration without
// holding the cache lock; the cache may change concurrently.
func (c *blobDiskCache) Keys() []string {
	c.mu.Lock()
	defer c.mu.Unlock()
	keys := make([]string, 0, len(c.items))
	for k := range c.items {
		keys = append(keys, k)
	}
	return keys
}

// Size returns current total cached bytes.
func (c *blobDiskCache) Size() int64 {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.curBytes
}

// Len returns number of cached entries.
func (c *blobDiskCache) Len() int {
	c.mu.Lock()
	defer c.mu.Unlock()
	return len(c.items)
}

// Close removes the cache directory and all cached blobs.
func (c *blobDiskCache) Close() error {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.items = nil
	c.head = nil
	c.tail = nil
	c.curBytes = 0
	return os.RemoveAll(c.dir)
}