pixa/internal/imgcache/pendingcounts.go

package imgcache

import (
	"context"
	"errors"
	"fmt"
	"time"

	"sneak.berlin/go/pixa/internal/imagerequest"
)

// pendingWriteInterval is how often the goroutine StartPendingCountWrites
// starts writes the pending access times and counts. SQLite's
// CURRENT_TIMESTAMP, which an access time is written as, counts whole
// seconds.
const pendingWriteInterval = time.Second

// StartPendingCountWrites starts the goroutine that writes, every
// pendingWriteInterval, the pending access times in one transaction and then
// the pending counts in one UPDATE. Requests only hand these to it, so none
// waits for the database to record what it did, and only this goroutine's
// write ever waits for it. It is a no-op when already started. The goroutine
// outlives the caller, so it runs with its own context, which
// StopPendingCountWrites cancels to tell it to finish.
func (c *Cache) StartPendingCountWrites() {
	if c.pendingCountsCancel != nil {
		return
	}

	ctx, cancel := context.WithCancel(context.Background())
	c.pendingCountsCancel = cancel

	go func() {
		c.pendingCountWriteLoop(ctx)
	}()
}

// StopPendingCountWrites tells the goroutine StartPendingCountWrites
// started, if it did, to finish, and waits for its write under way to end
// and for it to return. It then writes the pending access times and counts
// left, once, so that they are written before the database closes. It waits
// for all of this at most until ctx ends, then logs the counts not written
// and returns an error.
func (c *Cache) StopPendingCountWrites(ctx context.Context) error {
	var err error

	if c.pendingCountsCancel != nil {
		c.pendingCountsCancel()

		select {
		case <-c.pendingCountsDone:
		case <-ctx.Done():
			// The write given up on here may not reach the database:
			// database/sql closes without waiting for a query, and a
			// write still waiting for the connection fails when the
			// database closes.
			err = fmt.Errorf("pending counts still being written: %w", ctx.Err())
		}
	}

	if err == nil {
		err = errors.Join(c.writePendingAccessTimes(ctx), c.writePendingCounts(ctx))
	}

	if err != nil {
		c.log.Error("counts not written at shutdown",
			"hits", c.pendingHits.Load(),
			"misses", c.pendingMisses.Load(),
			"upstream_fetches", c.pendingUpstreamFetches.Load(),
			"upstream_fetch_bytes", c.pendingUpstreamFetchBytes.Load(),
			"transforms", c.pendingTransforms.Load(),
			"error", err,
		)

		return err
	}

	return nil
}

// pendingCountWriteLoop is the body of the goroutine
// StartPendingCountWrites starts. Once ctx is cancelled it starts no write
// and returns, but does not cut off a write under way then. Counts whose
// write fails stay pending, for the next write; access times whose write
// fails are dropped.
func (c *Cache) pendingCountWriteLoop(ctx context.Context) {
	defer close(c.pendingCountsDone)

	ticker := time.NewTicker(pendingWriteInterval)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
		}

		// select picks at random when ctx is cancelled and a tick waits
		// too.
		if ctx.Err() != nil {
			return
		}

		err := c.writePendingAccessTimes(context.WithoutCancel(ctx))
		if err != nil {
			c.log.Warn("failed to write access times", "error", err)
		}

		err = c.writePendingCounts(context.WithoutCancel(ctx))
		if err != nil {
			c.log.Warn("failed to write pending counts", "error", err)
		}
	}
}

// writePendingAccessTimes takes out the variants waiting in
// pendingAccessTimes and sets their last_accessed_at to now in one
// transaction, once for each variant however often it was served. Nothing
// else takes from pendingAccessTimes while it runs: only the goroutine
// StartPendingCountWrites starts calls it, or StopPendingCountWrites once
// that goroutine has returned.
func (c *Cache) writePendingAccessTimes(ctx context.Context) error {
	served := make(map[imagerequest.VariantKey]struct{})

	for range len(c.pendingAccessTimes) {
		served[<-c.pendingAccessTimes] = struct{}{}
	}

	if len(served) == 0 {
		return nil
	}

	tx, err := c.db.BeginTx(ctx, nil)
	if err != nil {
		return fmt.Errorf("failed to write %d access times: %w", len(served), err)
	}

	// Does nothing once the transaction is committed
	defer func() { _ = tx.Rollback() }()

	for cacheKey := range served {
		_, err = tx.ExecContext(ctx, `
			UPDATE variant_content SET last_accessed_at = CURRENT_TIMESTAMP
			WHERE cache_key = ?
		`, string(cacheKey))
		if err != nil {
			return fmt.Errorf("failed to write %d access times: %w", len(served), err)
		}
	}

	err = tx.Commit()
	if err != nil {
		return fmt.Errorf("failed to write %d access times: %w", len(served), err)
	}

	return nil
}

// writePendingCounts adds the pending counts to the cache_stats row in one
// UPDATE, then takes what it wrote out of them, leaving any added
// meanwhile. When the UPDATE fails, they all stay pending.
func (c *Cache) writePendingCounts(ctx context.Context) error {
	c.pendingCountsWriteMutex.Lock()
	defer c.pendingCountsWriteMutex.Unlock()

	hits := c.pendingHits.Load()
	misses := c.pendingMisses.Load()
	upstreamFetches := c.pendingUpstreamFetches.Load()
	upstreamFetchBytes := c.pendingUpstreamFetchBytes.Load()
	transforms := c.pendingTransforms.Load()

	if hits+misses+upstreamFetches+upstreamFetchBytes+transforms == 0 {
		return nil
	}

	_, err := c.db.ExecContext(ctx, `
		UPDATE cache_stats
		SET hit_count = hit_count + ?,
		    miss_count = miss_count + ?,
		    upstream_fetch_count = upstream_fetch_count + ?,
		    upstream_fetch_bytes = upstream_fetch_bytes + ?,
		    transform_count = transform_count + ?,
		    last_updated_at = CURRENT_TIMESTAMP
		WHERE id = 1
	`, hits, misses, upstreamFetches, upstreamFetchBytes, transforms)
	if err != nil {
		return fmt.Errorf("failed to write pending counts: %w", err)
	}

	c.pendingHits.Add(-hits)
	c.pendingMisses.Add(-misses)
	c.pendingUpstreamFetches.Add(-upstreamFetches)
	c.pendingUpstreamFetchBytes.Add(-upstreamFetchBytes)
	c.pendingTransforms.Add(-transforms)

	return nil
}