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.
//
// The table's PeerDown applies a session going down. It visits every prefix,
// which takes over a second on a full table. So the handler runs it on a
// goroutine of its own, one peer at a time, and goes on taking messages
// meanwhile. It holds back the peer's later messages until PeerDown returns,
// then applies them in order.
//
// 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: 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

	// Peers whose session went down. peerDownMu lets one PeerDown run at a
	// time, so the handler's own updates wait for the table's lock behind
	// at most one. heldBack, under mu, holds each such peer's later
	// messages, in order, until its PeerDown has returned and they are
	// applied. A peer with no such messages yet has an empty entry.
	// peerDowns counts the peerDown goroutines that have not returned.
	peerDownMu sync.Mutex
	heldBack   map[string][]*ristypes.RISMessage
	peerDowns  sync.WaitGroup
}

// 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{}),
		heldBack:             make(map[string][]*ristypes.RISMessage),
	}
	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. A message
// from a peer whose PeerDown has not returned is held back until it has.
func (h *PrefixHandler) HandleMessage(msg *ristypes.RISMessage) {
	h.mu.Lock()
	defer h.mu.Unlock()

	held, goingDown := h.heldBack[msg.Peer]
	if goingDown {
		h.heldBack[msg.Peer] = append(held, msg)

		return
	}

	h.applyMessage(msg)
}

// Stop waits until ctx ends for every PeerDown the handler started, and for
// the messages it held back meanwhile to be applied. Then it 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, and how many held-back messages are not
// applied. The caller calls Stop once its last HandleMessage has returned.
func (h *PrefixHandler) Stop(ctx context.Context) {
	peerDownsDone := make(chan struct{})

	go func() {
		h.peerDowns.Wait()
		close(peerDownsDone)
	}()

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

	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()

	heldMessages := 0
	for _, held := range h.heldBack {
		heldMessages += len(held)
	}
	h.mu.Unlock()

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

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

		return
	}

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

// applyMessage applies msg, with mu held. A peer's session going down starts
// the peer's PeerDown. An update is applied to the routing table, and its
// prefixes join the batch.
func (h *PrefixHandler) applyMessage(msg *ristypes.RISMessage) {
	if msg.Type == ristypes.MessageTypeState {
		if msg.State == ristypes.StateDown {
			h.startPeerDown(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)

	// 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
		}
	}
}

// startPeerDown runs peerDown for peer on a goroutine of its own, with mu
// held. The peer's messages are held back until it returns. Those held back
// for an earlier PeerDown of the peer and not yet applied stay held.
func (h *PrefixHandler) startPeerDown(peer string) {
	_, goingDown := h.heldBack[peer]
	if !goingDown {
		h.heldBack[peer] = []*ristypes.RISMessage{}
	}

	h.peerDowns.Add(1)

	go h.peerDown(peer)
}

// peerDown runs the table's PeerDown for peer once no other peer's runs. Then
// it applies the messages held back from the peer, in order. When the peer's
// session went down again among them, the PeerDown that starts applies the
// rest.
func (h *PrefixHandler) peerDown(peer string) {
	defer h.peerDowns.Done()

	h.peerDownMu.Lock()
	h.table.PeerDown(peer)
	h.peerDownMu.Unlock()

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

	// applyMessage lets mu go while it waits for the batch. HandleMessage
	// then adds the peer's new messages after these.
	for len(h.heldBack[peer]) > 0 {
		msg := h.heldBack[peer][0]
		h.heldBack[peer] = h.heldBack[peer][1:]
		h.applyMessage(msg)

		if msg.Type == ristypes.MessageTypeState &&
			msg.State == ristypes.StateDown {
			// applyMessage started the next PeerDown, which applies the
			// rest.
			return
		}
	}

	delete(h.heldBack, peer)
}

// 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()
}