routewatch/internal/database/counts_test.go

package database

import (
	"context"
	"slices"
	"sync"
	"testing"
	"time"

	"git.eeqj.de/sneak/routewatch/internal/config"
	"git.eeqj.de/sneak/routewatch/internal/logger"
	"github.com/google/uuid"
)

// 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:  maskLength,
		IPVersion:   ipVersionV4,
		OriginASN:   asn,
		PeerIP:      "192.0.2.1",
		ASPath:      []int{asn},
		NextHop:     "192.0.2.254",
		LastUpdated: ts,
	}
}

// mkV6Route builds an IPv6 live route.
func mkV6Route(prefix string, asn int, ts time.Time) *LiveRoute {
	return &LiveRoute{
		ID:          uuid.New(),
		Prefix:      prefix,
		MaskLength:  32,
		IPVersion:   ipVersionV6,
		OriginASN:   asn,
		PeerIP:      "2001:db8::1",
		ASPath:      []int{asn},
		NextHop:     "2001:db8::ffff",
		LastUpdated: ts,
	}
}

// TestLiveCountsTrackWritesInRealtime checks that the stats counts start at
// zero, reflect each write the moment it commits (no recompute, no timer), do
// not move when a route is merely re-announced, and drop when a route is
// deleted. These counts are what /api/v1/stats reports; before this change the
// endpoint recomputed them with a COUNT(*) over each table on every request.
func TestLiveCountsTrackWritesInRealtime(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() }()

	ctx := context.Background()

	empty, err := db.GetStatsContext(ctx)
	if err != nil {
		t.Fatalf("GetStatsContext on empty database: %v", err)
	}
	if empty.ASNs != 0 || empty.Prefixes != 0 || empty.Peerings != 0 ||
		empty.Peers != 0 || empty.LiveRoutes != 0 {
		t.Fatalf("empty database counts nonzero: %+v", empty)
	}

	ts := time.Date(2026, 1, 2, 3, 4, 5, 0, time.UTC)

	if err := db.GetOrCreateASNBatch(map[int]time.Time{64500: ts, 64501: ts}); err != nil {
		t.Fatalf("GetOrCreateASNBatch: %v", err)
	}
	if err := db.UpdatePrefixesBatch(map[string]time.Time{
		"198.51.100.0/24": ts,
		"2001:db8::/32":   ts,
	}); err != nil {
		t.Fatalf("UpdatePrefixesBatch: %v", err)
	}
	if err := db.UpdatePeerBatch(map[string]PeerUpdate{
		"192.0.2.1": {PeerIP: "192.0.2.1", PeerASN: 64500, MessageType: "UPDATE", Timestamp: ts},
	}); err != nil {
		t.Fatalf("UpdatePeerBatch: %v", err)
	}
	if err := db.RecordPeering(64500, 64501, ts); err != nil {
		t.Fatalf("RecordPeering: %v", err)
	}

	routes := []*LiveRoute{
		mkV4Route(t, "198.51.100.0/24", 64500, ts),
		mkV4Route(t, "203.0.113.0/24", 64501, ts.Add(time.Minute)),
		mkV6Route("2001:db8::/32", 64502, ts.Add(2*time.Minute)),
	}
	if err := db.UpsertLiveRouteBatch(routes); err != nil {
		t.Fatalf("UpsertLiveRouteBatch: %v", err)
	}

	stats, err := db.GetStatsContext(ctx)
	if err != nil {
		t.Fatalf("GetStatsContext: %v", err)
	}
	assertCounts(t, "after inserts", stats, wantCounts{
		asns: 2, prefixes: 2, peerings: 1, peers: 1,
		ipv4Routes: 2, ipv6Routes: 1, liveRoutes: 3,
	})

	// Re-announcing the same routes is an update, not an insert: counts hold.
	if err := db.UpsertLiveRouteBatch(routes); err != nil {
		t.Fatalf("UpsertLiveRouteBatch (re-announce): %v", err)
	}
	stats, err = db.GetStatsContext(ctx)
	if err != nil {
		t.Fatalf("GetStatsContext: %v", err)
	}
	assertCounts(t, "after re-announce", stats, wantCounts{
		asns: 2, prefixes: 2, peerings: 1, peers: 1,
		ipv4Routes: 2, ipv6Routes: 1, liveRoutes: 3,
	})

	// A withdrawal removes one route.
	if err := db.DeleteLiveRouteBatch([]LiveRouteDeletion{
		{Prefix: "203.0.113.0/24", OriginASN: 64501, PeerIP: "192.0.2.1", IPVersion: ipVersionV4},
	}); err != nil {
		t.Fatalf("DeleteLiveRouteBatch: %v", err)
	}
	stats, err = db.GetStatsContext(ctx)
	if err != nil {
		t.Fatalf("GetStatsContext: %v", err)
	}
	assertCounts(t, "after delete", stats, wantCounts{
		asns: 2, prefixes: 2, peerings: 1, peers: 1,
		ipv4Routes: 1, ipv6Routes: 1, liveRoutes: 2,
	})
}

