pixa/internal/imgcache/cache.go

package imgcache

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"path/filepath"
	"sync"
	"sync/atomic"
	"time"

	lru "github.com/hashicorp/golang-lru/v2"
	"sneak.berlin/go/pixa/internal/httpfetcher"
	"sneak.berlin/go/pixa/internal/imagerequest"
)

// Cache errors.
var (
	ErrCacheMiss     = errors.New("cache miss")
	ErrNegativeCache = errors.New("negative cache hit")
)

// metaCacheSize is how many variants' content types metaCache holds. A
// variant not among them is served as before, reading its .meta file.
const metaCacheSize = 10000

// cleanExpiredBatchSize is the most expired negative cache rows one
// statement of CleanExpired deletes. Each batch is a statement of its
// own, so a request waits behind one batch at most, however many rows
// have expired.
const cleanExpiredBatchSize = 1000

// CacheConfig holds cache configuration.
type CacheConfig struct {
	StateDir    string
	NegativeTTL time.Duration

	// MaxBytes is the disk cache size limit in bytes that eviction
	// enforces. Zero means no limit is enforced (no eviction).
	MaxBytes int64

	// UseDefaultMaxBytes makes NewCache replace MaxBytes with the
	// default limit: 75% of the sum of the space free on the filesystem
	// holding the cache and the bytes the cache already holds, at least
	// DefaultCacheMaxBytesFloor. The config layer sets this when the
	// operator omits cache_max_bytes.
	UseDefaultMaxBytes bool

	// DisableDiskCache turns the disk cache off entirely: no cache
	// directories are created, lookups always miss, stores are
	// no-ops, and the eviction goroutine only removes the expired
	// negative cache rows. The config layer sets this when the
	// operator configures cache_max_bytes: 0.
	DisableDiskCache bool

	// Logger receives accounting and eviction log output. A nil
	// Logger means slog.Default().
	Logger *slog.Logger
}

// Cache implements the caching layer for the image proxy.
type Cache struct {
	db         *sql.DB
	srcContent *ContentStorage // source images by source key
	variants   *VariantStorage // processed variants by cache key
	config     CacheConfig
	log        *slog.Logger

	// disabled means the disk cache is turned off entirely: lookups
	// always miss, stores are no-ops, and no eviction runs.
	disabled bool

	// Eviction machinery. The channels are created in NewCache so
	// stores can signal write pressure without racing StartEviction.
	// evictionCancel, set by StartEviction, cancels the eviction
	// goroutine's context.
	evictionPressure chan struct{}
	evictionDone     chan struct{}
	evictionCancel   context.CancelFunc

	// The pending counts: hits, misses, upstream fetches and transforms
	// whose write to the database missed its request's deadline, kept
	// here until they are written by the goroutine that
	// StartPendingCountWrites starts. As for eviction, the channels are
	// created in NewCache: pendingCountsAdded wakes that goroutine, and
	// pendingCountsDone is closed when it returns. pendingCountsCancel,
	// set by StartPendingCountWrites, cancels its context, which tells it
	// to finish.
	pendingHits               atomic.Int64
	pendingMisses             atomic.Int64
	pendingUpstreamFetches    atomic.Int64
	pendingUpstreamFetchBytes atomic.Int64
	pendingTransforms         atomic.Int64
	pendingCountsAdded        chan struct{}
	pendingCountsDone         chan struct{}
	pendingCountsCancel       context.CancelFunc

	// Held through each write of the pending counts, and by
	// readHitAndMissCounts to read them
	pendingCountsWriteMutex sync.Mutex

	// metaCache holds the content types of the variants most recently
	// stored or served, so a hit does not read the variant's .meta file.
	// It never stands in for the variant file, which is always opened.
	metaCache *lru.Cache[imagerequest.VariantKey, string]

	// reconciliationPageSize is the most rows one read of a content table
	// returns in the reconciliation pass and in Stats. newCache sets it to
	// defaultReconciliationPageSize; tests set it smaller.
	reconciliationPageSize int

	// reconciliationReadTestHook, when set, is called after each of those
	// reads with the number of rows the read covered, so tests can check
	// that no read covers more than one page; production code leaves it
	// nil.
	reconciliationReadTestHook func(rows int)
}

