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 the live-route writes. 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, inside the
write.

Model: opus-5-5

internal/database/counts_test.go

@@ -2,6 +2,7 @@ package database
 
 import (
 	"context"
+	"slices"
 	"sync"
 	"testing"
 	"time"
@@ -11,7 +12,8 @@ import (
 	"github.com/google/uuid"
 )
 
-// mkV4Route builds an IPv4 live route with its range columns populated.
+// mkV4Route builds an IPv4 live route with its mask length and range columns
+// taken from the prefix.
 func mkV4Route(t *testing.T, prefix string, asn int, ts time.Time) *LiveRoute {
 	t.Helper()
 
@@ -19,11 +21,15 @@ func mkV4Route(t *testing.T, prefix string, asn int, ts time.Time) *LiveRoute {
 	if err != nil {
 		t.Fatalf("CalculateIPv4Range(%s): %v", prefix, err)
 	}
+	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",
@@ -143,8 +149,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()}
 
@@ -184,6 +190,145 @@ 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 write, through both the
+// batch and the single-route methods, and that after every step it equals what
+// the distribution query reads from the route tables.
+func TestPrefixDistributionTracksWrites(t *testing.T) {
+	cfg := &config.Config{StateDir: t.TempDir()}
+
+	db, err := New(cfg, logger.New())
+	if err != nil {
+		t.Fatalf("failed to create database: %v", err)
+	}
+	defer func() { _ = db.Close() }()
+
+	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)
+
+	all := []PrefixDistribution{{MaskLength: 16, Count: 1}, {MaskLength: 24, Count: 2}}
+	v6Only := []PrefixDistribution{{MaskLength: 32, Count: 1}}
+
+	steps := []struct {
+		name   string
+		write  func() error
+		wantV4 []PrefixDistribution
+		wantV6 []PrefixDistribution
+	}{
+		{
+			name:  "empty database",
+			write: func() error { return nil },
+		},
+		{
+			name:   "new prefixes",
+			write:  func() error { return db.UpsertLiveRouteBatch([]*LiveRoute{shared, other, wide, v6}) },
+			wantV4: all, wantV6: v6Only,
+		},
+		{
+			name:   "re-announcement",
+			write:  func() error { return db.UpsertLiveRouteBatch([]*LiveRoute{shared, other, wide, v6}) },
+			wantV4: all, wantV6: v6Only,
+		},
+		{
+			name:   "second peer announces a prefix that already has a route",
+			write:  func() error { return db.UpsertLiveRoute(sharedSecondPeer) },
+			wantV4: all, wantV6: v6Only,
+		},
+		{
+			name: "withdrawal of a route that is not the last for its prefix",
+			write: func() error {
+				return db.DeleteLiveRouteBatch([]LiveRouteDeletion{
+					{Prefix: shared.Prefix, OriginASN: shared.OriginASN, PeerIP: shared.PeerIP, IPVersion: ipVersionV4},
+				})
+			},
+			wantV4: all, wantV6: v6Only,
+		},
+		{
+			name: "withdrawal of the last route for a prefix",
+			write: func() error {
+				return db.DeleteLiveRoute(sharedSecondPeer.Prefix, sharedSecondPeer.OriginASN, sharedSecondPeer.PeerIP)
+			},
+			wantV4: []PrefixDistribution{{MaskLength: 16, Count: 1}, {MaskLength: 24, Count: 1}},
+			wantV6: v6Only,
+		},
+		{
+			name: "withdrawal of every remaining route, one without an origin ASN",
+			write: func() error {
+				return db.DeleteLiveRouteBatch([]LiveRouteDeletion{
+					{Prefix: other.Prefix, PeerIP: other.PeerIP, IPVersion: ipVersionV4},
+					{Prefix: wide.Prefix, OriginASN: wide.OriginASN, PeerIP: wide.PeerIP, IPVersion: ipVersionV4},
+					{Prefix: v6.Prefix, OriginASN: v6.OriginASN, PeerIP: v6.PeerIP, IPVersion: ipVersionV6},
+				})
+			},
+		},
+	}
+
+	ctx := context.Background()
+	for _, step := range steps {
+		if err := step.write(); err != nil {
+			t.Fatalf("%s: %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)
+	}
+}
+
+// TestStatsDistributionNeedsNoQuery checks that the stats read serves the prefix
+// distribution from memory. The read's deadline has already passed, so any
+// query it ran would fail at once. The distribution query it used to run read
+// every live route and passed the /api/v1/stats deadline on a large database
+// (issue 30).
+func TestStatsDistributionNeedsNoQuery(t *testing.T) {
+	cfg := &config.Config{StateDir: t.TempDir()}
+
+	db, err := New(cfg, logger.New())
+	if err != nil {
+		t.Fatalf("failed to create database: %v", err)
+	}
+	defer func() { _ = db.Close() }()
+
+	ts := time.Date(2026, 1, 2, 3, 4, 5, 0, time.UTC)
+	if err := db.UpsertLiveRouteBatch([]*LiveRoute{
+		mkV4Route(t, "198.51.100.0/24", 64500, ts),
+		mkV6Route("2001:db8::/32", 64501, ts),
+	}); err != nil {
+		t.Fatalf("UpsertLiveRouteBatch: %v", err)
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	stats, err := db.GetStatsContext(ctx)
+	if err != nil {
+		t.Fatalf("GetStatsContext: %v", err)
+	}
+	assertDistribution(t, "IPv4 distribution", stats.IPv4PrefixDistribution,
+		[]PrefixDistribution{{MaskLength: 24, Count: 1}})
+	assertDistribution(t, "IPv6 distribution", stats.IPv6PrefixDistribution,
+		[]PrefixDistribution{{MaskLength: 32, Count: 1}})
 }
 
 // TestStatsRouteTimestamps checks the oldest/newest route timestamps are read
@@ -296,6 +441,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 {
@@ -333,3 +481,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)
+	}
+}