routewatch/internal/routewatch/prefixhandler.go

package routewatch

import (
	"context"
	"net/netip"
	"slices"
	"sync"
	"time"

	"git.eeqj.de/sneak/routewatch/internal/database"
	"git.eeqj.de/sneak/routewatch/internal/logger"
	"git.eeqj.de/sneak/routewatch/internal/metrics"
	"git.eeqj.de/sneak/routewatch/internal/ristypes"
	"git.eeqj.de/sneak/routewatch/internal/routetable"
)

const (
	// prefixHandlerQueueSize is the queue capacity for prefix tracking
	// operations, about 4 seconds of feed at peak. The streamer drops updates
	// rather than blocks when a queue is full (a STATE message is held until
	// there is room), so this bounds memory. Batches still flush on a timer
	// (prefixBatchTimeout).
	prefixHandlerQueueSize = 20000

	// prefixBatchSize is the number of prefix updates to batch together
	prefixBatchSize = 25000

	// prefixBatchMaxSize is the most prefix updates the batch holds while the
	// flush goroutine writes the previous one: ten batches, about 15 MiB
	// (README "Memory"). At this size HandleMessage waits until the flush
	// goroutine takes the batch.
	prefixBatchMaxSize = 250000

	// prefixBatchTimeout is the maximum time to wait before flushing a batch
	// DO NOT reduce this timeout - larger batches are more efficient
	prefixBatchTimeout = 1 * time.Second

	// maxTableChangesPerWrite is the most changed prefix and origin AS pairs
	// of the routing table that one flush writes; the rest wait for the next
	// flushes. The change of one pair takes about 430 bytes of memory and 15
	// microseconds to write. So a write holds about 9 MB and takes about a
	// third of the one-second flush interval, which leaves the rest to the
	// batch. The feed peaks at about 5,000 updates a second, each changing at
	// most one pair. So the million pairs that a full-table peer's session
	// going down changes are written in about a minute.
	maxTableChangesPerWrite = 20000
)

// PrefixHandler keeps the routing table in memory current from the feed. It
// applies each update to the table as it handles it. An announcement gives
// the peer a route to the prefix. A withdrawal removes the peer's route, as
// does an announcement whose AS path is empty or holds AS 0 anywhere. A peer
// whose BGP session goes down leaves every prefix.
//
// Each flush writes the table's changes to the database. It also stores each
// prefix the updates name with the time it was first seen.
type PrefixHandler struct {
	db      database.Store
	logger  *logger.Logger
	metrics *metrics.Tracker
	table   *routetable.Table

	// batchMaxSize is prefixBatchMaxSize. Only HandleMessage reads it, so a
	// test may lower it before it hands the handler a message.
	batchMaxSize int
	// tableChangesPerWrite is maxTableChangesPerWrite, except in tests that
	// give a smaller one.
	tableChangesPerWrite int

	// Batching. mu guards batch, writing and writingTableChanges only: the
	// flush goroutine takes them under mu and writes them after releasing
	// mu, so HandleMessage waits for a database write only while batch is at
	// batchMaxSize.
	mu    sync.Mutex
	batch []seenPrefix
	// writing is how many updates the flush goroutine took and has not
	// finished writing.
	writing int
	// writingTableChanges is how many changes the flush goroutine took from
	// the routing table and has not finished writing.
	writingTableChanges int
	// batchTaken is signalled when the flush goroutine takes batch; its L is
	// &mu.
	batchTaken  sync.Cond
	batchFullCh chan struct{} // wakes the flush goroutine when batch is full
	// stopCh hands Stop's context to the flush goroutine. It holds one, so
	// Stop does not wait for a write in progress to hand it over.
	stopCh chan context.Context
	done   chan struct{} // closed when the flush goroutine returns
}

// seenPrefix is what the batch keeps of an update: the prefix it names and its
// time. A flush needs only these to store each new prefix's first-seen time.
type seenPrefix struct {
	prefix    string
	timestamp time.Time
}

// prefixUpdate is one prefix of an update message, as applyToTable applies it
// to the routing table.
type prefixUpdate struct {
	prefix      string
	originASN   int
	peer        string
	messageType string
	timestamp   time.Time
	path        []int
	// endsInSet says the path ends in an AS set of several members, so it
	// gives no origin.
	endsInSet bool
}

// announcesRoute says whether u gives the peer a route to the prefix: an
// announcement whose path gives an origin, or ends in an AS set of several
// members, which gives origin 0 (none). Any other update removes the peer's
// route for the prefix: a withdrawal, or an announcement whose path is empty
// or holds AS 0 anywhere (the decoder empties such a path). Such an
// announcement replaces the peer's route, as any announcement does, with none
// (RFC 7607 treats AS 0 in a path as a withdrawal).
func (u prefixUpdate) announcesRoute() bool {
	return u.messageType == "announcement" && (u.originASN > 0 || u.endsInSet)
}

// NewPrefixHandler creates a new batched prefix tracking handler, which applies
// the updates to table, the routing table read from db at startup.
func NewPrefixHandler(
	db database.Store, logger *logger.Logger, table *routetable.Table,
) *PrefixHandler {
	return newPrefixHandler(db, logger, table, prefixBatchTimeout,
		maxTableChangesPerWrite)
}

