routewatch/internal/database/counts.go

package database

import (
	"context"
	"fmt"
	"sync"
)

// liveCounts holds the running row counts that the stats endpoints report. They
// are seeded once at startup from the tables and then adjusted on every write,
// so a stats read serves them from memory instead of running a COUNT(*) over
// each table. Those scans, once the database passed a few GiB, took the whole
// request timeout and made /api/v1/stats return 500 (issue 27).
//
// A single mutex guards all fields so the stats reader takes a consistent
// snapshot at one instant and writers, which already run under the database
// write lock, adjust the counts after their transaction commits.
type liveCounts struct {
	mu         sync.RWMutex
	asns       int
	prefixesV4 int
	prefixesV6 int
	peerings   int
	peers      int
	routesV4   int
	routesV6   int
}

// seed sets every count to the value read from the tables at startup. It runs
// before any writer, so it needs no coordination with the adjust methods.
func (c *liveCounts) seed(asns, prefixesV4, prefixesV6, peerings, peers, routesV4, routesV6 int) {
	c.mu.Lock()
	defer c.mu.Unlock()

	c.asns = asns
	c.prefixesV4 = prefixesV4
	c.prefixesV6 = prefixesV6
	c.peerings = peerings
	c.peers = peers
	c.routesV4 = routesV4
	c.routesV6 = routesV6
}

// addASNs adds n to the ASN count.
func (c *liveCounts) addASNs(n int) {
	c.mu.Lock()
	c.asns += n
	c.mu.Unlock()
}

// addPrefixes adds to the IPv4 and IPv6 prefix counts.
func (c *liveCounts) addPrefixes(v4, v6 int) {
	c.mu.Lock()
	c.prefixesV4 += v4
	c.prefixesV6 += v6
	c.mu.Unlock()
}

// addPeerings adds n to the peering count.
func (c *liveCounts) addPeerings(n int) {
	c.mu.Lock()
	c.peerings += n
	c.mu.Unlock()
}

// addPeers adds n to the BGP peer count.
func (c *liveCounts) addPeers(n int) {
	c.mu.Lock()
	c.peers += n
	c.mu.Unlock()
}

// addRoutes adds to the IPv4 and IPv6 live-route counts. Deletions pass
// negative values.
func (c *liveCounts) addRoutes(v4, v6 int) {
	c.mu.Lock()
	c.routesV4 += v4
	c.routesV6 += v6
	c.mu.Unlock()
}

// fill copies the counts into a Stats, including the derived totals, under a
// single read lock so the reader sees one consistent snapshot.
func (c *liveCounts) fill(s *Stats) {
	c.mu.RLock()
	defer c.mu.RUnlock()

	s.ASNs = c.asns
	s.IPv4Prefixes = c.prefixesV4
	s.IPv6Prefixes = c.prefixesV6
	s.Prefixes = c.prefixesV4 + c.prefixesV6
	s.Peerings = c.peerings
	s.Peers = c.peers
	s.IPv4Routes = c.routesV4
	s.IPv6Routes = c.routesV6
	s.LiveRoutes = c.routesV4 + c.routesV6
}

// countRows returns the number of rows in the named table. It is used only at
// startup to seed the in-memory counters, so a full COUNT(*) is acceptable.
func (d *Database) countRows(ctx context.Context, table string) (int, error) {
	var n int
	// table is one of a fixed set of literals below, never external input.
	if err := d.db.QueryRowContext(ctx, "SELECT COUNT(*) FROM "+table).Scan(&n); err != nil {
		return 0, fmt.Errorf("failed to count %s: %w", table, err)
	}

	return n, nil
}

// seedCounts reads the current row counts from the tables into the in-memory
// counters. It runs once at startup, before the streamer begins writing.
func (d *Database) seedCounts(ctx context.Context) error {
	asns, err := d.countRows(ctx, "asns")
	if err != nil {
		return err
	}
	prefixesV4, err := d.countRows(ctx, "prefixes_v4")
	if err != nil {
		return err
	}
	prefixesV6, err := d.countRows(ctx, "prefixes_v6")
	if err != nil {
		return err
	}
	peerings, err := d.countRows(ctx, "peerings")
	if err != nil {
		return err
	}
	peers, err := d.countRows(ctx, "bgp_peers")
	if err != nil {
		return err
	}
	routesV4, err := d.countRows(ctx, "live_routes_v4")
	if err != nil {
		return err
	}
	routesV6, err := d.countRows(ctx, "live_routes_v6")
	if err != nil {
		return err
	}

	d.counts.seed(asns, prefixesV4, prefixesV6, peerings, peers, routesV4, routesV6)

	return nil
}