routewatch/internal/geo/cache.go

package geo

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"maps"
	"time"
)

const (
	// entryLifetime is how long a geojs answer is kept after its fetch.
	entryLifetime = 7 * 24 * time.Hour

	// flushInterval is how often new answers are written to cache.db.
	flushInterval = 5 * time.Second

	// sweepInterval is how often expired answers are deleted.
	sweepInterval = 10 * time.Minute

	// rowColumns are the columns of the geolocations table in the order
	// scanRow reads them.
	rowColumns = "address, country_code, country, continent_code, region, " +
		"city, latitude, longitude, accuracy_radius_km, time_zone, fetched_at"
)

// rowScanner is the Scan method of *sql.Row and *sql.Rows.
type rowScanner interface {
	Scan(dest ...any) error
}

// cached returns the current answer held for address, from memory, then from
// the answers waiting to be written, then from cache.db. It returns nil when
// there is none.
func (l *Locator) cached(ctx context.Context, address string) *Location {
	found, ok := l.entries.Get(address)
	if ok {
		if !l.expired(found) {
			return found
		}

		l.entries.Remove(address)
	}

	l.pendingMu.Lock()
	waiting := l.pending[address]
	l.pendingMu.Unlock()

	if waiting != nil && !l.expired(waiting) {
		l.entries.Add(address, waiting)

		return waiting
	}

	stored, err := l.readRow(ctx, address)
	if err != nil {
		l.logger.Warn("Failed to read a geolocation cache row",
			"address", address, "error", err)

		return nil
	}

	if stored != nil {
		l.entries.Add(address, stored)
	}

	return stored
}

// expired says whether the answer is at the end of its lifetime.
func (l *Locator) expired(answer *Location) bool {
	return !l.now().Before(answer.Updated.Add(entryLifetime))
}

// cutoff returns the fetch time at or before which an answer has expired.
func (l *Locator) cutoff() int64 {
	return l.now().Add(-entryLifetime).Unix()
}

// store keeps an answer in memory and queues it for cache.db. While cache.db
// cannot be written, the queue stops growing at maxEntries answers.
func (l *Locator) store(address string, answer *Location) {
	l.entries.Add(address, answer)

	l.pendingMu.Lock()
	defer l.pendingMu.Unlock()

	if len(l.pending) < l.maxEntries {
		l.pending[address] = answer
	}
}

// keepCache writes new answers and deletes expired ones until ctx ends.
func (l *Locator) keepCache(ctx context.Context) {
	flushes := time.NewTicker(flushInterval)
	defer flushes.Stop()

	sweeps := time.NewTicker(sweepInterval)
	defer sweeps.Stop()

	for {
		select {
		case <-ctx.Done():
			l.flush(context.WithoutCancel(ctx))

			return
		case <-flushes.C:
			l.flush(ctx)
		case <-sweeps.C:
			l.sweep(ctx)
		}
	}
}

// load reads the answers that have not expired into memory, at most maxEntries
// of them, the newest last so that they are the last to be dropped.
func (l *Locator) load(ctx context.Context) error {
	rows, err := l.db.QueryContext(ctx,
		"SELECT "+rowColumns+" FROM (SELECT "+rowColumns+" FROM geolocations "+
			"WHERE fetched_at > ? ORDER BY fetched_at DESC LIMIT ?) "+
			"ORDER BY fetched_at",
		l.cutoff(), l.maxEntries)
	if err != nil {
		return fmt.Errorf("reading the geolocation cache: %w", err)
	}

	defer func() { _ = rows.Close() }()

	for rows.Next() {
		address, answer, err := scanRow(rows)
		if err != nil {
			return fmt.Errorf("reading a geolocation cache row: %w", err)
		}

		l.entries.Add(address, answer)
	}

	err = rows.Err()
	if err != nil {
		return fmt.Errorf("reading the geolocation cache: %w", err)
	}

	return nil
}

// readRow returns the current answer for address in cache.db, or nil when it
// has none.
func (l *Locator) readRow(
	ctx context.Context, address string,
) (*Location, error) {
	row := l.db.QueryRowContext(ctx,
		"SELECT "+rowColumns+" FROM geolocations "+
			"WHERE address = ? AND fetched_at > ?",
		address, l.cutoff())

	_, answer, err := scanRow(row)
	if errors.Is(err, sql.ErrNoRows) {
		return nil, nil //nolint:nilnil // no row is not a failure
	}

	return answer, err
}

// scanRow reads one row of the geolocations table.
func scanRow(row rowScanner) (string, *Location, error) {
	var (
		address string
		fetched int64
		answer  Location
	)

	err := row.Scan(&address, &answer.CountryCode, &answer.Country,
		&answer.ContinentCode, &answer.Region, &answer.City, &answer.Latitude,
		&answer.Longitude, &answer.AccuracyRadiusKm, &answer.TimeZone, &fetched)
	if err != nil {
		return "", nil, err
	}

	answer.Source = sourceGeoJS
	answer.Updated = time.Unix(fetched, 0).UTC()

	return address, &answer, nil
}

// flush writes the pending answers to cache.db. Those it cannot write stay
// pending, unless a newer answer for the address has arrived meanwhile.
func (l *Locator) flush(ctx context.Context) {
	l.pendingMu.Lock()
	batch := l.pending
	l.pending = make(map[string]*Location)
	l.pendingMu.Unlock()

	if len(batch) == 0 {
		return
	}

	err := l.writeRows(ctx, batch)
	if err == nil {
		return
	}

	l.logger.Error("Failed to write the geolocation cache", "error", err)

	l.pendingMu.Lock()
	defer l.pendingMu.Unlock()

	maps.Insert(batch, maps.All(l.pending))
	l.pending = batch
}

// writeRows writes batch to cache.db in one transaction.
func (l *Locator) writeRows(
	ctx context.Context, batch map[string]*Location,
) error {
	tx, err := l.db.BeginTx(ctx, nil)
	if err != nil {
		return err
	}

	defer func() { _ = tx.Rollback() }()

	for address, answer := range batch {
		_, err = tx.ExecContext(ctx,
			"INSERT OR REPLACE INTO geolocations ("+rowColumns+") "+
				"VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
			address, answer.CountryCode, answer.Country, answer.ContinentCode,
			answer.Region, answer.City, answer.Latitude, answer.Longitude,
			answer.AccuracyRadiusKm, answer.TimeZone, answer.Updated.Unix())
		if err != nil {
			return err
		}
	}

	return tx.Commit()
}

// sweep deletes the expired answers from memory and cache.db.
func (l *Locator) sweep(ctx context.Context) {
	for _, address := range l.entries.Keys() {
		answer, ok := l.entries.Peek(address)
		if ok && l.expired(answer) {
			l.entries.Remove(address)
		}
	}

	_, err := l.db.ExecContext(ctx,
		"DELETE FROM geolocations WHERE fetched_at <= ?", l.cutoff())
	if err != nil {
		l.logger.Error("Failed to delete expired geolocation cache rows",
			"error", err)
	}
}