scratchbox/internal/http/raw_cache.go

package httpapi

import (
	"container/list"
	"sync"
)

const (
	defaultRawCacheMaxEntries = 32
	defaultRawCacheMaxBytes   = 64 * 1024 * 1024
)

type rawContentCache struct {
	mu         sync.Mutex
	maxEntries int
	maxBytes   int64
	totalBytes int64
	items      map[string]*list.Element
	lru        *list.List
	inFlight   map[string]*rawContentLoad
}

type rawContentItem struct {
	id   string
	data []byte
	size int64
}

type rawContentLoad struct {
	wait chan struct{}
	data []byte
	err  error
}

// copyBytes returns a defensive copy so callers cannot mutate the cache's backing
// slice (which would affect concurrent/future hits for the same ID). Cost is
// acceptable: cache is small (default 32 entries) and holds decompressed raw
// content for range requests only.
func copyBytes(b []byte) []byte {
	if b == nil {
		return nil
	}
	c := make([]byte, len(b))
	copy(c, b)
	return c
}

func newRawContentCache(maxEntries int, maxBytes int64) *rawContentCache {
	if maxEntries <= 0 {
		maxEntries = defaultRawCacheMaxEntries
	}
	if maxBytes <= 0 {
		maxBytes = defaultRawCacheMaxBytes
	}
	return &rawContentCache{
		maxEntries: maxEntries,
		maxBytes:   maxBytes,
		items:      make(map[string]*list.Element, maxEntries),
		lru:        list.New(),
		inFlight:   make(map[string]*rawContentLoad),
	}
}

func (c *rawContentCache) GetOrLoad(id string, loader func() ([]byte, error)) ([]byte, error, bool) {
	c.mu.Lock()
	if elem, ok := c.items[id]; ok {
		c.lru.MoveToFront(elem)
		item := elem.Value.(*rawContentItem)
		data := copyBytes(item.data)
		c.mu.Unlock()
		return data, nil, true
	}
	if load, ok := c.inFlight[id]; ok {
		c.mu.Unlock()
		<-load.wait
		return copyBytes(load.data), load.err, false
	}

	load := &rawContentLoad{wait: make(chan struct{})}
	c.inFlight[id] = load
	c.mu.Unlock()

	data, err := loader()
	load.data = data
	load.err = err

	c.mu.Lock()
	delete(c.inFlight, id)
	if err == nil {
		c.setLocked(id, data)
	}
	c.mu.Unlock()

	close(load.wait)
	return copyBytes(data), err, false
}

func (c *rawContentCache) Delete(id string) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.deleteLocked(id)
}

func (c *rawContentCache) setLocked(id string, data []byte) {
	size := int64(len(data))
	if size <= 0 || size > c.maxBytes {
		c.deleteLocked(id)
		return
	}

	if elem, ok := c.items[id]; ok {
		item := elem.Value.(*rawContentItem)
		c.totalBytes -= item.size
		item.data = data
		item.size = size
		c.totalBytes += size
		c.lru.MoveToFront(elem)
		c.evictLocked()
		return
	}

	item := &rawContentItem{
		id:   id,
		data: data,
		size: size,
	}
	elem := c.lru.PushFront(item)
	c.items[id] = elem
	c.totalBytes += size
	c.evictLocked()
}

func (c *rawContentCache) deleteLocked(id string) {
	elem, ok := c.items[id]
	if !ok {
		return
	}
	item := elem.Value.(*rawContentItem)
	c.totalBytes -= item.size
	delete(c.items, id)
	c.lru.Remove(elem)
}

func (c *rawContentCache) evictLocked() {
	for c.totalBytes > c.maxBytes || len(c.items) > c.maxEntries {
		last := c.lru.Back()
		if last == nil {
			return
		}
		item := last.Value.(*rawContentItem)
		c.totalBytes -= item.size
		delete(c.items, item.id)
		c.lru.Remove(last)
	}
}