routewatch/internal/routewatch/prefixhandler.go

package routewatch

import (
	"context"
	"net"
	"net/netip"
	"slices"
	"strings"
	"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"
)

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 60 MiB
	// (README "Memory"). At this size HandleMessage waits until the flush
	// goroutine takes the batch. On top of these, the handler holds at most
	// prefixBatchSize updates, about 6 MiB, that a flush kept for the peers
	// whose route delete failed (held).
	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

	// IP version constants
	ipv4Version = 4
	ipv6Version = 6
)

// PrefixHandler tracks BGP prefixes and maintains a live routing table in the
// database. Routes are added on announcement and deleted on withdrawal, and all
// routes of a peer are deleted when its BGP session goes down.
type PrefixHandler struct {
	db      database.Store
	logger  *logger.Logger
	metrics *metrics.Tracker

	// Batching. mu guards batch, held, peersDown, writing and deleting 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
	// prefixBatchMaxSize.
	mu    sync.Mutex
	batch []prefixUpdate
	// held holds the updates that the peers whose route delete failed in the
	// last flush sent after their session went down, for the next flush to
	// write after it retries the delete. They are not in batch, so they do
	// not count toward a full batch: a delete that keeps failing is retried
	// at the pace of the flush timer.
	held []prefixUpdate
	// peersDown holds, each once, the peers whose BGP session went down
	// since the batch was taken, and those whose route delete failed in the
	// last flush.
	peersDown []string
	// writing is how many updates the flush goroutine took and has not
	// finished writing.
	writing int
	// deleting is how many peers from peersDown the flush goroutine took and
	// has not yet deleted the routes of.
	deleting 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
	lastFlush   time.Time     // used only by the flush goroutine
	// 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
}

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
}

// NewPrefixHandler creates a new batched prefix tracking handler
func NewPrefixHandler(db database.Store, logger *logger.Logger) *PrefixHandler {
	h := &PrefixHandler{
		db:          db,
		logger:      logger,
		batch:       make([]prefixUpdate, 0, prefixBatchSize),
		batchFullCh: make(chan struct{}, 1),
		lastFlush:   time.Now(),
		stopCh:      make(chan context.Context, 1),
		done:        make(chan struct{}),
	}
	h.batchTaken.L = &h.mu

	// Start the flush timer goroutine
	go h.flushLoop()

	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.peerSessionDown(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) >= prefixBatchMaxSize {
		h.batchTaken.Wait()
	}

	// Withdrawals are queued before announcements, because the batch keeps
	// each peer's last update for a prefix: 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 {
		h.batch = append(h.batch, prefixUpdate{
			prefix:      canonicalPrefix(prefix),
			peer:        msg.Peer,
			messageType: "withdrawal",
			timestamp:   timestamp,
		})
		// Record withdrawal in metrics
		if h.metrics != nil {
			h.metrics.RecordWithdrawal()
		}
	}

	for _, announcement := range msg.Announcements {
		for _, prefix := range announcement.Prefixes {
			h.batch = append(h.batch, prefixUpdate{
				prefix:      canonicalPrefix(prefix),
				originASN:   originASN,
				peer:        msg.Peer,
				messageType: "announcement",
				timestamp:   timestamp,
				path:        path,
				endsInSet:   len(msg.Path.Set) > 0,
			})
			// 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 prefixBatchMaxSize.
	if len(h.batch) >= prefixBatchSize {
		select {
		case h.batchFullCh <- struct{}{}:
		default: // already woken
		}
	}
}

// Stop stops the flush goroutine, which writes the updates and peer downs the
// handler holds before it returns, and waits for it until ctx ends. Once ctx
// has ended the flush goroutine skips that last write, and a write still
// running is lost when the process exits. A peer whose route delete fails in
// the last write is not retried either. Stop logs how many updates are not
// written and how many peer downs have not had their routes deleted.
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) + len(h.held)
	peersDown := h.deleting + len(h.peersDown)
	h.mu.Unlock()

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

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

		return
	}

	h.logger.Error("The prefix handler stopped without deleting the routes "+
		"of peers whose session went down", "updates", unwritten,
		"peers_down", peersDown)
}