// NewCache creates a new cache instance.
func NewCache(db *sql.DB, config CacheConfig) (*Cache, error) {
	return newCache(db, config, defaultFreeSpaceProbe)
}

// newCache is NewCache with the free-space probe passed in, so tests
// can fake the free space the default limit is worked out from.
func newCache(
	db *sql.DB, config CacheConfig, probe FreeSpaceProbeFunc,
) (*Cache, error) {
	log := config.Logger
	if log == nil {
		log = slog.Default()
	}

	metaCache, err := lru.New[imagerequest.VariantKey, string](metaCacheSize)
	if err != nil {
		return nil, fmt.Errorf("failed to create variant content type cache: %w", err)
	}

	c := &Cache{
		db:               db,
		config:           config,
		log:              log,
		disabled:         config.DisableDiskCache,
		evictionPressure: make(chan struct{}, 1),
		evictionDone:     make(chan struct{}),
		metaCache:        metaCache,

		pendingCountsAdded: make(chan struct{}, 1),
		pendingCountsDone:  make(chan struct{}),

		reconciliationPageSize: defaultReconciliationPageSize,
	}

	if c.disabled {
		return c, nil
	}

	srcContent, err := NewContentStorage(
		filepath.Join(config.StateDir, "cache", "sources"),
	)
	if err != nil {
		return nil, fmt.Errorf("failed to create source content storage: %w", err)
	}

	variants, err := NewVariantStorage(
		filepath.Join(config.StateDir, "cache", "variants"), log,
	)
	if err != nil {
		return nil, fmt.Errorf("failed to create variant storage: %w", err)
	}

	c.srcContent = srcContent
	c.variants = variants

	if config.UseDefaultMaxBytes {
		limit, err := c.computeDefaultMaxBytes(context.Background(), probe)
		if err != nil {
			return nil, err
		}

		c.config.MaxBytes = limit

		log.Info("computed default cache size limit from free space and cache contents",
			"cache_max_bytes", limit,
			"cache_dir", filepath.Join(config.StateDir, "cache"),
		)
	}

	return c, nil
}

// LookupResult contains the result of a cache lookup.
type LookupResult struct {
	Hit         bool
	CacheKey    imagerequest.VariantKey
	ContentType string
	SizeBytes   int64
	CacheStatus CacheStatus
}

// Lookup checks if a processed variant exists on disk. Hits touch the
// variant's LRU timestamp; a disabled cache always misses.
func (c *Cache) Lookup(
	ctx context.Context, req *imagerequest.Request,
) (*LookupResult, error) {
	cacheKey := imagerequest.CacheKey(req)

	// Check variant storage directly - no DB needed for cache hits
	if !c.disabled && c.variants.Exists(cacheKey) {
		c.touchVariant(ctx, cacheKey)

		return &LookupResult{
			Hit:         true,
			CacheKey:    cacheKey,
			CacheStatus: CacheHit,
		}, nil
	}

	return &LookupResult{
		Hit:         false,
		CacheKey:    cacheKey,
		CacheStatus: CacheMiss,
	}, nil
}

// GetVariant returns a reader, size, and content type for a cached
// variant. The content type comes from metaCache, or else from the
// variant's .meta file and is then kept in metaCache. A variant with
// no .meta file is served as application/octet-stream, which is not
// kept.
func (c *Cache) GetVariant(
	cacheKey imagerequest.VariantKey,
) (io.ReadCloser, int64, string, error) {
	if c.disabled {
		return nil, 0, "", ErrNotFound
	}

	contentType, known := c.metaCache.Get(cacheKey)
	if !known {
		return c.loadVariantWithMeta(cacheKey)
	}

	reader, size, err := c.variants.LoadWithSize(cacheKey)
	if err != nil {
		// The file is gone, e.g. deleted outside pixa
		c.metaCache.Remove(cacheKey)

		return nil, 0, "", err
	}

	return reader, size, contentType, nil
}