// newPrefixHandler is NewPrefixHandler with the flush timer's interval,
// batchTimeout, in place of prefixBatchTimeout, and tableChangesPerWrite in
// place of maxTableChangesPerWrite. A test that the timer must not disturb
// gives an interval longer than the test runs. A test of writes split over
// several flushes gives a small tableChangesPerWrite, so that it exceeds it
// with few changes.
func newPrefixHandler(
	db database.Store, logger *logger.Logger, table *routetable.Table,
	batchTimeout time.Duration, tableChangesPerWrite int,
) *PrefixHandler {
	h := &PrefixHandler{
		db:                   db,
		logger:               logger,
		table:                table,
		batch:                make([]seenPrefix, 0, prefixBatchSize),
		batchMaxSize:         prefixBatchMaxSize,
		tableChangesPerWrite: tableChangesPerWrite,
		batchFullCh:          make(chan struct{}, 1),
		stopCh:               make(chan context.Context, 1),
		done:                 make(chan struct{}),
	}
	h.batchTaken.L = &h.mu

	// Start the flush timer goroutine
	go flushLoop(batchTimeout, h.batchFullCh, h.stopCh, h.done, h.flushBatch)

	return h
}

// SetMetricsTracker sets the metrics tracker for recording route updates
func (h *PrefixHandler) SetMetricsTracker(metrics *metrics.Tracker) {
	h.metrics = metrics
}

// WantsMessage returns true if this handler wants to process messages of the
// given type
func (h *PrefixHandler) WantsMessage(messageType string) bool {
	// UPDATE messages change routes; a STATE message can say a peer's
	// session went down, which removes all of its routes.
	return messageType == ristypes.MessageTypeUpdate ||
		messageType == ristypes.MessageTypeState
}

// QueueCapacity returns the desired queue capacity for this handler
func (h *PrefixHandler) QueueCapacity() int {
	// Batching allows us to use a larger queue
	return prefixHandlerQueueSize
}

// HandleMessage processes a message to track prefix information
func (h *PrefixHandler) HandleMessage(msg *ristypes.RISMessage) {
	if msg.Type == ristypes.MessageTypeState {
		if msg.State == ristypes.StateDown {
			h.table.PeerDown(msg.Peer)
		}

		return
	}

	// Use the pre-parsed timestamp
	timestamp := msg.ParsedTimestamp

	originASN := msg.Path.Origin()

	// The stored and shown AS path lists the members of an AS set that ends
	// the path after the rest of it.
	path := slices.Concat(msg.Path.Sequence, msg.Path.Set)

	h.mu.Lock()
	defer h.mu.Unlock()

	// While the batch is at its maximum size, wait until the flush goroutine
	// takes it. Meanwhile the handler takes nothing off its queue, which
	// fills, and the streamer drops messages for it.
	for len(h.batch) >= h.batchMaxSize {
		h.batchTaken.Wait()
	}

	// Withdrawals are applied to the routing table before announcements, so
	// a message that lists a prefix as both withdrawn and announced keeps the
	// route (RFC 4271 section 4.3). A withdrawal names only the prefix and
	// peer: the message's AS path, and so its origin AS, belongs to the
	// announced prefixes.
	for _, prefix := range msg.Withdrawals {
		update := prefixUpdate{
			prefix:      canonicalPrefix(prefix),
			peer:        msg.Peer,
			messageType: "withdrawal",
			timestamp:   timestamp,
		}
		h.batch = append(h.batch, seenPrefix{update.prefix, update.timestamp})
		h.applyToTable(update)

		// Record withdrawal in metrics
		if h.metrics != nil {
			h.metrics.RecordWithdrawal()
		}
	}

	for _, announcement := range msg.Announcements {
		for _, prefix := range announcement.Prefixes {
			update := prefixUpdate{
				prefix:      canonicalPrefix(prefix),
				originASN:   originASN,
				peer:        msg.Peer,
				messageType: "announcement",
				timestamp:   timestamp,
				path:        path,
				endsInSet:   len(msg.Path.Set) > 0,
			}
			h.batch = append(h.batch, seenPrefix{update.prefix, update.timestamp})
			h.applyToTable(update)

			// Record announcement in metrics
			if h.metrics != nil {
				h.metrics.RecordAnnouncement()
			}
		}
	}

	// A full batch wakes the flush goroutine. If it is writing, the wake-up
	// waits in batchFullCh and the batch keeps growing until the write ends
	// or the batch reaches batchMaxSize.
	if len(h.batch) >= prefixBatchSize {
		select {
		case h.batchFullCh <- struct{}{}:
		default: // already woken
		}
	}
}

