routewatch/internal/database/counts.go

package database

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

// liveCounts holds the running row counts and the prefix distribution 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 query over the tables. The COUNT(*) scans (issue 27) and then
// the prefix distribution query (issue 30) each grew with the database until
// they took the whole request timeout and made /api/v1/stats return 500.
//
// 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
	// The prefix distribution: for each mask length, the number of distinct
	// prefixes that have at least one live route.
	distributionV4 [ipv4Bits + 1]int
	distributionV6 [ipv6Bits + 1]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,
	distributionV4, distributionV6 []PrefixDistribution,
) {
	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

	for _, entry := range distributionV4 {
		addAtMaskLength(c.distributionV4[:], entry.MaskLength, entry.Count)
	}

	for _, entry := range distributionV6 {
		addAtMaskLength(c.distributionV6[:], entry.MaskLength, entry.Count)
	}
}

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

// addToDistribution adds n to the IPv4 and IPv6 prefix distributions once for
// each listed mask length. A write lists the mask lengths of the prefixes it
// gave their first live route with n = 1, and of the prefixes it left with no
// live route with n = -1.
func (c *liveCounts) addToDistribution(maskLengthsV4, maskLengthsV6 []int, n int) {
	c.mu.Lock()
	defer c.mu.Unlock()

	for _, maskLength := range maskLengthsV4 {
		addAtMaskLength(c.distributionV4[:], maskLength, n)
	}

	for _, maskLength := range maskLengthsV6 {
		addAtMaskLength(c.distributionV6[:], maskLength, n)
	}
}

// addAtMaskLength adds n to counts[maskLength]. A mask length the array has no
// entry for is ignored, so a malformed route cannot crash the daemon.
func addAtMaskLength(counts []int, maskLength, n int) {
	if maskLength >= 0 && maskLength < len(counts) {
		counts[maskLength] += n
	}
}

// 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
	s.IPv4PrefixDistribution = distributionList(c.distributionV4[:])
	s.IPv6PrefixDistribution = distributionList(c.distributionV6[:])
}

// distributionList lists the mask lengths that have at least one prefix, in
// ascending order, the way the distribution query returns them.
func distributionList(counts []int) []PrefixDistribution {
	var list []PrefixDistribution

	for maskLength, count := range counts {
		if count > 0 {
			list = append(list,
				PrefixDistribution{MaskLength: maskLength, Count: count})
		}
	}

	return list
}

// 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.
	err := d.db.QueryRowContext(ctx, "SELECT COUNT(*) FROM "+table).Scan(&n)
	if err != nil {
		return 0, fmt.Errorf("failed to count %s: %w", table, err)
	}

	return n, nil
}

// seedCounts reads the current row counts and prefix distribution 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, tableLiveRoutesV4)
	if err != nil {
		return err
	}

	routesV6, err := d.countRows(ctx, tableLiveRoutesV6)
	if err != nil {
		return err
	}

	distributionV4, distributionV6, err := d.GetPrefixDistributionContext(ctx)
	if err != nil {
		return err
	}

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

	return nil
}