// StoreSource stores fetched source content: it writes the file named by
// the source's key, then the source's row. When the row cannot be written
// the file is removed and the store fails, as a file with no row is never
// served or evicted; if that removal fails too, the next reconciliation
// pass removes the file. On a disabled cache it is a no-op returning an
// empty key.
func (c *Cache) StoreSource(
	ctx context.Context,
	req *imagerequest.Request,
	content io.Reader,
	result *httpfetcher.FetchResult,
) (imagerequest.SourceKey, error) {
	if c.disabled {
		return "", nil
	}

	data, err := io.ReadAll(content)
	if err != nil {
		return "", fmt.Errorf("failed to read source content: %w", err)
	}

	key := req.Source.Key()

	size, err := c.srcContent.Store(key, data)
	if err != nil {
		return "", fmt.Errorf("failed to store source content: %w", err)
	}

	_, err = c.db.ExecContext(ctx, `
		INSERT INTO source_content
			(source_key, source_host, source_path, source_query,
			 content_type, size_bytes)
		VALUES (?, ?, ?, ?, ?, ?)
		ON CONFLICT(source_key) DO UPDATE SET
			content_type = excluded.content_type,
			size_bytes = excluded.size_bytes,
			fetched_at = CURRENT_TIMESTAMP,
			last_accessed_at = CURRENT_TIMESTAMP
	`, string(key), req.Source.Host, req.Source.Path, req.Source.Query,
		result.ContentType, size)
	if err != nil {
		deleteErr := c.srcContent.Delete(key)
		if deleteErr != nil {
			c.log.Warn("failed to remove source file whose row was not written",
				"source_key", key, "error", deleteErr)
		}

		return "", fmt.Errorf("failed to insert source content: %w", err)
	}

	c.notifyWritePressure()

	return key, nil
}

// StoreVariant stores a processed variant by its cache key and records
// it in the size accounting. On a disabled cache it is a no-op. The
// accounting insert is best-effort (the startup reconciliation pass
// adopts any variant file that misses its accounting row).
func (c *Cache) StoreVariant(
	ctx context.Context, cacheKey imagerequest.VariantKey,
	content io.Reader, contentType string,
) error {
	if c.disabled {
		return nil
	}

	size, err := c.variants.Store(cacheKey, content, contentType)
	if err != nil {
		return err
	}

	c.metaCache.Add(cacheKey, contentType)

	_, err = c.db.ExecContext(ctx, `
		INSERT INTO variant_content (cache_key, size_bytes, content_type)
		VALUES (?, ?, ?)
		ON CONFLICT(cache_key) DO UPDATE SET
			size_bytes = excluded.size_bytes,
			content_type = excluded.content_type,
			last_accessed_at = CURRENT_TIMESTAMP
	`, string(cacheKey), size, contentType)
	if err != nil {
		c.log.Warn("failed to record variant in size accounting",
			"cache_key", cacheKey, "error", err)
	}

	c.notifyWritePressure()

	return nil
}

// LookupSource checks if we have cached source content for a request.
// Returns the source key and content type if found, or empty values
// if not; a row whose file is missing is not found. Hits touch the
// source's LRU timestamp; a disabled cache always reports no cached
// source.
func (c *Cache) LookupSource(
	ctx context.Context, req *imagerequest.Request,
) (imagerequest.SourceKey, string, error) {
	if c.disabled {
		return "", "", nil
	}

	key := req.Source.Key()

	var contentType string

	err := c.db.QueryRowContext(ctx, `
		SELECT content_type FROM source_content WHERE source_key = ?
	`, string(key)).Scan(&contentType)

	if errors.Is(err, sql.ErrNoRows) {
		return "", "", nil
	}

	if err != nil {
		return "", "", fmt.Errorf("failed to lookup source: %w", err)
	}

	// Verify the content file exists
	if !c.srcContent.Exists(key) {
		return "", "", nil
	}

	c.touchSourceContent(ctx, key)

	return key, contentType, nil
}

