Compare commits
2
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
4d8d477ae0 | ||
|
|
b4dee48401 |
+178
-29
@@ -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
|
||||
@@ -1704,45 +1716,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,
|
||||
@@ -1750,17 +1900,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)
|
||||
}
|
||||
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -20,6 +20,8 @@ const (
|
||||
debugInterval = 60 * time.Second
|
||||
// bytesPerMB is bytes per megabyte
|
||||
bytesPerMB = 1024 * 1024
|
||||
// failureExitCode is the exit status when RouteWatch stops on an error
|
||||
failureExitCode = 1
|
||||
)
|
||||
|
||||
// logDebugStats logs goroutine count and memory usage
|
||||
@@ -83,6 +85,9 @@ func CLIEntry() {
|
||||
go func() {
|
||||
if err := rw.Run(ctx); err != nil {
|
||||
logger.Error("RouteWatch error", "error", err)
|
||||
if err := shutdowner.Shutdown(fx.ExitCode(failureExitCode)); err != nil {
|
||||
logger.Error("Failed to shutdown gracefully", "error", err)
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
|
||||
+17
-60
@@ -201,16 +201,12 @@ func (s *Server) handleStatusJSON() http.HandlerFunc {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
s.logger.Error("Database stats timeout in status.json")
|
||||
writeJSONError(w, http.StatusGatewayTimeout, requestTimeoutMessage)
|
||||
writeJSONError(w, http.StatusRequestTimeout, "Database timeout")
|
||||
|
||||
return
|
||||
case err := <-errChan:
|
||||
s.logger.Error("Failed to get database stats", "error", err)
|
||||
if deadlinePassed(ctx) {
|
||||
writeJSONError(w, http.StatusGatewayTimeout, requestTimeoutMessage)
|
||||
} else {
|
||||
writeJSONError(w, http.StatusInternalServerError, err.Error())
|
||||
}
|
||||
writeJSONError(w, http.StatusInternalServerError, err.Error())
|
||||
|
||||
return
|
||||
case dbStats = <-statsChan:
|
||||
@@ -390,7 +386,7 @@ func (s *Server) handleStats() http.HandlerFunc {
|
||||
// Check if context is already cancelled
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
writeJSONError(w, http.StatusGatewayTimeout, requestTimeoutMessage)
|
||||
http.Error(w, "Request timeout", http.StatusRequestTimeout)
|
||||
|
||||
return
|
||||
default:
|
||||
@@ -419,16 +415,11 @@ func (s *Server) handleStats() http.HandlerFunc {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
s.logger.Error("Database stats timeout")
|
||||
writeJSONError(w, http.StatusGatewayTimeout, requestTimeoutMessage)
|
||||
|
||||
// Don't write response here - timeout middleware already handles it
|
||||
return
|
||||
case err := <-errChan:
|
||||
s.logger.Error("Failed to get database stats", "error", err)
|
||||
if deadlinePassed(ctx) {
|
||||
writeJSONError(w, http.StatusGatewayTimeout, requestTimeoutMessage)
|
||||
} else {
|
||||
writeJSONError(w, http.StatusInternalServerError, err.Error())
|
||||
}
|
||||
writeJSONError(w, http.StatusInternalServerError, err.Error())
|
||||
|
||||
return
|
||||
case dbStats = <-statsChan:
|
||||
@@ -631,13 +622,6 @@ func (s *Server) handleIPInfo() http.HandlerFunc {
|
||||
}
|
||||
}
|
||||
|
||||
// The time limit passed while resolving host names.
|
||||
if deadlinePassed(ctx) {
|
||||
writeJSONError(w, http.StatusGatewayTimeout, requestTimeoutMessage)
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
if len(ipsToLookup) == 0 {
|
||||
writeJSONError(w, http.StatusBadRequest, "No valid IPs or hostnames provided")
|
||||
|
||||
@@ -680,13 +664,6 @@ func (s *Server) handleIPInfo() http.HandlerFunc {
|
||||
}
|
||||
}
|
||||
|
||||
// The time limit passed while looking up the IPs.
|
||||
if deadlinePassed(ctx) {
|
||||
writeJSONError(w, http.StatusGatewayTimeout, requestTimeoutMessage)
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
// Return response (even if no results, include errors)
|
||||
if len(response.Results) == 0 && len(response.Errors) > 0 {
|
||||
writeJSONError(w, http.StatusNotFound, "No routes found: "+response.Errors[0])
|
||||
@@ -726,12 +703,9 @@ func (s *Server) handleASDetailJSON() http.HandlerFunc {
|
||||
|
||||
asInfo, prefixes, err := s.db.GetASDetailsContext(r.Context(), asn)
|
||||
if err != nil {
|
||||
switch {
|
||||
case errors.Is(err, database.ErrNoRoute):
|
||||
if errors.Is(err, database.ErrNoRoute) {
|
||||
writeJSONError(w, http.StatusNotFound, err.Error())
|
||||
case deadlinePassed(r.Context()):
|
||||
writeJSONError(w, http.StatusGatewayTimeout, requestTimeoutMessage)
|
||||
default:
|
||||
} else {
|
||||
writeJSONError(w, http.StatusInternalServerError, err.Error())
|
||||
}
|
||||
|
||||
@@ -787,12 +761,9 @@ func (s *Server) handlePrefixDetailJSON() http.HandlerFunc {
|
||||
|
||||
routes, err := s.db.GetPrefixDetailsContext(r.Context(), prefix)
|
||||
if err != nil {
|
||||
switch {
|
||||
case errors.Is(err, database.ErrNoRoute):
|
||||
if errors.Is(err, database.ErrNoRoute) {
|
||||
writeJSONError(w, http.StatusNotFound, err.Error())
|
||||
case deadlinePassed(r.Context()):
|
||||
writeJSONError(w, http.StatusGatewayTimeout, requestTimeoutMessage)
|
||||
default:
|
||||
} else {
|
||||
writeJSONError(w, http.StatusInternalServerError, err.Error())
|
||||
}
|
||||
|
||||
@@ -832,12 +803,9 @@ func (s *Server) handleASDetail() http.HandlerFunc {
|
||||
|
||||
asInfo, prefixes, err := s.db.GetASDetailsContext(r.Context(), asn)
|
||||
if err != nil {
|
||||
switch {
|
||||
case errors.Is(err, database.ErrNoRoute):
|
||||
if errors.Is(err, database.ErrNoRoute) {
|
||||
http.Error(w, "AS not found", http.StatusNotFound)
|
||||
case deadlinePassed(r.Context()):
|
||||
http.Error(w, requestTimeoutMessage, http.StatusGatewayTimeout)
|
||||
default:
|
||||
} else {
|
||||
s.logger.Error("Failed to get AS details", "error", err)
|
||||
http.Error(w, "Internal server error", http.StatusInternalServerError)
|
||||
}
|
||||
@@ -964,12 +932,9 @@ func (s *Server) handlePrefixDetail() http.HandlerFunc {
|
||||
|
||||
routes, err := s.db.GetPrefixDetailsContext(r.Context(), prefix)
|
||||
if err != nil {
|
||||
switch {
|
||||
case errors.Is(err, database.ErrNoRoute):
|
||||
if errors.Is(err, database.ErrNoRoute) {
|
||||
http.Error(w, "Prefix not found", http.StatusNotFound)
|
||||
case deadlinePassed(r.Context()):
|
||||
http.Error(w, requestTimeoutMessage, http.StatusGatewayTimeout)
|
||||
default:
|
||||
} else {
|
||||
s.logger.Error("Failed to get prefix details", "error", err)
|
||||
http.Error(w, "Internal server error", http.StatusInternalServerError)
|
||||
}
|
||||
@@ -1123,12 +1088,8 @@ func (s *Server) handlePrefixLength() http.HandlerFunc {
|
||||
const maxPrefixes = 500
|
||||
prefixes, err := s.db.GetRandomPrefixesByLengthContext(r.Context(), maskLength, ipVersion, maxPrefixes)
|
||||
if err != nil {
|
||||
if deadlinePassed(r.Context()) {
|
||||
http.Error(w, requestTimeoutMessage, http.StatusGatewayTimeout)
|
||||
} else {
|
||||
s.logger.Error("Failed to get prefixes by length", "error", err)
|
||||
http.Error(w, "Internal server error", http.StatusInternalServerError)
|
||||
}
|
||||
s.logger.Error("Failed to get prefixes by length", "error", err)
|
||||
http.Error(w, "Internal server error", http.StatusInternalServerError)
|
||||
|
||||
return
|
||||
}
|
||||
@@ -1218,12 +1179,8 @@ func (s *Server) handlePrefixLength6() http.HandlerFunc {
|
||||
const maxPrefixes = 500
|
||||
prefixes, err := s.db.GetRandomPrefixesByLengthContext(r.Context(), maskLength, ipVersion, maxPrefixes)
|
||||
if err != nil {
|
||||
if deadlinePassed(r.Context()) {
|
||||
http.Error(w, requestTimeoutMessage, http.StatusGatewayTimeout)
|
||||
} else {
|
||||
s.logger.Error("Failed to get prefixes by length", "error", err)
|
||||
http.Error(w, "Internal server error", http.StatusInternalServerError)
|
||||
}
|
||||
s.logger.Error("Failed to get prefixes by length", "error", err)
|
||||
http.Error(w, "Internal server error", http.StatusInternalServerError)
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
@@ -200,179 +200,6 @@ func TestPrefixDetailHandlersParsePrefix(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// shortRequestDeadline is the deadline the tests below give a request. The
|
||||
// server's 30-second limit and the stats handlers' 4-second limit inherit it,
|
||||
// so the tests do not wait for either.
|
||||
const shortRequestDeadline = 50 * time.Millisecond
|
||||
|
||||
// TestRequestPastTimeLimitAnswers504 sends a request through the server's
|
||||
// router to a page and to a JSON route whose handlers wait for the time limit
|
||||
// to pass and return without answering. The page answers 504; the JSON route
|
||||
// answers 504 with the JSON error body. Before the fix both answered 408,
|
||||
// written by the time limit middleware from its own goroutine while the
|
||||
// handler could still be writing.
|
||||
func TestRequestPastTimeLimitAnswers504(t *testing.T) {
|
||||
s := New(nil, streamer.New(logger.New(), metrics.New()), logger.New(), &config.Config{})
|
||||
|
||||
waitForTimeLimit := func(_ http.ResponseWriter, r *http.Request) {
|
||||
<-r.Context().Done()
|
||||
}
|
||||
// Registered the way setupRoutes registers pages and JSON routes.
|
||||
s.router.Get("/slow-page", waitForTimeLimit)
|
||||
s.router.Get("/slow.json", JSONValidationMiddleware(http.HandlerFunc(waitForTimeLimit)).ServeHTTP)
|
||||
|
||||
tests := []struct {
|
||||
path string
|
||||
wantJSONBody bool
|
||||
}{
|
||||
{path: "/slow-page"},
|
||||
{path: "/slow.json", wantJSONBody: true},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), shortRequestDeadline)
|
||||
rec := httptest.NewRecorder()
|
||||
s.router.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, tt.path, nil).WithContext(ctx))
|
||||
cancel()
|
||||
|
||||
if rec.Code != http.StatusGatewayTimeout {
|
||||
t.Errorf("%s: status %d, want %d; body %s", tt.path, rec.Code, http.StatusGatewayTimeout, rec.Body)
|
||||
}
|
||||
if tt.wantJSONBody {
|
||||
checkJSONErrorCode(t, tt.path, rec, http.StatusGatewayTimeout)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestStatsHandlersAnswer504WhenDatabaseReadRunsOut calls each stats handler
|
||||
// with a store that does not answer before the handler's 4-second limit runs
|
||||
// out, and checks that the handler itself answers 504 with the JSON error
|
||||
// body. Before the fix status.json answered 408, and /api/v1/stats wrote
|
||||
// nothing, which the JSON check on its route turned into 500.
|
||||
func TestStatsHandlersAnswer504WhenDatabaseReadRunsOut(t *testing.T) {
|
||||
release := make(chan struct{})
|
||||
defer close(release)
|
||||
db := blockingStatsDB{release: release}
|
||||
s := New(db, streamer.New(logger.New(), metrics.New()), logger.New(), &config.Config{})
|
||||
|
||||
handlers := map[string]http.HandlerFunc{
|
||||
"status.json": s.handleStatusJSON(),
|
||||
"stats": s.handleStats(),
|
||||
}
|
||||
|
||||
for name, handler := range handlers {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), shortRequestDeadline)
|
||||
rec := httptest.NewRecorder()
|
||||
handler(rec, httptest.NewRequest(http.MethodGet, "/", nil).WithContext(ctx))
|
||||
cancel()
|
||||
|
||||
if rec.Code != http.StatusGatewayTimeout {
|
||||
t.Errorf("%s: status %d, want %d; body %s", name, rec.Code, http.StatusGatewayTimeout, rec.Body)
|
||||
}
|
||||
checkJSONErrorCode(t, name, rec, http.StatusGatewayTimeout)
|
||||
}
|
||||
}
|
||||
|
||||
// slowDB embeds database.Store (left nil) and overrides the reads the routes
|
||||
// in TestRoutesAnswer504WhenDatabaseReadIsCutOff make. Each read waits until
|
||||
// its context ends and returns the context's error, as a slow database read
|
||||
// cut off by a time limit does.
|
||||
type slowDB struct {
|
||||
database.Store
|
||||
}
|
||||
|
||||
func (slowDB) GetStatsContext(ctx context.Context) (database.Stats, error) {
|
||||
<-ctx.Done()
|
||||
|
||||
return database.Stats{}, ctx.Err()
|
||||
}
|
||||
|
||||
func (slowDB) GetIPInfoContext(ctx context.Context, _ string) (*database.IPInfo, error) {
|
||||
<-ctx.Done()
|
||||
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
|
||||
func (slowDB) GetASDetailsContext(ctx context.Context, _ int) (*database.ASN, []database.LiveRoute, error) {
|
||||
<-ctx.Done()
|
||||
|
||||
return nil, nil, ctx.Err()
|
||||
}
|
||||
|
||||
func (slowDB) GetPrefixDetailsContext(ctx context.Context, _ string) ([]database.LiveRoute, error) {
|
||||
<-ctx.Done()
|
||||
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
|
||||
func (slowDB) GetRandomPrefixesByLengthContext(ctx context.Context, _, _, _ int) ([]database.LiveRoute, error) {
|
||||
<-ctx.Done()
|
||||
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
|
||||
// TestRoutesAnswer504WhenDatabaseReadIsCutOff sends a request through the
|
||||
// server's router for each route that reads the database, with a store whose
|
||||
// reads fail when the request's time limit passes. Each route answers 504, the
|
||||
// JSON routes with the JSON error body. Before the fix they answered 408 or,
|
||||
// when the handler wrote first, its own error: 500, or 404 on the IP lookups.
|
||||
// The health check is left out: it answers 503 when its own 2-second ping runs
|
||||
// out.
|
||||
func TestRoutesAnswer504WhenDatabaseReadIsCutOff(t *testing.T) {
|
||||
s := New(slowDB{}, streamer.New(logger.New(), metrics.New()), logger.New(), &config.Config{})
|
||||
|
||||
tests := []struct {
|
||||
path string
|
||||
wantJSONBody bool
|
||||
}{
|
||||
{path: "/as/64500"},
|
||||
{path: "/prefix/192.0.2.0/24"},
|
||||
{path: "/prefixlength/24"},
|
||||
{path: "/prefixlength6/48"},
|
||||
{path: "/status.json", wantJSONBody: true},
|
||||
{path: "/ip/192.0.2.1", wantJSONBody: true},
|
||||
{path: "/api/v1/stats", wantJSONBody: true},
|
||||
{path: "/api/v1/ip/192.0.2.1", wantJSONBody: true},
|
||||
{path: "/api/v1/as/64500", wantJSONBody: true},
|
||||
{path: "/api/v1/prefix/192.0.2.0/24", wantJSONBody: true},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), shortRequestDeadline)
|
||||
rec := httptest.NewRecorder()
|
||||
s.router.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, tt.path, nil).WithContext(ctx))
|
||||
cancel()
|
||||
|
||||
if rec.Code != http.StatusGatewayTimeout {
|
||||
t.Errorf("%s: status %d, want %d; body %s", tt.path, rec.Code, http.StatusGatewayTimeout, rec.Body)
|
||||
}
|
||||
if tt.wantJSONBody {
|
||||
checkJSONErrorCode(t, tt.path, rec, http.StatusGatewayTimeout)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// checkJSONErrorCode checks that rec holds the JSON error body with the given
|
||||
// code.
|
||||
func checkJSONErrorCode(t *testing.T, name string, rec *httptest.ResponseRecorder, want int) {
|
||||
t.Helper()
|
||||
|
||||
var body struct {
|
||||
Status string `json:"status"`
|
||||
Error struct {
|
||||
Code int `json:"code"`
|
||||
} `json:"error"`
|
||||
}
|
||||
if err := json.Unmarshal(rec.Body.Bytes(), &body); err != nil {
|
||||
t.Errorf("%s: decoding the answer %q: %v", name, rec.Body, err)
|
||||
|
||||
return
|
||||
}
|
||||
if body.Status != "error" || body.Error.Code != want {
|
||||
t.Errorf("%s: answer %s, want status \"error\" and error code %d", name, rec.Body, want)
|
||||
}
|
||||
}
|
||||
|
||||
// settledGoroutineCount lets transient goroutines finish, then reports the
|
||||
// current count.
|
||||
func settledGoroutineCount() int {
|
||||
|
||||
+131
-17
@@ -4,7 +4,6 @@ import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
@@ -110,24 +109,142 @@ func JSONResponseMiddleware(next http.Handler) http.Handler {
|
||||
})
|
||||
}
|
||||
|
||||
// requestTimeoutMessage is the error message sent with 504 Gateway Timeout.
|
||||
const requestTimeoutMessage = "Request timeout"
|
||||
// timeoutWriter wraps ResponseWriter to prevent concurrent writes after timeout
|
||||
type timeoutWriter struct {
|
||||
http.ResponseWriter
|
||||
mu sync.Mutex
|
||||
written bool
|
||||
header http.Header // cached header to prevent concurrent access
|
||||
}
|
||||
|
||||
// deadlinePassed reports whether ctx ended because its deadline passed. When
|
||||
// a handler's database read or other lookup fails, it uses this to tell a
|
||||
// request cut off by its time limit, which answers 504, from other failures.
|
||||
func deadlinePassed(ctx context.Context) bool {
|
||||
return errors.Is(ctx.Err(), context.DeadlineExceeded)
|
||||
func (tw *timeoutWriter) Write(b []byte) (int, error) {
|
||||
tw.mu.Lock()
|
||||
defer tw.mu.Unlock()
|
||||
|
||||
if tw.written {
|
||||
return 0, nil // Discard writes after timeout
|
||||
}
|
||||
|
||||
return tw.ResponseWriter.Write(b)
|
||||
}
|
||||
|
||||
func (tw *timeoutWriter) WriteHeader(statusCode int) {
|
||||
tw.mu.Lock()
|
||||
defer tw.mu.Unlock()
|
||||
|
||||
if tw.written {
|
||||
return // Discard writes after timeout
|
||||
}
|
||||
|
||||
tw.ResponseWriter.WriteHeader(statusCode)
|
||||
}
|
||||
|
||||
func (tw *timeoutWriter) Header() http.Header {
|
||||
tw.mu.Lock()
|
||||
defer tw.mu.Unlock()
|
||||
|
||||
if tw.written {
|
||||
// Return a copy to prevent modifications after timeout
|
||||
if tw.header == nil {
|
||||
tw.header = make(http.Header)
|
||||
}
|
||||
|
||||
return tw.header
|
||||
}
|
||||
|
||||
return tw.ResponseWriter.Header()
|
||||
}
|
||||
|
||||
func (tw *timeoutWriter) markWritten() {
|
||||
tw.mu.Lock()
|
||||
defer tw.mu.Unlock()
|
||||
tw.written = true
|
||||
}
|
||||
|
||||
// TimeoutMiddleware creates an HTTP middleware that enforces a request timeout.
|
||||
// If the handler does not complete within the specified duration, the middleware
|
||||
// returns a JSON error response with HTTP status 408 (Request Timeout).
|
||||
//
|
||||
// The timeout parameter specifies the maximum duration allowed for request processing.
|
||||
// The returned middleware handles panics from the wrapped handler by re-panicking
|
||||
// after cleanup, and prevents concurrent writes to the response after timeout occurs.
|
||||
func TimeoutMiddleware(timeout time.Duration, logger *slog.Logger) func(http.Handler) http.Handler {
|
||||
return func(next http.Handler) http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
startTime := time.Now()
|
||||
|
||||
ctx, cancel := context.WithTimeout(r.Context(), timeout)
|
||||
defer cancel()
|
||||
|
||||
tw := &timeoutWriter{
|
||||
ResponseWriter: w,
|
||||
header: make(http.Header),
|
||||
}
|
||||
|
||||
done := make(chan struct{})
|
||||
panicChan := make(chan interface{}, 1)
|
||||
|
||||
go func() {
|
||||
defer func() {
|
||||
if p := recover(); p != nil {
|
||||
panicChan <- p
|
||||
}
|
||||
}()
|
||||
|
||||
next.ServeHTTP(tw, r.WithContext(ctx))
|
||||
close(done)
|
||||
}()
|
||||
|
||||
select {
|
||||
case p := <-panicChan:
|
||||
panic(p)
|
||||
case <-done:
|
||||
return
|
||||
case <-ctx.Done():
|
||||
tw.markWritten() // Prevent the handler from writing after timeout
|
||||
execTime := time.Since(startTime)
|
||||
|
||||
// Log the timeout as a warning
|
||||
logger.Warn("Request timeout",
|
||||
"method", r.Method,
|
||||
"path", r.URL.Path,
|
||||
"duration_ms", execTime.Milliseconds(),
|
||||
"remote_addr", r.RemoteAddr,
|
||||
)
|
||||
|
||||
// Write directly to the underlying writer since we've marked tw as written
|
||||
// This is safe because markWritten() prevents the handler from writing
|
||||
tw.mu.Lock()
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusRequestTimeout)
|
||||
tw.mu.Unlock()
|
||||
|
||||
response := map[string]interface{}{
|
||||
"status": "error",
|
||||
"error": map[string]interface{}{
|
||||
"msg": "Request timeout",
|
||||
"code": http.StatusRequestTimeout,
|
||||
},
|
||||
"@meta": map[string]interface{}{
|
||||
"time_zone": "UTC",
|
||||
"api_version": 1,
|
||||
"execution_time": fmt.Sprintf("%d ms", execTime.Milliseconds()),
|
||||
},
|
||||
}
|
||||
|
||||
_ = json.NewEncoder(w).Encode(response)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// JSONValidationMiddleware is an HTTP middleware that validates JSON API responses.
|
||||
// It ensures that responses with Content-Type "application/json" contain valid JSON.
|
||||
//
|
||||
// If a response is not valid JSON or is empty when JSON is expected, the middleware
|
||||
// returns a properly formatted JSON error response. If the request's time limit
|
||||
// has passed, it answers 504 with the error message "Request timeout". For other
|
||||
// errors, it returns "Internal server error" with status 500 if the original
|
||||
// status was 200.
|
||||
// returns a properly formatted JSON error response. For timeout errors (status 408),
|
||||
// the error message will be "Request timeout". For other errors, it returns
|
||||
// "Internal server error" with status 500 if the original status was 200.
|
||||
//
|
||||
// Non-JSON responses are passed through unchanged.
|
||||
func JSONValidationMiddleware(next http.Handler) http.Handler {
|
||||
@@ -176,14 +293,11 @@ func JSONValidationMiddleware(next http.Handler) http.Handler {
|
||||
if statusCode == http.StatusOK {
|
||||
statusCode = http.StatusInternalServerError
|
||||
}
|
||||
if deadlinePassed(r.Context()) {
|
||||
statusCode = http.StatusGatewayTimeout
|
||||
}
|
||||
w.WriteHeader(statusCode)
|
||||
|
||||
errorMsg := "Internal server error"
|
||||
if statusCode == http.StatusGatewayTimeout {
|
||||
errorMsg = requestTimeoutMessage
|
||||
if statusCode == http.StatusRequestTimeout {
|
||||
errorMsg = "Request timeout"
|
||||
}
|
||||
|
||||
response := map[string]interface{}{
|
||||
|
||||
@@ -16,12 +16,9 @@ func (s *Server) setupRoutes() {
|
||||
r.Use(middleware.RealIP)
|
||||
r.Use(RequestLoggerMiddleware(s.logger.Logger)) // Structured request logging
|
||||
r.Use(middleware.Recoverer)
|
||||
r.Use(JSONResponseMiddleware)
|
||||
// Timeout must come after JSONResponseMiddleware, which sends its answer
|
||||
// when the handler returns and would send a 200 ahead of Timeout's 504 for
|
||||
// a handler that wrote nothing.
|
||||
const requestTimeout = 30 * time.Second // Increased from 8s for slow queries
|
||||
r.Use(middleware.Timeout(requestTimeout))
|
||||
r.Use(TimeoutMiddleware(requestTimeout, s.logger.Logger))
|
||||
r.Use(JSONResponseMiddleware)
|
||||
|
||||
// Routes
|
||||
r.Get("/", s.handleIndex())
|
||||
|
||||
@@ -3,6 +3,8 @@ package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"time"
|
||||
@@ -53,9 +55,12 @@ func New(db database.Store, streamer *streamer.Streamer, logger *logger.Logger,
|
||||
return s
|
||||
}
|
||||
|
||||
// Start starts the HTTP server
|
||||
// Start listens on the configured port and serves HTTP in the background.
|
||||
// It returns an error if it cannot listen, for example because the port is
|
||||
// already in use.
|
||||
func (s *Server) Start() error {
|
||||
const (
|
||||
network = "tcp"
|
||||
readHeaderTimeout = 40 * time.Second
|
||||
readTimeout = 60 * time.Second
|
||||
writeTimeout = 60 * time.Second
|
||||
@@ -73,10 +78,15 @@ func (s *Server) Start() error {
|
||||
|
||||
s.logger.Info("Starting HTTP server", "port", s.port, "addr", s.srv.Addr)
|
||||
|
||||
// Start in goroutine but log when actually listening
|
||||
listener, err := net.Listen(network, s.srv.Addr)
|
||||
if err != nil {
|
||||
return fmt.Errorf("starting HTTP server: %w", err)
|
||||
}
|
||||
|
||||
s.logger.Info("HTTP server listening", "addr", listener.Addr().String())
|
||||
|
||||
go func() {
|
||||
s.logger.Info("HTTP server listening", "addr", s.srv.Addr)
|
||||
if err := s.srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
|
||||
if err := s.srv.Serve(listener); err != nil && err != http.ErrServerClosed {
|
||||
s.logger.Error("HTTP server error", "error", err)
|
||||
}
|
||||
}()
|
||||
|
||||
@@ -0,0 +1,33 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net"
|
||||
"testing"
|
||||
|
||||
"git.eeqj.de/sneak/routewatch/internal/config"
|
||||
"git.eeqj.de/sneak/routewatch/internal/logger"
|
||||
"git.eeqj.de/sneak/routewatch/internal/metrics"
|
||||
"git.eeqj.de/sneak/routewatch/internal/streamer"
|
||||
)
|
||||
|
||||
// TestStartFailsWhenPortIsHeld holds a port with a listener of its own and
|
||||
// checks that Start on that port returns an error. Before the fix Start
|
||||
// listened in a goroutine, only logged the failure and returned nil, so the
|
||||
// daemon kept running without its web server.
|
||||
func TestStartFailsWhenPortIsHeld(t *testing.T) {
|
||||
held, err := net.Listen("tcp", ":0")
|
||||
if err != nil {
|
||||
t.Fatalf("failed to open the test listener: %v", err)
|
||||
}
|
||||
defer func() { _ = held.Close() }()
|
||||
|
||||
port := held.Addr().(*net.TCPAddr).Port
|
||||
log := logger.New()
|
||||
s := New(nil, streamer.New(log, metrics.New()), log, &config.Config{Port: port})
|
||||
defer func() { _ = s.Stop(context.Background()) }()
|
||||
|
||||
if err := s.Start(); err == nil {
|
||||
t.Fatalf("Start returned nil while port %d was held by another listener", port)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user