// Stop stops the flush goroutine, which makes one last flush before it returns,
// and waits for it until ctx ends. That flush writes at most
// tableChangesPerWrite of the routing table's changes, and the prefixes of
// the updates the handler holds. Once ctx has ended the flush goroutine skips
// that last flush, and a write still running is lost when the process exits. A
// failed write of the table's changes is not retried either. Stop logs how many
// updates and how many of the table's changes are not written.
func (h *PrefixHandler) Stop(ctx context.Context) {
	h.stopCh <- ctx

	select {
	case <-h.done:
	case <-ctx.Done():
	}

	h.mu.Lock()
	unwritten := h.writing + len(h.batch)
	tableChanges := h.writingTableChanges + h.table.ChangeCount()
	h.mu.Unlock()

	if unwritten == 0 && tableChanges == 0 {
		return
	}

	if ctx.Err() != nil {
		h.logger.Error("Stop deadline passed before the prefix handler wrote "+
			"its updates", "updates", unwritten, "table_changes", tableChanges)

		return
	}

	h.logger.Error("The prefix handler stopped before it wrote everything",
		"updates", unwritten, "table_changes", tableChanges)
}

// takeBatch returns the pending updates and leaves an empty batch in their
// place for HandleMessage to fill. The updates count in writing until
// flushBatch has written them.
func (h *PrefixHandler) takeBatch() []seenPrefix {
	h.mu.Lock()
	defer h.mu.Unlock()

	if len(h.batch) == 0 {
		return nil
	}

	batch := h.batch
	h.writing = len(batch)
	h.batch = make([]seenPrefix, 0, prefixBatchSize)
	h.batchTaken.Broadcast()

	return batch
}

// flushBatch writes the routing table's changes, then takes the pending batch
// and stores the prefixes it names, without holding the mutex during the
// writes.
func (h *PrefixHandler) flushBatch() {
	h.writeTableChanges()

	batch := h.takeBatch()
	if len(batch) == 0 {
		return
	}

	startTime := time.Now()

	// Each prefix the batch names, announced or withdrawn, with the earliest
	// timestamp among its updates: the prefix tables store only first_seen,
	// written when the prefix is new and never changed after.
	prefixesToUpdate := make(map[string]time.Time)

	for _, update := range batch {
		seen, exists := prefixesToUpdate[update.prefix]
		if !exists || update.timestamp.Before(seen) {
			prefixesToUpdate[update.prefix] = update.timestamp
		}
	}

	err := h.db.UpdatePrefixesBatch(prefixesToUpdate)
	if err != nil {
		h.logger.Error("Failed to update prefix batch",
			"error", err, "count", len(prefixesToUpdate))
	}

	h.mu.Lock()
	h.writing = 0
	h.mu.Unlock()

	h.logger.Debug("Flushed prefix batch",
		"batch_size", len(batch),
		"unique_prefixes", len(prefixesToUpdate),
		"duration_ms", time.Since(startTime).Milliseconds(),
	)
}

// applyToTable applies update to the routing table in memory: as an
// announcement when it gives the peer a route, otherwise as a withdrawal. A
// prefix that does not parse is logged and left out. Each announcement counts
// in the IPv4 or IPv6 route update rate.
func (h *PrefixHandler) applyToTable(update prefixUpdate) {
	prefix, err := netip.ParsePrefix(update.prefix)
	if err != nil {
		h.logger.Error("Failed to parse prefix",
			"prefix", update.prefix, "error", err)

		return
	}

	if !update.announcesRoute() {
		h.table.Withdraw(prefix, update.peer)

		return
	}

	if h.metrics != nil {
		if prefix.Addr().Is4() {
			h.metrics.RecordIPv4Update()
		} else {
			h.metrics.RecordIPv6Update()
		}
	}

	h.table.Announce(prefix, update.originASN, routetable.Route{
		Peer:   update.peer,
		ASPath: update.path,
		Time:   update.timestamp,
	})
}

// writeTableChanges writes at most tableChangesPerWrite changes of the routing
// table in memory, and every peer the table numbered since the last write.
// When the write fails, they go back to the table, so that a later write writes
// them again, each pair with its content then.
func (h *PrefixHandler) writeTableChanges() {
	h.mu.Lock()
	changes, newPeers := h.table.TakeChanges(h.tableChangesPerWrite)
	h.writingTableChanges = len(changes)
	h.mu.Unlock()

	if len(changes) == 0 && len(newPeers) == 0 {
		return
	}

	err := h.db.WriteRouteChanges(changes, newPeers)

	h.mu.Lock()
	defer h.mu.Unlock()

	h.writingTableChanges = 0

	if err != nil {
		h.logger.Error("Failed to write the routing table's changes; the next "+
			"flush retries", "changes", len(changes), "new_peers", len(newPeers),
			"error", err)

		h.table.PutBackChanges(changes, newPeers)
	}
}

// canonicalPrefix returns prefix in the text form net/netip prints, the form
// the prefix tables store and the IP lookup reads them in. The feed sends IPv6
// announcements compressed ("2001:db8::/32") but IPv6 withdrawals uncompressed
// ("2001:db8:0:0:0:0:0:0/32"). A prefix that does not parse is returned
// unchanged, and applyToTable logs it.
func canonicalPrefix(prefix string) string {
	p, err := netip.ParsePrefix(prefix)
	if err != nil {
		return prefix
	}

	return p.Masked().String()
}