// StoreNegative stores a negative cache entry for a failed fetch.
func (c *Cache) StoreNegative(
	ctx context.Context, req *imagerequest.Request, statusCode int, errMsg string,
) error {
	expiresAt := time.Now().UTC().Add(c.config.NegativeTTL)

	_, err := c.db.ExecContext(ctx, `
		INSERT INTO negative_cache
			(source_host, source_path, source_query, status_code,
			 error_message, expires_at)
		VALUES (?, ?, ?, ?, ?, ?)
		ON CONFLICT(source_host, source_path, source_query) DO UPDATE SET
			status_code = excluded.status_code,
			error_message = excluded.error_message,
			fetched_at = CURRENT_TIMESTAMP,
			expires_at = excluded.expires_at
	`, req.Source.Host, req.Source.Path, req.Source.Query, statusCode, errMsg, expiresAt)
	if err != nil {
		return fmt.Errorf("failed to insert negative cache: %w", err)
	}

	return nil
}

// GetSourceContent returns a reader for cached source content by its key,
// and the content's size in bytes.
func (c *Cache) GetSourceContent(
	key imagerequest.SourceKey,
) (io.ReadCloser, int64, error) {
	if c.disabled {
		return nil, 0, ErrNotFound
	}

	return c.srcContent.LoadWithSize(key)
}

// CleanExpired removes the negative cache rows past their expiry, one
// batch per statement, until none is left. Once ctx ends it starts no
// further batch and returns ctx's error. The eviction goroutine calls it
// at startup and on every periodic tick, so a row stays about
// NegativeTTL plus one eviction interval, whether or not its URL is
// requested again.
func (c *Cache) CleanExpired(ctx context.Context) error {
	for {
		if ctx.Err() != nil {
			return ctx.Err()
		}

		deleted, err := c.cleanExpiredBatch(ctx)
		if err != nil {
			return err
		}

		if deleted < cleanExpiredBatchSize {
			return nil
		}
	}
}

// Stats returns cache statistics.
func (c *Cache) Stats(ctx context.Context) (*CacheStats, error) {
	var stats CacheStats

	err := c.readHitAndMissCounts(ctx, &stats)
	if err != nil && !errors.Is(err, sql.ErrNoRows) {
		return nil, fmt.Errorf("failed to get cache stats: %w", err)
	}

	// Count and size the cached source images and processed variants from
	// their tables. A disabled cache holds none, whatever rows an earlier
	// run left.
	if !c.disabled {
		err = c.db.QueryRowContext(ctx, `
			SELECT (SELECT COUNT(*) FROM source_content)
			     + (SELECT COUNT(*) FROM variant_content)
		`).Scan(&stats.TotalItems)
		if err != nil {
			c.log.Warn("failed to count cache items for stats", "error", err)
		}

		stats.TotalSizeBytes, err = c.sumContentSizeBytes(ctx)
		if err != nil {
			c.log.Warn("failed to sum cache size for stats", "error", err)
		}
	}

	// Compute hit rate as a ratio
	if stats.HitCount+stats.MissCount > 0 {
		stats.HitRate = float64(stats.HitCount) / float64(stats.HitCount+stats.MissCount)
	}

	return &stats, nil
}

