pixa/internal/imgcache/cache.go

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"
)

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

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

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

	// 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 what the cache does, made and registered by
	// registerCounters (see counters.go)
	cacheHits          prometheus.Counter
	cacheMisses        prometheus.Counter
	upstreamFetches    prometheus.Counter
	upstreamFetchBytes prometheus.Counter
	transforms         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.
func NewCache(
	db *sql.DB, config CacheConfig, metricsRegistry prometheus.Registerer,
) (*Cache, error) {
	return newCache(db, config, metricsRegistry, 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,
	metricsRegistry prometheus.Registerer,
	probe FreeSpaceProbeFunc,
) (*Cache, error) {
	log := config.Logger
	if log == nil {
		log = slog.Default()
	}

	c := &Cache{
		db:               db,
		config:           config,
		log:              log,
		disabled:         config.DisableDiskCache,
		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 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"),
	)
	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
}

// GetVariant 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 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) GetVariant(
	cacheKey imagerequest.VariantKey,
) (io.ReadCloser, int64, string, error) {
	if c.disabled {
		return nil, 0, "", ErrNotFound
	}

	return c.variants.LoadWithContentType(cacheKey)
}

// 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,
) (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, 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
}

// 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,
) error {
	if c.disabled {
		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
}

// LookupSource checks if we have cached source content for a request.
// Returns the source key if found, or an empty key if not; a row whose
// file is missing is not found. A source found has its access time
// handed to the background writer without waiting (see
// touchSourceContent); a disabled cache always reports no cached source.
func (c *Cache) LookupSource(
	ctx context.Context, req *imagerequest.Request,
) (imagerequest.SourceKey, error) {
	if c.disabled {
		return "", nil
	}

	key := req.Source.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 "", fmt.Errorf("failed to lookup source: %w", err)
	}

	if !rowExists {
		return "", nil
	}

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

	c.touchSourceContent(key)

	return key, nil
}

// StoreNegative stores a negative cache entry for a failed fetch.
func (c *Cache) StoreNegative(ctx context.Context, req *imagerequest.Request) 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
	`, req.Source.Host, req.Source.Path, req.Source.Query, 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
		}
	}
}

// IncrementStats counts a cache hit or miss, and an upstream fetch that read
// fetchBytes bytes, as IncrementUpstreamFetch does. Like the other Increment
// methods, 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, fetchBytes int64) {
	if hit {
		c.cacheHits.Inc()
	} else {
		c.cacheMisses.Inc()
	}

	c.IncrementUpstreamFetch(fetchBytes)
}

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

	c.upstreamFetches.Inc()
	c.upstreamFetchBytes.Add(float64(fetchBytes))
}

// IncrementTransformCount counts one image transcoded by the image processor.
func (c *Cache) IncrementTransformCount() {
	c.transforms.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:
	}
}

// checkNegativeCache checks if a request 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) 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)
	}

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

	return !expired, 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
}