diff --git a/internal/database/database.go b/internal/database/database.go index 5393f74..acaa6c5 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" @@ -53,6 +54,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, which then seldom land on a prefix twice. + 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 @@ -1696,45 +1708,176 @@ 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. 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. A prefix with more routes, or with more rows of other mask +// lengths just before its own, comes up more often: the page needs a varied +// sample, not an even one. 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, @@ -1742,17 +1885,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 82753a9..8ee6463 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,19 @@ 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" +) + // 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 @@ -317,3 +331,141 @@ 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(), "live_routes_v4", 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) + } +}