// peerSessionDown notes that the peer's BGP session went down, so that the next
// flush deletes all of the peer's live routes before it writes the batch. The
// peer's updates still in the batch or held came before the down, which
// replaces them, so they are dropped; updates the peer sends later stay in the
// batch and are written after the delete. A peer already in peersDown is not
// added again: one delete removes all of its routes.
func (h *PrefixHandler) peerSessionDown(peer string) {
	h.mu.Lock()
	defer h.mu.Unlock()

	isPeers := func(update prefixUpdate) bool { return update.peer == peer }
	h.batch = slices.DeleteFunc(h.batch, isPeers)
	h.held = slices.DeleteFunc(h.held, isPeers)

	if !slices.Contains(h.peersDown, peer) {
		h.peersDown = append(h.peersDown, peer)
	}
}

// flushLoop runs in a goroutine, the only one that writes the batch: it
// flushes when the batch is full, periodically, and once more on Stop.
func (h *PrefixHandler) flushLoop() {
	defer close(h.done)

	ticker := time.NewTicker(prefixBatchTimeout)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			if time.Since(h.lastFlush) >= prefixBatchTimeout {
				h.flushBatch()
			}
		case <-h.batchFullCh:
			h.flushBatch()
		case ctx := <-h.stopCh:
			// Final flush, unless the stop deadline has passed
			if ctx.Err() == nil {
				h.flushBatch()
			}

			return
		}
	}
}

// takeBatch returns the pending updates, the held ones first, and the peers in
// peersDown, and leaves an empty batch, no held updates and no peers in their
// place for HandleMessage to fill. The updates count in writing until
// flushBatch has written them, and the peers in deleting until their routes
// are deleted.
func (h *PrefixHandler) takeBatch() ([]prefixUpdate, []string) {
	h.mu.Lock()
	defer h.mu.Unlock()

	if len(h.held) == 0 && len(h.batch) == 0 && len(h.peersDown) == 0 {
		return nil, nil
	}

	batch, peersDown := h.batch, h.peersDown
	if len(h.held) > 0 {
		batch = slices.Concat(h.held, h.batch)
	}

	h.writing = len(batch)
	h.deleting = len(peersDown)
	h.held = nil
	h.batch = make([]prefixUpdate, 0, prefixBatchSize)
	h.peersDown = nil
	h.batchTaken.Broadcast()

	return batch, peersDown
}