// IncrementStats counts a cache hit or miss, and an upstream fetch that read
// fetchBytes bytes, as IncrementUpstreamFetch does. Like the other Increment
// methods, it writes to the database with ctx's deadline but not its
// cancellation, so an ended request is still counted but does not wait for
// the database past its deadline; a count not written by then becomes a
// pending count (see StartPendingCountWrites).
func (c *Cache) IncrementStats(ctx context.Context, hit bool, fetchBytes int64) {
	countCtx, cancel := withoutCancelKeepingDeadline(ctx)
	defer cancel()

	var err error

	pendingCount := &c.pendingMisses

	if hit {
		pendingCount = &c.pendingHits

		_, err = c.db.ExecContext(countCtx, `
			UPDATE cache_stats
			SET hit_count = hit_count + 1,
			    last_updated_at = CURRENT_TIMESTAMP
			WHERE id = 1
		`)
	} else {
		_, err = c.db.ExecContext(countCtx, `
			UPDATE cache_stats
			SET miss_count = miss_count + 1,
			    last_updated_at = CURRENT_TIMESTAMP
			WHERE id = 1
		`)
	}

	switch {
	case errors.Is(err, context.DeadlineExceeded):
		c.addPendingCount(pendingCount, 1)
	case err != nil:
		c.log.Warn("failed to count cache hit or miss", "hit", hit, "error", err)
	}

	c.IncrementUpstreamFetch(ctx, fetchBytes)
}

// IncrementUpstreamFetch counts one upstream fetch that read fetchBytes bytes.
// A fetch that read no bytes is not counted.
func (c *Cache) IncrementUpstreamFetch(ctx context.Context, fetchBytes int64) {
	if fetchBytes <= 0 {
		return
	}

	countCtx, cancel := withoutCancelKeepingDeadline(ctx)
	defer cancel()

	_, err := c.db.ExecContext(countCtx, `
		UPDATE cache_stats
		SET upstream_fetch_count = upstream_fetch_count + 1,
		    upstream_fetch_bytes = upstream_fetch_bytes + ?,
		    last_updated_at = CURRENT_TIMESTAMP
		WHERE id = 1
	`, fetchBytes)

	switch {
	case errors.Is(err, context.DeadlineExceeded):
		c.addPendingCount(&c.pendingUpstreamFetches, 1)
		c.addPendingCount(&c.pendingUpstreamFetchBytes, fetchBytes)
	case err != nil:
		c.log.Warn("failed to count upstream fetch",
			"fetch_bytes", fetchBytes, "error", err)
	}
}

// IncrementTransformCount counts one image transcoded by the image processor.
func (c *Cache) IncrementTransformCount(ctx context.Context) {
	countCtx, cancel := withoutCancelKeepingDeadline(ctx)
	defer cancel()

	_, err := c.db.ExecContext(countCtx, `
		UPDATE cache_stats
		SET transform_count = transform_count + 1,
		    last_updated_at = CURRENT_TIMESTAMP
		WHERE id = 1
	`)

	switch {
	case errors.Is(err, context.DeadlineExceeded):
		c.addPendingCount(&c.pendingTransforms, 1)
	case err != nil:
		c.log.Warn("failed to count transform", "error", err)
	}
}

// readHitAndMissCounts sets the hit and miss counts of stats to those of the
// cache_stats row plus the pending hits and misses (see
// StartPendingCountWrites). It holds pendingCountsWriteMutex, so that no
// write of the pending counts falls between the two reads.
func (c *Cache) readHitAndMissCounts(ctx context.Context, stats *CacheStats) error {
	c.pendingCountsWriteMutex.Lock()
	defer c.pendingCountsWriteMutex.Unlock()

	err := c.db.QueryRowContext(ctx, `
		SELECT hit_count, miss_count
		FROM cache_stats WHERE id = 1
	`).Scan(&stats.HitCount, &stats.MissCount)

	stats.HitCount += c.pendingHits.Load()
	stats.MissCount += c.pendingMisses.Load()

	return err
}

