pixa/internal/imgcache/pendingcounts.go

package imgcache

import (
	"context"
	"fmt"
	"sync/atomic"
)

// StartPendingCountWrites starts the goroutine that writes the pending
// counts to the database whenever there are some, one UPDATE at a time:
// however many requests pass their deadline before their counts are
// written, at most this one write waits for the database. 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 counts left, 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 under way is left to finish: the database's
			// close waits for a query under way.
			err = fmt.Errorf("pending counts still being written: %w", ctx.Err())
		}
	}

	if err == nil {
		err = 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. It returns when ctx is cancelled, but
// does not cut off a write under way then. A write that fails leaves the
// counts pending, for the next write.
func (c *Cache) pendingCountWriteLoop(ctx context.Context) {
	defer close(c.pendingCountsDone)

	for {
		select {
		case <-ctx.Done():
			return
		case <-c.pendingCountsAdded:
		}

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

// addPendingCount adds n to pendingCount, one of the pending counts, and
// wakes the goroutine that writes them. pendingCountsAdded has capacity one,
// so a wakeup already waiting is enough.
func (c *Cache) addPendingCount(pendingCount *atomic.Int64, n int64) {
	pendingCount.Add(n)

	select {
	case c.pendingCountsAdded <- struct{}{}:
	default:
	}
}

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