// flushBatch takes the pending batch and writes it to the database, without
// holding the mutex during the write.
func (h *PrefixHandler) flushBatch() {
	batch, peersDown := h.takeBatch()
	if len(batch) == 0 && len(peersDown) == 0 {
		return
	}

	startTime := time.Now()
	batchSize := len(batch)

	// The routes of a peer whose session went down are deleted first: the
	// batch holds only the updates the peer sent after the down. A peer whose
	// delete failed is kept, with those updates, for the next flush to retry.
	failedPeers := h.deleteLiveRoutesOfPeers(peersDown)
	if len(failedPeers) > 0 {
		batch = h.keepForNextFlush(batch, failedPeers)
	}

	// For each prefix and peer, keep only the last update the feed delivered.
	// Updates from different peers for the same prefix are all kept. Arrival
	// order decides, not the timestamp: timestamps are whole seconds, so
	// updates within one second carry the same one.
	type prefixAndPeer struct{ prefix, peer string }

	lastUpdates := make(map[prefixAndPeer]prefixUpdate)

	for _, update := range batch {
		lastUpdates[prefixAndPeer{update.prefix, update.peer}] = update
	}

	// Collect routes to upsert and delete
	var (
		routesToUpsert []*database.LiveRoute
		routesToDelete []database.LiveRouteDeletion
	)

	// Collect unique prefixes to update
	prefixesToUpdate := make(map[string]time.Time)

	for _, update := range lastUpdates {
		// Track prefix for both announcements and withdrawals
		seen, exists := prefixesToUpdate[update.prefix]
		if !exists || update.timestamp.After(seen) {
			prefixesToUpdate[update.prefix] = update.timestamp
		}

		// An announcement is stored when its path gives an origin, or when it
		// ends in an AS set of several members: that route is stored with
		// origin 0 (none). An announcement with an empty path is not stored.
		if update.messageType == "announcement" &&
			(update.originASN > 0 || update.endsInSet) {
			// Create live route for batch upsert
			route := h.createLiveRoute(update)
			if route != nil {
				routesToUpsert = append(routesToUpsert, route)
			}
		} else if update.messageType == "withdrawal" {
			// Parse CIDR to get IP version
			_, ipVersion, err := parseCIDR(update.prefix)
			if err != nil {
				h.logger.Error("Failed to parse CIDR for withdrawal",
					"prefix", update.prefix, "error", err)

				continue
			}

			// Create deletion record for batch delete
			routesToDelete = append(routesToDelete, database.LiveRouteDeletion{
				Prefix:    update.prefix,
				PeerIP:    update.peer,
				IPVersion: ipVersion,
			})
		}
	}

	successCount := h.writeBatch(routesToUpsert, routesToDelete, prefixesToUpdate)

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

	elapsed := time.Since(startTime)
	h.logger.Debug("Flushed prefix batch",
		"batch_size", batchSize,
		"peers_down", len(peersDown),
		"unique_prefixes", len(prefixesToUpdate),
		"success", successCount,
		"duration_ms", elapsed.Milliseconds(),
	)

	h.lastFlush = time.Now()
}

// deleteLiveRoutesOfPeers deletes all live routes of each of peers, the peers
// whose BGP session went down, and returns those whose delete failed. Most
// downs come from peers that have no routes, so only a delete that removed
// routes is logged. A peer whose delete succeeded no longer counts in deleting;
// a failed one counts until keepForNextFlush puts it back in peersDown.
func (h *PrefixHandler) deleteLiveRoutesOfPeers(peers []string) []string {
	var failed []string

	for _, peer := range peers {
		deleted, err := h.db.DeleteLiveRoutesOfPeer(peer)
		if err != nil {
			h.logger.Error(
				"Failed to delete the live routes of a peer whose session "+
					"went down; the next flush retries",
				"peer", peer, "routes_deleted", deleted, "error", err)

			failed = append(failed, peer)

			continue
		}

		h.mu.Lock()
		h.deleting--
		h.mu.Unlock()

		if deleted > 0 {
			h.logger.Info(
				"Deleted the live routes of a peer whose session went down",
				"peer", peer, "routes", deleted)
		}
	}

	return failed
}

// keepForNextFlush puts the peers whose route delete failed back in peersDown,
// and their updates from batch in held, so that the next flush retries each
// delete before it writes the peer's updates; written now, they would be
// deleted by the retry. The updates of a peer whose session went down again
// since the batch was taken came before that down, so they are dropped, as
// peerSessionDown drops them, and the peer is already in peersDown. It returns
// the updates of the other peers, for this flush to write.
//
// While a delete keeps failing, the held updates would grow with every flush,
// so only the last prefixBatchSize of them are kept and the older ones are
// dropped. This never leaves in place a route the peer withdrew: a prefix with
// a kept update keeps its last one, and the retry deletes all of the peer's
// routes, so a prefix whose updates were all dropped has no route from the peer
// until the peer announces it again.
func (h *PrefixHandler) keepForNextFlush(
	batch []prefixUpdate, failedPeers []string,
) []prefixUpdate {
	var kept, toWrite []prefixUpdate

	for _, update := range batch {
		if slices.Contains(failedPeers, update.peer) {
			kept = append(kept, update)
		} else {
			toWrite = append(toWrite, update)
		}
	}

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

	kept = slices.DeleteFunc(kept, func(update prefixUpdate) bool {
		return slices.Contains(h.peersDown, update.peer)
	})

	if len(kept) > prefixBatchSize {
		h.logger.Error(
			"Dropped the oldest updates of peers whose route delete failed",
			"peers", failedPeers, "dropped", len(kept)-prefixBatchSize)

		// A copy, so that the dropped updates are freed.
		kept = slices.Clone(kept[len(kept)-prefixBatchSize:])
	}

	// takeBatch left held empty, and only this flush fills it. The kept
	// updates now count in held, not in writing, and the failed peers in
	// peersDown, not in deleting.
	h.held = kept
	h.writing = len(toWrite)
	h.deleting = 0

	for _, peer := range failedPeers {
		if !slices.Contains(h.peersDown, peer) {
			h.peersDown = append(h.peersDown, peer)
		}
	}

	return toWrite
}

