pixa/internal/imgcache/cache.go

// Package imgcache is the image cache. It keeps:
//   - the source images fetched from upstream
//   - the variants made from them
//   - the negative cache: the source URLs whose fetch failed recently
//
// Each image is a file with a row in the database. Eviction removes the
// least recently used images to keep the cache under its size limit.
package imgcache

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

	"github.com/prometheus/client_golang/prometheus"
	"sneak.berlin/go/pixa/internal/imagerequest"
)

// maxPendingAccessTimes is how many access times of variants served from
// the cache, and how many of cached sources a miss used, may each wait to be
// written, at about 80 bytes each. An access time handed over while that
// many of its kind wait is dropped.
const maxPendingAccessTimes = 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

// DefaultNegativeTTL is how long pixad keeps a failed fetch in the negative
// cache.
const DefaultNegativeTTL = 5 * time.Minute

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

	// MaxBytes is the disk cache size limit in bytes that eviction
	// enforces. Zero 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.
	MaxBytes int64

	// 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     Config
	log        *slog.Logger

	// 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

	// pendingVariantAccessTimes holds the variants served from the cache,
	// and pendingSourceAccessTimes the cached sources a miss used, whose
	// access time is still to be written by the goroutine that
	// StartPendingAccessTimeWrites starts, so that no request waits for the
	// database to record it. As for eviction, the channels are created in
	// NewCache, so that requests can hand over access times before that
	// goroutine starts; pendingAccessTimesDone is closed when it returns.
	// pendingAccessTimesCancel, set by StartPendingAccessTimeWrites,
	// cancels its context, which tells it to finish.
	pendingVariantAccessTimes chan imagerequest.VariantKey
	pendingSourceAccessTimes  chan imagerequest.SourceKey
	pendingAccessTimesDone    chan struct{}
	pendingAccessTimesCancel  context.CancelFunc

	// The counts of cache hits and misses, made and registered by
	// registerCounters (see counters.go)
	cacheHits   prometheus.Counter
	cacheMisses prometheus.Counter

	// reconciliationPageSize is the most rows one read of a content table
	// returns in the reconciliation pass. 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 and registers its counters on
// metricsRegistry. pixad passes Prometheus' default registry, the one
// /metrics serves. It makes no database call.
func NewCache(
	db *sql.DB, config Config, metricsRegistry prometheus.Registerer,
) (*Cache, error) {
	log := config.Logger
	if log == nil {
		log = slog.Default()
	}

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

		pendingVariantAccessTimes: make(chan imagerequest.VariantKey, maxPendingAccessTimes),
		pendingSourceAccessTimes:  make(chan imagerequest.SourceKey, maxPendingAccessTimes),
		pendingAccessTimesDone:    make(chan struct{}),

		reconciliationPageSize: defaultReconciliationPageSize,
	}

	err := c.registerCounters(metricsRegistry)
	if err != nil {
		return nil, err
	}

	if config.MaxBytes == 0 {
		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"),
	)
	if err != nil {
		return nil, fmt.Errorf("failed to create variant storage: %w", err)
	}

	c.srcContent = srcContent
	c.variants = variants

	return c, nil
}

// Variant opens a cached variant and returns a reader, its size and its
// content type, read from the variant file's first bytes. It opens that one
// file and makes no database call. A variant it opens has its access time
// handed to the background writer without waiting (see touchVariant). A
// variant not in the cache, or any variant of a disabled cache, is
// ErrNotFound; a variant file whose first bytes are not those of an output
// format is another error.
func (c *Cache) Variant(
	cacheKey imagerequest.VariantKey,
) (io.ReadCloser, int64, string, error) {
	if c.config.MaxBytes == 0 {
		return nil, 0, "", ErrNotFound
	}

	reader, size, contentType, err := c.variants.LoadWithContentType(cacheKey)
	if err != nil {
		return nil, 0, "", err
	}

	c.touchVariant(cacheKey)

	return reader, size, contentType, nil
}

// PutVariant 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) PutVariant(
	ctx context.Context, cacheKey imagerequest.VariantKey, content io.Reader,
) error {
	if c.config.MaxBytes == 0 {
		return nil
	}

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

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

	c.notifyWritePressure()

	return nil
}

// Source opens the cached source of src and returns a reader and its size
// in bytes. A source whose row or file is missing is ErrNotFound, as is any
// source of a disabled cache. A source it opens has its access time handed
// to the background writer without waiting (see touchSourceContent).
func (c *Cache) Source(
	ctx context.Context, src imagerequest.Source,
) (io.ReadCloser, int64, error) {
	if c.config.MaxBytes == 0 {
		return nil, 0, ErrNotFound
	}

	key := src.Key()

	var rowExists bool

	err := c.db.QueryRowContext(ctx, `
		SELECT EXISTS(SELECT 1 FROM source_content WHERE source_key = ?)
	`, string(key)).Scan(&rowExists)
	if err != nil {
		return nil, 0, fmt.Errorf("failed to lookup source: %w", err)
	}

	if !rowExists {
		return nil, 0, ErrNotFound
	}

	reader, size, err := c.srcContent.LoadWithSize(key)
	if err != nil {
		return nil, 0, err
	}

	c.touchSourceContent(key)

	return reader, size, nil
}

// PutSource 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) PutSource(
	ctx context.Context,
	src imagerequest.Source,
	content io.Reader,
) (imagerequest.SourceKey, error) {
	if c.config.MaxBytes == 0 {
		return "", nil
	}

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

	key := src.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, size_bytes)
		VALUES (?, ?)
		ON CONFLICT(source_key) DO UPDATE SET
			size_bytes = excluded.size_bytes,
			last_accessed_at = CURRENT_TIMESTAMP
	`, string(key), 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
}

// Failure reports whether src is in the negative cache. It only reads: a
// row past its expiry is not a hit, and the eviction goroutine deletes it
// with the other expired rows (see CleanExpired).
func (c *Cache) Failure(
	ctx context.Context, src imagerequest.Source,
) (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 = ?
	`, src.Host, src.Path, src.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)
	}

	expired := time.Now().After(expiresAt)

	return !expired, nil
}

// PutFailure stores a negative cache entry for src, whose fetch failed.
func (c *Cache) PutFailure(ctx context.Context, src imagerequest.Source) error {
	expiresAt := time.Now().UTC().Add(c.config.NegativeTTL)

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

	return nil
}

// 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
		}
	}
}

// IncrementStats counts a cache hit or miss. It only adds to the cache's
// Prometheus counters in memory (see counters.go), so it never waits for the
// database or fails.
func (c *Cache) IncrementStats(hit bool) {
	if hit {
		c.cacheHits.Inc()
	} else {
		c.cacheMisses.Inc()
	}
}

// touchVariant hands the access time of a variant served from the cache to
// the goroutine StartPendingAccessTimeWrites starts, which writes it in its
// next batch. It never waits: while maxPendingAccessTimes access times of
// variants wait to be written, this one is dropped, and the variant only
// looks colder to eviction.
func (c *Cache) touchVariant(cacheKey imagerequest.VariantKey) {
	select {
	case c.pendingVariantAccessTimes <- cacheKey:
	default:
	}
}

// touchSourceContent hands the access time of a cached source a miss used to
// the goroutine StartPendingAccessTimeWrites starts, which writes it in its
// next batch. It never waits: while maxPendingAccessTimes access times of
// sources wait to be written, this one is dropped, and the source only looks
// colder to eviction.
func (c *Cache) touchSourceContent(key imagerequest.SourceKey) {
	select {
	case c.pendingSourceAccessTimes <- key:
	default:
	}
}

// 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
}