diff --git a/internal/database/database.go b/internal/database/database.go index 468bd2b..9da60f1 100644 --- a/internal/database/database.go +++ b/internal/database/database.go @@ -8,6 +8,7 @@ import ( "encoding/json" "errors" "fmt" + "math/rand/v2" "net" "net/netip" "os" @@ -61,6 +62,17 @@ const ( sqliteHardHeapLimitBytes = 1610612736 ) +// Bounds on the work GetRandomPrefixesByLength does to pick its sample. +const ( + // maxPrefixesToShuffle is the most prefixes of one mask length that are + // read in full and shuffled to pick the sample. A mask length with more is + // sampled with random lookups. + maxPrefixesToShuffle = 5000 + // sampleLookupsPerPrefix is the most random lookups made for each prefix + // asked for. + sampleLookupsPerPrefix = 4 +) + // Common errors var ( // ErrInvalidIP is returned when an IP address is malformed @@ -1711,45 +1723,183 @@ func (d *Database) GetRandomPrefixesByLength(maskLength, ipVersion, limit int) ( return d.GetRandomPrefixesByLengthContext(context.Background(), maskLength, ipVersion, limit) } -// GetRandomPrefixesByLengthContext returns a random sample of prefixes with context support +// GetRandomPrefixesByLengthContext returns limit prefixes with the given mask +// length picked at random, or all of them when there are fewer, each once with +// its most recently updated live route, for a limit up to +// maxPrefixesToShuffle. However many routes there are, it reads at most +// maxPrefixesToShuffle+1 prefixes and makes at most sampleLookupsPerPrefix +// lookups for each prefix asked for. func (d *Database) GetRandomPrefixesByLengthContext( ctx context.Context, maskLength, ipVersion, limit int) ([]LiveRoute, error) { - // Select unique prefixes with their most recent route information tableName := "live_routes_v4" if ipVersion == ipVersionV6 { tableName = "live_routes_v6" } + // Reading one prefix more than maxPrefixesToShuffle tells whether these + // are all the prefixes with this mask length. + prefixes, err := d.firstPrefixesByLength(ctx, tableName, maskLength, maxPrefixesToShuffle+1) + if err != nil { + return nil, err + } + rand.Shuffle(len(prefixes), func(i, j int) { + prefixes[i], prefixes[j] = prefixes[j], prefixes[i] + }) + if len(prefixes) > maxPrefixesToShuffle { + prefixes, err = d.samplePrefixesByLength(ctx, tableName, maskLength, limit, prefixes) + if err != nil { + return nil, err + } + } + prefixes = prefixes[:min(limit, len(prefixes))] + + return d.newestRouteOfEachPrefix(ctx, tableName, ipVersion, prefixes) +} + +// firstPrefixesByLength returns the first n distinct prefixes, in text order, +// with the given mask length. It reads only their entries in the mask length +// and prefix index. +func (d *Database) firstPrefixesByLength( + ctx context.Context, tableName string, maskLength, n int) ([]string, error) { //nolint:gosec // Table name is hardcoded based on IP version query := fmt.Sprintf(` - WITH unique_prefixes AS ( - SELECT prefix, MAX(last_updated) as max_updated - FROM %s - WHERE mask_length = ? - GROUP BY prefix - ORDER BY RANDOM() - LIMIT ? - ) - SELECT lr.prefix, lr.mask_length, lr.origin_asn, lr.as_path, - lr.peer_ip, lr.last_updated - FROM %s lr - INNER JOIN unique_prefixes up ON lr.prefix = up.prefix AND lr.last_updated = up.max_updated - WHERE lr.mask_length = ? - `, tableName, tableName) + SELECT DISTINCT prefix FROM %s + WHERE mask_length = ? + ORDER BY prefix + LIMIT ? + `, tableName) - rows, err := d.db.QueryContext(ctx, query, maskLength, limit, maskLength) + rows, err := d.db.QueryContext(ctx, query, maskLength, n) if err != nil { - return nil, fmt.Errorf("failed to query random prefixes: %w", err) + return nil, fmt.Errorf("failed to list prefixes: %w", err) } - defer func() { - _ = rows.Close() - }() + defer func() { _ = rows.Close() }() - var routes []LiveRoute + var prefixes []string for rows.Next() { - var route LiveRoute + var prefix string + if err := rows.Scan(&prefix); err != nil { + return nil, fmt.Errorf("failed to scan prefix: %w", err) + } + prefixes = append(prefixes, prefix) + } + if err := rows.Err(); err != nil { + return nil, fmt.Errorf("failed to list prefixes: %w", err) + } + + return prefixes, nil +} + +// samplePrefixesByLength picks limit distinct prefixes with the given mask +// length at random, in at most sampleLookupsPerPrefix lookups for each prefix +// asked for. ORDER BY RANDOM() would read every route with that mask length. +// Instead each lookup draws a row id between those of the first and last route +// with that mask length and takes the prefix of the first such route at or +// after it. SQLite keeps the row id at the end of every index entry, so the +// mask length index is ordered by mask length, then row id, and the lookup is +// one search of it. +// +// The sample is uneven, as the page needs a varied sample, not an even one: a +// route is picked in proportion to the row ids between it and the previous +// route with that mask length, including the row ids of routes with other mask +// lengths and of withdrawn routes. A route that follows a long run of +// withdrawn routes is picked far more often than one that directly follows +// another route with that mask length. +// +// Should the lookups keep landing on prefixes already picked, the rest are +// taken from shuffled, prefixes of that length in random order. +func (d *Database) samplePrefixesByLength( + ctx context.Context, tableName string, maskLength, limit int, shuffled []string) ([]string, error) { + //nolint:gosec // Table name is hardcoded based on IP version + rangeQuery := fmt.Sprintf(` + SELECT + (SELECT MIN(rowid) FROM %[1]s WHERE mask_length = ?), + (SELECT MAX(rowid) FROM %[1]s WHERE mask_length = ?) + `, tableName) + + var first, last sql.NullInt64 + if err := d.db.QueryRowContext(ctx, rangeQuery, maskLength, maskLength).Scan(&first, &last); err != nil { + return nil, fmt.Errorf("failed to read row id range: %w", err) + } + if !first.Valid { + // Every route with this mask length was withdrawn since they were listed. + return nil, nil + } + + //nolint:gosec // Table name is hardcoded based on IP version + lookupQuery := fmt.Sprintf(` + SELECT prefix FROM %s + WHERE mask_length = ? AND rowid >= ? + ORDER BY rowid + LIMIT 1 + `, tableName) + + stmt, err := d.db.PrepareContext(ctx, lookupQuery) + if err != nil { + return nil, fmt.Errorf("failed to prepare prefix lookup: %w", err) + } + defer func() { _ = stmt.Close() }() + + picked := make(map[string]bool, limit) + prefixes := make([]string, 0, limit) + pick := func(prefix string) { + if !picked[prefix] { + picked[prefix] = true + prefixes = append(prefixes, prefix) + } + } + + for lookups := 0; lookups < limit*sampleLookupsPerPrefix && len(prefixes) < limit; lookups++ { + //nolint:gosec // The sample needs variety, not unpredictability. + rowID := first.Int64 + rand.Int64N(last.Int64-first.Int64+1) + + var prefix string + err := stmt.QueryRowContext(ctx, maskLength, rowID).Scan(&prefix) + if errors.Is(err, sql.ErrNoRows) { + // The routes from rowID on were withdrawn since the range was read. + continue + } + if err != nil { + return nil, fmt.Errorf("failed to look up prefix: %w", err) + } + pick(prefix) + } + + for _, prefix := range shuffled { + if len(prefixes) == limit { + break + } + pick(prefix) + } + + return prefixes, nil +} + +// newestRouteOfEachPrefix returns the most recently updated live route of each +// prefix, one lookup on the prefix index per prefix. A prefix whose routes were +// all withdrawn since it was picked is left out. +func (d *Database) newestRouteOfEachPrefix( + ctx context.Context, tableName string, ipVersion int, prefixes []string) ([]LiveRoute, error) { + //nolint:gosec // Table name is hardcoded based on IP version + query := fmt.Sprintf(` + SELECT prefix, mask_length, origin_asn, as_path, peer_ip, last_updated + FROM %s + WHERE prefix = ? + ORDER BY last_updated DESC + LIMIT 1 + `, tableName) + + stmt, err := d.db.PrepareContext(ctx, query) + if err != nil { + return nil, fmt.Errorf("failed to prepare route query: %w", err) + } + defer func() { _ = stmt.Close() }() + + routes := make([]LiveRoute, 0, len(prefixes)) + for _, prefix := range prefixes { + route := LiveRoute{IPVersion: ipVersion} var pathJSON string - err := rows.Scan( + err := stmt.QueryRowContext(ctx, prefix).Scan( &route.Prefix, &route.MaskLength, &route.OriginASN, @@ -1757,17 +1907,16 @@ func (d *Database) GetRandomPrefixesByLengthContext( &route.PeerIP, &route.LastUpdated, ) - // Set IP version based on which table we queried - route.IPVersion = ipVersion - if err != nil { + if errors.Is(err, sql.ErrNoRows) { continue } + if err != nil { + return nil, fmt.Errorf("failed to query route of %s: %w", prefix, err) + } - // Decode AS path if err := json.Unmarshal([]byte(pathJSON), &route.ASPath); err != nil { route.ASPath = []int{} } - routes = append(routes, route) } diff --git a/internal/database/database_test.go b/internal/database/database_test.go index 1b52c08..93a1c48 100644 --- a/internal/database/database_test.go +++ b/internal/database/database_test.go @@ -6,6 +6,7 @@ import ( "errors" "fmt" "net/netip" + "slices" "sync" "testing" "time" @@ -45,6 +46,32 @@ const ( asnSecondBand = 100 ) +// Parameters for the prefix-length sample tests. +const ( + // sampleLimit is how many prefixes the tests ask the sample for. + sampleLimit = 10 + // sampledPrefixCount is how many prefixes the sampling test stores, more + // than sampleLimit so the sample has to choose among them. + sampledPrefixCount = 30 + // secondV4PeerIP and secondV6PeerIP are the second peer each prefix is + // learned from, besides the one mkV4Route and mkV6Route set. + secondV4PeerIP = "192.0.2.2" + secondV6PeerIP = "2001:db8::2" + // largeSamplePrefixCount is how many prefixes the test of a mask length + // with more than maxPrefixesToShuffle prefixes stores: twice that, so about + // half of them come after the first maxPrefixesToShuffle+1 in text order. + largeSamplePrefixCount = 2 * maxPrefixesToShuffle + // largeSampleLimit is how many prefixes that test asks for. Each random + // lookup lands after the first maxPrefixesToShuffle+1 in text order about + // half the time, so none of them doing so is about as likely as + // largeSampleLimit coin tosses all coming up heads. + largeSampleLimit = 30 + // liveRoutesV4Table and liveRoutesV6Table are the live route tables the + // tests pass to the sampling steps. + liveRoutesV4Table = "live_routes_v4" + liveRoutesV6Table = "live_routes_v6" +) + // TestGetIPInfoFindsMostSpecificLiveRoute stores nested live prefixes for both // families and checks that a lookup returns the most specific one covering the // address, ErrNoRoute when none covers it, and the next less specific prefix @@ -389,3 +416,281 @@ func TestBatchWriteDuringCheckpoint(t *testing.T) { cancel() wg.Wait() } + +// TestGetRandomPrefixesByLengthListsPrefixOnce stores one prefix of each +// family learned from two peers in the same second and checks that the +// prefix-length sample lists each of them once. +func TestGetRandomPrefixesByLengthListsPrefixOnce(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) + v4First := mkV4Route(t, "198.51.100.0/24", 64500, ts) + v4Second := mkV4Route(t, "198.51.100.0/24", 64500, ts) + v4Second.PeerIP = secondV4PeerIP + v6First := mkV6Route("2001:db8::/32", 64500, ts) + v6Second := mkV6Route("2001:db8::/32", 64500, ts) + v6Second.PeerIP = secondV6PeerIP + if err := db.UpsertLiveRouteBatch([]*LiveRoute{v4First, v4Second, v6First, v6Second}); err != nil { + t.Fatalf("UpsertLiveRouteBatch: %v", err) + } + + for _, want := range []*LiveRoute{v4First, v6First} { + routes, err := db.GetRandomPrefixesByLength(want.MaskLength, want.IPVersion, sampleLimit) + if err != nil { + t.Fatalf("GetRandomPrefixesByLength(%d, %d): %v", want.MaskLength, want.IPVersion, err) + } + if len(routes) != 1 || routes[0].Prefix != want.Prefix { + t.Errorf("GetRandomPrefixesByLength(%d, %d) = %+v, want %s listed once", + want.MaskLength, want.IPVersion, routes, want.Prefix) + } + } +} + +// TestGetRandomPrefixesByLengthListsNewestRoute stores a prefix learned from +// two peers at different times, the older route first, and checks that the +// prefix-length sample lists the newer route. +func TestGetRandomPrefixesByLengthListsNewestRoute(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() }() + + older := time.Date(2026, 1, 2, 3, 4, 5, 0, time.UTC) + olderRoute := mkV4Route(t, "198.51.100.0/24", 64500, older) + newerRoute := mkV4Route(t, "198.51.100.0/24", 64500, older.Add(time.Minute)) + newerRoute.PeerIP = secondV4PeerIP + if err := db.UpsertLiveRouteBatch([]*LiveRoute{olderRoute, newerRoute}); err != nil { + t.Fatalf("UpsertLiveRouteBatch: %v", err) + } + + routes, err := db.GetRandomPrefixesByLength(newerRoute.MaskLength, ipVersionV4, sampleLimit) + if err != nil { + t.Fatalf("GetRandomPrefixesByLength: %v", err) + } + if len(routes) != 1 || routes[0].PeerIP != newerRoute.PeerIP { + t.Errorf("GetRandomPrefixesByLength = %+v, want only the newer route, from %s", + routes, newerRoute.PeerIP) + } +} + +// TestGetRandomPrefixesByLengthSamplesManyPrefixes stores more /24 prefixes +// than the sample asks for, each followed by a /16 so that routes of another +// mask length sit between the /24 rows, and each prefix learned from two peers +// in the same second. It checks both the sample the pages get, picked from the +// shuffled prefixes, and the random lookups used for a mask length with more +// than maxPrefixesToShuffle prefixes. +func TestGetRandomPrefixesByLengthSamplesManyPrefixes(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) + var routes []*LiveRoute + var stored []string + for i := range sampledPrefixCount { + prefix := fmt.Sprintf("10.0.%d.0/24", i) + for _, p := range []string{prefix, fmt.Sprintf("172.%d.0.0/16", i)} { + first := mkV4Route(t, p, 64500, ts) + second := mkV4Route(t, p, 64500, ts) + second.PeerIP = secondV4PeerIP + routes = append(routes, first, second) + } + stored = append(stored, prefix) + } + if err := db.UpsertLiveRouteBatch(routes); err != nil { + t.Fatalf("UpsertLiveRouteBatch: %v", err) + } + slices.Sort(stored) + + maskLength := routes[0].MaskLength + sample, err := db.GetRandomPrefixesByLength(maskLength, ipVersionV4, sampleLimit) + if err != nil { + t.Fatalf("GetRandomPrefixesByLength: %v", err) + } + var listed []string + for _, route := range sample { + listed = append(listed, route.Prefix) + } + checkPrefixSample(t, "GetRandomPrefixesByLength", listed, stored) + + looked, err := db.samplePrefixesByLength( + context.Background(), liveRoutesV4Table, maskLength, sampleLimit, nil) + if err != nil { + t.Fatalf("samplePrefixesByLength: %v", err) + } + checkPrefixSample(t, "samplePrefixesByLength", looked, stored) +} + +// checkPrefixSample checks that sample lists sampleLimit prefixes, each once +// and each one of stored, and that it is not just the first sampleLimit of +// stored, which is in text order. +func checkPrefixSample(t *testing.T, name string, sample, stored []string) { + t.Helper() + + if len(sample) != sampleLimit { + t.Errorf("%s listed %d prefixes, want %d", name, len(sample), sampleLimit) + } + listed := make(map[string]bool) + for _, prefix := range sample { + if listed[prefix] || !slices.Contains(stored, prefix) { + t.Errorf("%s lists %s twice or it is not one of the /24s", name, prefix) + } + listed[prefix] = true + } + if slices.Equal(slices.Sorted(slices.Values(sample)), stored[:sampleLimit]) { + t.Errorf("%s lists just the first %d prefixes in text order", name, sampleLimit) + } +} + +// TestGetRandomPrefixesByLengthSamplesPastFirstPrefixes stores more prefixes +// of one mask length than maxPrefixesToShuffle. It checks that +// firstPrefixesByLength reads just the first maxPrefixesToShuffle+1 in text +// order, and that the sample lists the prefixes asked for, each once, some of +// them after those first ones, which only the random lookups reach. +func TestGetRandomPrefixesByLengthSamplesPastFirstPrefixes(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) + var routes []*LiveRoute + var stored []string + for i := range largeSamplePrefixCount { + route := mkV6Route(fmt.Sprintf("2001:%x::/32", i), 64500, ts) + routes = append(routes, route) + stored = append(stored, route.Prefix) + } + if err := db.UpsertLiveRouteBatch(routes); err != nil { + t.Fatalf("UpsertLiveRouteBatch: %v", err) + } + slices.Sort(stored) + firstPrefixes := stored[:maxPrefixesToShuffle+1] + + maskLength := routes[0].MaskLength + read, err := db.firstPrefixesByLength( + context.Background(), liveRoutesV6Table, maskLength, maxPrefixesToShuffle+1) + if err != nil { + t.Fatalf("firstPrefixesByLength: %v", err) + } + if !slices.Equal(read, firstPrefixes) { + t.Errorf("firstPrefixesByLength read %d prefixes, want the first %d in text order", + len(read), len(firstPrefixes)) + } + + sample, err := db.GetRandomPrefixesByLength(maskLength, ipVersionV6, largeSampleLimit) + if err != nil { + t.Fatalf("GetRandomPrefixesByLength: %v", err) + } + if len(sample) != largeSampleLimit { + t.Errorf("GetRandomPrefixesByLength listed %d prefixes, want %d", len(sample), largeSampleLimit) + } + listed := make(map[string]bool) + pastFirstPrefixes := 0 + for _, route := range sample { + if listed[route.Prefix] { + t.Errorf("GetRandomPrefixesByLength lists %s twice", route.Prefix) + } + listed[route.Prefix] = true + if !slices.Contains(firstPrefixes, route.Prefix) { + pastFirstPrefixes++ + } + } + if pastFirstPrefixes == 0 { + t.Errorf("GetRandomPrefixesByLength lists only prefixes among the first %d in text order", + len(firstPrefixes)) + } +} + +// TestSamplePrefixesByLengthFillsUpFromShuffled stores a single prefix, so +// every random lookup lands on it, and checks that the sample still lists +// sampleLimit prefixes: the stored one, then the shuffled ones in order. The +// stored prefix is the second shuffled one, so the fill-up must skip it. +func TestSamplePrefixesByLengthFillsUpFromShuffled(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() }() + + var shuffled []string + for i := range sampledPrefixCount { + shuffled = append(shuffled, fmt.Sprintf("10.0.%d.0/24", i)) + } + route := mkV4Route(t, shuffled[1], 64500, time.Date(2026, 1, 2, 3, 4, 5, 0, time.UTC)) + if err := db.UpsertLiveRouteBatch([]*LiveRoute{route}); err != nil { + t.Fatalf("UpsertLiveRouteBatch: %v", err) + } + + sample, err := db.samplePrefixesByLength( + context.Background(), liveRoutesV4Table, route.MaskLength, sampleLimit, shuffled) + if err != nil { + t.Fatalf("samplePrefixesByLength: %v", err) + } + want := append([]string{shuffled[1], shuffled[0]}, shuffled[2:sampleLimit]...) + if !slices.Equal(sample, want) { + t.Errorf("samplePrefixesByLength = %v, want %v", sample, want) + } +} + +// TestGetRandomPrefixesByLengthSkipsWithdrawnRoutes withdraws a route after +// storing it, as can happen between the steps of GetRandomPrefixesByLength. It +// checks that the random lookups pick nothing for a mask length with no route +// left, and that only the prefixes that still have a route get their newest +// route listed. +func TestGetRandomPrefixesByLengthSkipsWithdrawnRoutes(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) + kept := mkV4Route(t, "198.51.100.0/24", 64500, ts) + withdrawn := mkV4Route(t, "172.16.0.0/16", 64500, ts) + if err := db.UpsertLiveRouteBatch([]*LiveRoute{kept, withdrawn}); err != nil { + t.Fatalf("UpsertLiveRouteBatch: %v", err) + } + if err := db.DeleteLiveRoute(withdrawn.Prefix, withdrawn.OriginASN, withdrawn.PeerIP); err != nil { + t.Fatalf("DeleteLiveRoute: %v", err) + } + + ctx := context.Background() + sample, err := db.samplePrefixesByLength( + ctx, liveRoutesV4Table, withdrawn.MaskLength, sampleLimit, []string{withdrawn.Prefix}) + if err != nil { + t.Fatalf("samplePrefixesByLength: %v", err) + } + if len(sample) != 0 { + t.Errorf("samplePrefixesByLength = %v, want none for a mask length with no route left", sample) + } + + routes, err := db.newestRouteOfEachPrefix( + ctx, liveRoutesV4Table, ipVersionV4, []string{withdrawn.Prefix, kept.Prefix}) + if err != nil { + t.Fatalf("newestRouteOfEachPrefix: %v", err) + } + if len(routes) != 1 || routes[0].Prefix != kept.Prefix { + t.Errorf("newestRouteOfEachPrefix = %+v, want only the route of %s", routes, kept.Prefix) + } +}