// loadVariantWithMeta is GetVariant for a variant metaCache does not
// hold: it reads the content type from the variant's .meta file and
// keeps it in metaCache, unless a StoreVariant has put one there
// meanwhile, as the store's is newer. A read that finds no .meta file,
// as one can when the store's .meta write failed or an eviction removed
// the file after the read opened the variant, serves
// application/octet-stream and keeps nothing, so metaCache only ever
// holds a type read from a .meta file or passed to StoreVariant.
func (c *Cache) loadVariantWithMeta(
	cacheKey imagerequest.VariantKey,
) (io.ReadCloser, int64, string, error) {
	reader, size, contentType, err := c.variants.LoadWithMeta(cacheKey)
	if err != nil {
		return nil, 0, "", err
	}

	if contentType == "" {
		return reader, size, fallbackContentType, nil
	}

	c.metaCache.ContainsOrAdd(cacheKey, contentType)

	return reader, size, contentType, nil
}

// touchVariant updates the LRU timestamp of a variant, best-effort:
// a failed touch only makes the entry look colder to eviction.
func (c *Cache) touchVariant(ctx context.Context, cacheKey imagerequest.VariantKey) {
	_, err := c.db.ExecContext(ctx, `
		UPDATE variant_content SET last_accessed_at = CURRENT_TIMESTAMP
		WHERE cache_key = ?
	`, string(cacheKey))
	if err != nil {
		c.log.Debug("failed to touch variant LRU timestamp",
			"cache_key", cacheKey, "error", err)
	}
}

// touchSourceContent updates the LRU timestamp of a cached source,
// best-effort: a failed touch only makes the source look colder.
func (c *Cache) touchSourceContent(ctx context.Context, key imagerequest.SourceKey) {
	_, err := c.db.ExecContext(ctx, `
		UPDATE source_content SET last_accessed_at = CURRENT_TIMESTAMP
		WHERE source_key = ?
	`, string(key))
	if err != nil {
		c.log.Debug("failed to touch source content LRU timestamp",
			"source_key", key, "error", err)
	}
}

// checkNegativeCache checks if a request is in the negative cache.
func (c *Cache) checkNegativeCache(
	ctx context.Context, req *imagerequest.Request,
) (bool, error) {
	var expiresAt time.Time

	err := c.db.QueryRowContext(ctx, `
		SELECT expires_at FROM negative_cache
		WHERE source_host = ? AND source_path = ? AND source_query = ?
	`, req.Source.Host, req.Source.Path, req.Source.Query).Scan(&expiresAt)

	if errors.Is(err, sql.ErrNoRows) {
		return false, nil
	}

	if err != nil {
		return false, fmt.Errorf("failed to check negative cache: %w", err)
	}

	// Check if expired
	if time.Now().After(expiresAt) {
		// Clean up expired entry
		_, err = c.db.ExecContext(ctx, `
			DELETE FROM negative_cache
			WHERE source_host = ? AND source_path = ? AND source_query = ?
		`, req.Source.Host, req.Source.Path, req.Source.Query)
		if err != nil {
			c.log.Warn("failed to delete expired negative cache entry",
				"host", req.Source.Host, "path", req.Source.Path, "error", err)
		}

		return false, nil
	}

	return true, nil
}

// cleanExpiredBatch deletes up to cleanExpiredBatchSize negative cache
// rows past their expiry in one statement, and returns how many it
// deleted. DELETE ... LIMIT needs a SQLite build option, so the subquery
// picks the rows.
func (c *Cache) cleanExpiredBatch(ctx context.Context) (int64, error) {
	result, err := c.db.ExecContext(ctx, `
		DELETE FROM negative_cache WHERE id IN (
			SELECT id FROM negative_cache
			WHERE expires_at < CURRENT_TIMESTAMP
			LIMIT ?
		)
	`, cleanExpiredBatchSize)
	if err != nil {
		return 0, fmt.Errorf("failed to clean negative cache: %w", err)
	}

	deleted, err := result.RowsAffected()
	if err != nil {
		return 0, fmt.Errorf("failed to count cleaned negative cache rows: %w", err)
	}

	return deleted, nil
}