// writeBatch writes the collected route upserts, route deletions and prefix
// updates to the database and returns how many routes it wrote.
func (h *PrefixHandler) writeBatch(
	routesToUpsert []*database.LiveRoute,
	routesToDelete []database.LiveRouteDeletion,
	prefixesToUpdate map[string]time.Time,
) int {
	successCount := 0

	if len(routesToUpsert) > 0 {
		err := h.db.UpsertLiveRouteBatch(routesToUpsert)
		if err != nil {
			h.logger.Error("Failed to upsert route batch",
				"error", err, "count", len(routesToUpsert))
		} else {
			successCount += len(routesToUpsert)
		}
	}

	if len(routesToDelete) > 0 {
		err := h.db.DeleteLiveRouteBatch(routesToDelete)
		if err != nil {
			h.logger.Error("Failed to delete route batch",
				"error", err, "count", len(routesToDelete))
		} else {
			successCount += len(routesToDelete)
		}
	}

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

	return successCount
}

// canonicalPrefix returns prefix in the text form net/netip prints, the form
// the IP lookup builds when it looks a prefix up. The feed sends IPv6
// announcements compressed ("2001:db8::/32") but IPv6 withdrawals uncompressed
// ("2001:db8:0:0:0:0:0:0/32"); stored as received, a withdrawal would not match
// the route its announcement stored. A prefix that does not parse is returned
// unchanged, and the batch flush reports it.
func canonicalPrefix(prefix string) string {
	p, err := netip.ParsePrefix(prefix)
	if err != nil {
		return prefix
	}

	return p.Masked().String()
}

// parseCIDR returns the mask length and the IP version of a prefix string
func parseCIDR(prefix string) (int, int, error) {
	_, ipNet, err := net.ParseCIDR(prefix)
	if err != nil {
		return 0, 0, err
	}

	ones, _ := ipNet.Mask.Size()
	if strings.Contains(prefix, ":") {
		return ones, ipv6Version, nil
	}

	return ones, ipv4Version, nil
}

// createLiveRoute creates a LiveRoute from a prefix update
func (h *PrefixHandler) createLiveRoute(update prefixUpdate) *database.LiveRoute {
	// Parse CIDR to get mask length
	maskLength, ipVersion, err := parseCIDR(update.prefix)
	if err != nil {
		h.logger.Error("Failed to parse CIDR",
			"prefix", update.prefix,
			"error", err,
		)

		return nil
	}

	// Track route update metrics
	if h.metrics != nil {
		if ipVersion == ipv4Version {
			h.metrics.RecordIPv4Update()
		} else {
			h.metrics.RecordIPv6Update()
		}
	}

	// Create live route record
	liveRoute := &database.LiveRoute{
		Prefix:      update.prefix,
		MaskLength:  maskLength,
		IPVersion:   ipVersion,
		OriginASN:   update.originASN,
		PeerIP:      update.peer,
		ASPath:      update.path,
		NextHop:     update.peer, // Using peer as next hop
		LastUpdated: update.timestamp,
	}

	return liveRoute
}