// TestLiveCountsSeededFromDatabaseAtStartup writes rows, reopens the same
// 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()}

	db, err := New(cfg, logger.New())
	if err != nil {
		t.Fatalf("failed to create database: %v", err)
	}

	ts := time.Date(2026, 1, 2, 3, 4, 5, 0, time.UTC)
	if err := db.GetOrCreateASNBatch(map[int]time.Time{64500: ts, 64501: ts, 64502: ts}); err != nil {
		t.Fatalf("GetOrCreateASNBatch: %v", err)
	}
	if err := db.UpsertLiveRouteBatch([]*LiveRoute{
		mkV4Route(t, "198.51.100.0/24", 64500, ts),
		mkV6Route("2001:db8::/32", 64502, ts),
	}); err != nil {
		t.Fatalf("UpsertLiveRouteBatch: %v", err)
	}
	if err := db.Close(); err != nil {
		t.Fatalf("Close: %v", err)
	}

	reopened, err := New(cfg, logger.New())
	if err != nil {
		t.Fatalf("failed to reopen database: %v", err)
	}
	defer func() { _ = reopened.Close() }()

	stats, err := reopened.GetStatsContext(context.Background())
	if err != nil {
		t.Fatalf("GetStatsContext after reopen: %v", err)
	}
	if stats.ASNs != 3 {
		t.Errorf("seeded ASNs = %d, want 3", stats.ASNs)
	}
	if stats.IPv4Routes != 1 || stats.IPv6Routes != 1 || stats.LiveRoutes != 2 {
		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
// from the right rows across both tables and parse into time.Time. The old
// MIN/MAX union query read its result into *time.Time, which the driver could
// not parse, so it logged a warning every call and left both timestamps nil.
func TestStatsRouteTimestamps(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() }()

	ctx := context.Background()

	empty, err := db.GetStatsContext(ctx)
	if err != nil {
		t.Fatalf("GetStatsContext on empty database: %v", err)
	}
	if empty.OldestRoute != nil || empty.NewestRoute != nil {
		t.Fatalf("empty database timestamps = (%v, %v), want (nil, nil)",
			empty.OldestRoute, empty.NewestRoute)
	}

	base := time.Date(2026, 1, 2, 3, 4, 5, 0, time.UTC)
	oldest := base
	newest := base.Add(2 * time.Minute)
	if err := db.UpsertLiveRouteBatch([]*LiveRoute{
		mkV4Route(t, "198.51.100.0/24", 64500, base.Add(time.Minute)),
		mkV4Route(t, "203.0.113.0/24", 64501, oldest),
		mkV6Route("2001:db8::/32", 64502, newest),
	}); err != nil {
		t.Fatalf("UpsertLiveRouteBatch: %v", err)
	}

	stats, err := db.GetStatsContext(ctx)
	if err != nil {
		t.Fatalf("GetStatsContext: %v", err)
	}
	if stats.OldestRoute == nil || !stats.OldestRoute.Equal(oldest) {
		t.Errorf("OldestRoute = %v, want %v", stats.OldestRoute, oldest)
	}
	if stats.NewestRoute == nil || !stats.NewestRoute.Equal(newest) {
		t.Errorf("NewestRoute = %v, want %v", stats.NewestRoute, newest)
	}
}

// TestLiveCountsConcurrentReadWrite runs writers and stats readers at once so
// the race detector proves the counters are safe under concurrent use.
func TestLiveCountsConcurrentReadWrite(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)

	const writers = 4
	var wg sync.WaitGroup

	wg.Add(writers)
	for w := range writers {
		go func(base int) {
			defer wg.Done()
			for i := range 25 {
				asn := 65000 + base*100 + i
				route := mkV6Route("2001:db8::/32", asn, ts)
				route.PeerIP = "2001:db8::" + uuid.NewString()
				if err := db.UpsertLiveRoute(route); err != nil {
					t.Errorf("UpsertLiveRoute: %v", err)

					return
				}
			}
		}(w)
	}

	var readerWG sync.WaitGroup
	readerWG.Add(1)
	stop := make(chan struct{})
	go func() {
		defer readerWG.Done()
		for {
			select {
			case <-stop:
				return
			default:
				if _, err := db.GetStatsContext(context.Background()); err != nil {
					t.Errorf("GetStatsContext: %v", err)

					return
				}
			}
		}
	}()

	wg.Wait()
	close(stop)
	readerWG.Wait()

	stats, err := db.GetStatsContext(context.Background())
	if err != nil {
		t.Fatalf("GetStatsContext: %v", err)
	}
	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 {
	asns       int
	prefixes   int
	peerings   int
	peers      int
	ipv4Routes int
	ipv6Routes int
	liveRoutes int
}

func assertCounts(t *testing.T, when string, got Stats, want wantCounts) {
	t.Helper()

	if got.ASNs != want.asns {
		t.Errorf("%s: ASNs = %d, want %d", when, got.ASNs, want.asns)
	}
	if got.Prefixes != want.prefixes {
		t.Errorf("%s: Prefixes = %d, want %d", when, got.Prefixes, want.prefixes)
	}
	if got.Peerings != want.peerings {
		t.Errorf("%s: Peerings = %d, want %d", when, got.Peerings, want.peerings)
	}
	if got.Peers != want.peers {
		t.Errorf("%s: Peers = %d, want %d", when, got.Peers, want.peers)
	}
	if got.IPv4Routes != want.ipv4Routes {
		t.Errorf("%s: IPv4Routes = %d, want %d", when, got.IPv4Routes, want.ipv4Routes)
	}
	if got.IPv6Routes != want.ipv6Routes {
		t.Errorf("%s: IPv6Routes = %d, want %d", when, got.IPv6Routes, want.ipv6Routes)
	}
	if got.LiveRoutes != want.liveRoutes {
		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)
	}
}