Serve the /api/v1/stats prefix distribution from memory (closes #30)

The prefix distribution was the last query on the stats path. It read one
index entry per live route on every request, so on a large database it took
the whole 4-second deadline and /api/v1/stats answered 500.

The distribution now lives next to the in-memory counts: seeded once at
startup from the same query, then kept exact by every live-route write. A new
route whose prefix had no live route adds one at its mask length; a delete
that leaves a prefix with no live route takes one away. Each check is one
lookup on the prefix index, made only for new and removed routes, within the
same write. The single-route delete now runs its delete and that lookup in one
transaction, so a failed lookup cannot leave the counts off.

Model: opus-5-5

internal/database/counts_test.go

@@ -2,6 +2,7 @@ package database
 
 import (
 	"context"
+	"slices"
 	"sync"
 	"testing"
 	"time"
@@ -11,14 +12,20 @@ import (
 	"github.com/google/uuid"
 )
 
-// mkV4Route builds an IPv4 live route.
+// mkV4Route builds an IPv4 live route with its mask length taken from the
+// prefix.
 func mkV4Route(t *testing.T, prefix string, asn int, ts time.Time) *LiveRoute {
 	t.Helper()
 
+	maskLength, err := prefixMaskLength(prefix)
+	if err != nil {
+		t.Fatalf("prefixMaskLength(%s): %v", prefix, err)
+	}
+
 	return &LiveRoute{
 		ID:          uuid.New(),
 		Prefix:      prefix,
-		MaskLength:  24,
+		MaskLength:  maskLength,
 		IPVersion:   ipVersionV4,
 		OriginASN:   asn,
 		PeerIP:      "192.0.2.1",
@@ -136,8 +143,8 @@ func TestLiveCountsTrackWritesInRealtime(t *testing.T) {
 }
 
 // TestLiveCountsSeededFromDatabaseAtStartup writes rows, reopens the same
-// database file, and checks the counts come back from the seed scan rather than
-// starting at zero.
+// database file, and checks the counts and the prefix distribution come back
+// from the seed scan rather than starting at zero.
 func TestLiveCountsSeededFromDatabaseAtStartup(t *testing.T) {
 	cfg := &config.Config{StateDir: t.TempDir()}
 
@@ -177,6 +184,171 @@ func TestLiveCountsSeededFromDatabaseAtStartup(t *testing.T) {
 		t.Errorf("seeded routes = (v4 %d, v6 %d, total %d), want (1, 1, 2)",
 			stats.IPv4Routes, stats.IPv6Routes, stats.LiveRoutes)
 	}
+	assertDistribution(t, "seeded IPv4 distribution", stats.IPv4PrefixDistribution,
+		[]PrefixDistribution{{MaskLength: 24, Count: 1}})
+	assertDistribution(t, "seeded IPv6 distribution", stats.IPv6PrefixDistribution,
+		[]PrefixDistribution{{MaskLength: 32, Count: 1}})
+}
+
+// TestPrefixDistributionTracksWrites checks that the prefix distribution the
+// stats read reports stays exact across each kind of live-route write, and that
+// after every step it equals what the distribution query reads from the route
+// tables. The steps run once through the batch methods the prefix handler uses
+// and once through the single-route methods.
+func TestPrefixDistributionTracksWrites(t *testing.T) {
+	ts := time.Date(2026, 1, 2, 3, 4, 5, 0, time.UTC)
+	shared := mkV4Route(t, "198.51.100.0/24", 64500, ts)
+	sharedSecondPeer := mkV4Route(t, "198.51.100.0/24", 64500, ts)
+	sharedSecondPeer.PeerIP = "192.0.2.2"
+	other := mkV4Route(t, "203.0.113.0/24", 64501, ts)
+	wide := mkV4Route(t, "172.16.0.0/16", 64502, ts)
+	v6 := mkV6Route("2001:db8::/32", 64503, ts)
+	v6SecondPeer := mkV6Route("2001:db8::/32", 64503, ts)
+	v6SecondPeer.PeerIP = "2001:db8::2"
+
+	all := []PrefixDistribution{{MaskLength: 16, Count: 1}, {MaskLength: 24, Count: 2}}
+	v6Only := []PrefixDistribution{{MaskLength: 32, Count: 1}}
+
+	steps := []struct {
+		name     string
+		announce []*LiveRoute
+		withdraw []*LiveRoute
+		wantV4   []PrefixDistribution
+		wantV6   []PrefixDistribution
+	}{
+		{
+			name:     "new routes",
+			announce: []*LiveRoute{shared, other, wide, v6},
+			wantV4:   all, wantV6: v6Only,
+		},
+		{
+			name:     "re-announcement",
+			announce: []*LiveRoute{shared, other, wide, v6},
+			wantV4:   all, wantV6: v6Only,
+		},
+		{
+			name:     "second peer announces a prefix that has a live route",
+			announce: []*LiveRoute{sharedSecondPeer, v6SecondPeer},
+			wantV4:   all, wantV6: v6Only,
+		},
+		{
+			name:     "withdrawal of a route that is not the last for its prefix",
+			withdraw: []*LiveRoute{shared, v6},
+			wantV4:   all, wantV6: v6Only,
+		},
+		{
+			name:     "withdrawal of the last route for a prefix",
+			withdraw: []*LiveRoute{sharedSecondPeer, v6SecondPeer},
+			wantV4:   []PrefixDistribution{{MaskLength: 16, Count: 1}, {MaskLength: 24, Count: 1}},
+		},
+		{
+			name:     "withdrawal of every remaining route",
+			withdraw: []*LiveRoute{other, wide},
+		},
+		{
+			name:     "two peers announce a new prefix together",
+			announce: []*LiveRoute{shared, sharedSecondPeer, v6, v6SecondPeer},
+			wantV4:   []PrefixDistribution{{MaskLength: 24, Count: 1}},
+			wantV6:   v6Only,
+		},
+		{
+			name:     "both routes for a prefix withdrawn together",
+			withdraw: []*LiveRoute{shared, sharedSecondPeer, v6, v6SecondPeer},
+		},
+		// The feed often withdraws a route that is not live. That must not take
+		// the prefix out of the distribution. A count wrongly taken below zero is
+		// left out of the answer, so the next step shows it: its announcement
+		// at the same mask length would then not be counted.
+		{
+			name:     "withdrawal of routes that are not live",
+			withdraw: []*LiveRoute{other, v6},
+		},
+		{
+			name:     "announcement after a withdrawal of routes that are not live",
+			announce: []*LiveRoute{other, v6},
+			wantV4:   []PrefixDistribution{{MaskLength: 24, Count: 1}},
+			wantV6:   v6Only,
+		},
+	}
+
+	for _, batch := range []bool{true, false} {
+		name := "single-route writes"
+		if batch {
+			name = "batch writes"
+		}
+
+		t.Run(name, func(t *testing.T) {
+			db, err := New(&config.Config{StateDir: t.TempDir()}, logger.New())
+			if err != nil {
+				t.Fatalf("failed to create database: %v", err)
+			}
+			defer func() { _ = db.Close() }()
+
+			ctx := context.Background()
+			for _, step := range steps {
+				if err := announceRoutes(db, batch, step.announce); err != nil {
+					t.Fatalf("%s: announce: %v", step.name, err)
+				}
+				if err := withdrawRoutes(db, batch, step.withdraw); err != nil {
+					t.Fatalf("%s: withdraw: %v", step.name, err)
+				}
+
+				stats, err := db.GetStatsContext(ctx)
+				if err != nil {
+					t.Fatalf("%s: GetStatsContext: %v", step.name, err)
+				}
+				queryV4, queryV6, err := db.GetPrefixDistributionContext(ctx)
+				if err != nil {
+					t.Fatalf("%s: GetPrefixDistributionContext: %v", step.name, err)
+				}
+
+				assertDistribution(t, step.name+": IPv4 distribution", stats.IPv4PrefixDistribution, step.wantV4)
+				assertDistribution(t, step.name+": IPv6 distribution", stats.IPv6PrefixDistribution, step.wantV6)
+				assertDistribution(t, step.name+": IPv4 distribution query", queryV4, step.wantV4)
+				assertDistribution(t, step.name+": IPv6 distribution query", queryV6, step.wantV6)
+			}
+		})
+	}
+}
+
+// announceRoutes writes routes in one UpsertLiveRouteBatch, or with one
+// UpsertLiveRoute each.
+func announceRoutes(db *Database, batch bool, routes []*LiveRoute) error {
+	if batch {
+		return db.UpsertLiveRouteBatch(routes)
+	}
+	for _, route := range routes {
+		if err := db.UpsertLiveRoute(route); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+// withdrawRoutes removes routes in one DeleteLiveRouteBatch, or with one
+// DeleteLiveRoute each. It names each route by prefix and peer only, with no
+// origin ASN, as a withdrawal from the feed does when its message carries no AS
+// path.
+func withdrawRoutes(db *Database, batch bool, routes []*LiveRoute) error {
+	if !batch {
+		for _, route := range routes {
+			if err := db.DeleteLiveRoute(route.Prefix, 0, route.PeerIP); err != nil {
+				return err
+			}
+		}
+
+		return nil
+	}
+
+	deletions := make([]LiveRouteDeletion, 0, len(routes))
+	for _, route := range routes {
+		deletions = append(deletions, LiveRouteDeletion{
+			Prefix: route.Prefix, PeerIP: route.PeerIP, IPVersion: route.IPVersion,
+		})
+	}
+
+	return db.DeleteLiveRouteBatch(deletions)
 }
 
 // TestStatsRouteTimestamps checks the oldest/newest route timestamps are read
@@ -289,6 +461,9 @@ func TestLiveCountsConcurrentReadWrite(t *testing.T) {
 	if want := writers * 25; stats.IPv6Routes != want {
 		t.Errorf("IPv6Routes = %d, want %d", stats.IPv6Routes, want)
 	}
+	// Every writer announced the same prefix, so it counts once.
+	assertDistribution(t, "IPv6 distribution", stats.IPv6PrefixDistribution,
+		[]PrefixDistribution{{MaskLength: 32, Count: 1}})
 }
 
 type wantCounts struct {
@@ -326,3 +501,11 @@ func assertCounts(t *testing.T, when string, got Stats, want wantCounts) {
 		t.Errorf("%s: LiveRoutes = %d, want %d", when, got.LiveRoutes, want.liveRoutes)
 	}
 }
+
+func assertDistribution(t *testing.T, what string, got, want []PrefixDistribution) {
+	t.Helper()
+
+	if !slices.Equal(got, want) {
+		t.Errorf("%s = %v, want %v", what, got, want)
+	}
+}