routewatch/internal/routewatch/peerhandler.go

package routewatch

// peerhandler.go provides batched peer tracking functionality for BGP route
// monitoring. It tracks BGP peers from incoming UPDATE messages and maintains
// peer state in the database.

import (
	"context"
	"strconv"
	"sync"
	"time"

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

const (
	// peerHandlerQueueSize is the queue capacity for peer tracking operations,
	// about 4 seconds of feed at peak. The streamer drops rather than blocks
	// when a queue is full, so this bounds memory. Batches still flush on a
	// timer (peerBatchTimeout).
	peerHandlerQueueSize = 20000

	// peerBatchSize is the number of peer updates to batch together
	peerBatchSize = 10000

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

	// peerBatchTimeout is the maximum time to wait before flushing a batch
	peerBatchTimeout = 2 * time.Second
)

// PeerHandler tracks BGP peers from UPDATE messages using batched operations.
// It maintains a queue of peer updates and periodically flushes them to the database
// in batches to improve performance. The handler writes one row per peer IP
// address for each batch, with the earliest and the latest time of its updates.
type PeerHandler struct {
	db     database.Store
	logger *logger.Logger

	// batchMaxSize is peerBatchMaxSize. Only HandleMessage reads it, so a test
	// may lower it before it hands the handler a message.
	batchMaxSize int

	// Batching. mu guards peerBatch and writing only: the flush goroutine
	// takes the batch under mu and writes it after releasing mu, so
	// HandleMessage waits for a database write only while peerBatch is at
	// batchMaxSize.
	mu        sync.Mutex
	peerBatch []peerUpdate
	// writing is how many updates the flush goroutine took from peerBatch
	// and has not finished writing.
	writing int
	// batchTaken is signalled when the flush goroutine takes peerBatch; its L
	// is &mu.
	batchTaken  sync.Cond
	batchFullCh chan struct{} // wakes the flush goroutine when peerBatch 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
}

type peerUpdate struct {
	peerIP      string
	peerASN     int
	messageType string
	timestamp   time.Time
}

// NewPeerHandler creates a new PeerHandler with the given database store and logger.
// It initializes the peer batch buffer and starts a background goroutine that
// periodically flushes accumulated peer updates to the database. The handler
// should be stopped by calling Stop when it is no longer needed.
func NewPeerHandler(db database.Store, logger *logger.Logger) *PeerHandler {
	return newPeerHandler(db, logger, peerBatchTimeout)
}

// newPeerHandler is NewPeerHandler with the flush timer's interval,
// batchTimeout, in place of peerBatchTimeout. A test that the timer must not
// disturb gives an interval longer than the test runs.
func newPeerHandler(
	db database.Store, logger *logger.Logger, batchTimeout time.Duration,
) *PeerHandler {
	h := &PeerHandler{
		db:           db,
		logger:       logger,
		batchMaxSize: peerBatchMaxSize,
		peerBatch:    make([]peerUpdate, 0, peerBatchSize),
		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
}

// WantsMessage returns true for UPDATE messages only. A peer whose session
// never comes up sends STATE messages and nothing else, so a STATE message is
// not a sign of a peer that carries routes.
func (h *PeerHandler) WantsMessage(messageType string) bool {
	return messageType == ristypes.MessageTypeUpdate
}

// QueueCapacity returns the desired queue capacity for this handler.
// The PeerHandler uses a large queue capacity because batching allows
// for efficient processing of many updates at once.
func (h *PeerHandler) QueueCapacity() int {
	// Batching allows us to use a larger queue
	return peerHandlerQueueSize
}

// HandleMessage processes a RIS message to track peer information.
// It extracts the peer IP address and ASN from the message and adds
// the update to an internal batch. The flush goroutine writes the batch to the
// database when it reaches peerBatchSize or the batch timeout expires.
func (h *PeerHandler) HandleMessage(msg *ristypes.RISMessage) {
	// Parse peer ASN from string
	peerASN := 0

	if msg.PeerASN != "" {
		asn, err := strconv.Atoi(msg.PeerASN)
		if err == nil {
			peerASN = asn
		}
	}

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

	// Add to batch
	h.peerBatch = append(h.peerBatch, peerUpdate{
		peerIP:      msg.Peer,
		peerASN:     peerASN,
		messageType: msg.Type,
		timestamp:   msg.ParsedTimestamp,
	})

	// 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.peerBatch) >= peerBatchSize {
		select {
		case h.batchFullCh <- struct{}{}:
		default: // already woken
		}
	}
}

// Stop stops the flush goroutine, which writes the updates 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, so Stop logs how many updates are not written.
func (h *PeerHandler) Stop(ctx context.Context) {
	h.stopCh <- ctx

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

	h.mu.Lock()
	unwritten := h.writing + len(h.peerBatch)
	h.mu.Unlock()

	if unwritten > 0 {
		h.logger.Error("Stop deadline passed before the peer handler wrote "+
			"its updates", "updates", unwritten)
	}
}

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

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

	batch := h.peerBatch
	h.writing = len(batch)
	h.peerBatch = make([]peerUpdate, 0, peerBatchSize)
	h.batchTaken.Broadcast()

	return batch
}

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

	// The earliest and the latest time the batch saw each peer, with the peer
	// AS and message type of its latest update. The batch holds several
	// updates for a peer, so a new peer is stored first seen at the earliest of
	// them and last seen at the latest.
	peers := make(map[string]database.PeerUpdate)
	for _, update := range batch {
		peer, ok := peers[update.peerIP]
		if !ok || update.timestamp.After(peer.Seen.Last) {
			peer.PeerIP = update.peerIP
			peer.PeerASN = update.peerASN
			peer.MessageType = update.messageType
		}

		peer.Seen.Add(update.timestamp)
		peers[update.peerIP] = peer
	}

	// Process all peers in a single batch transaction
	err := h.db.UpdatePeerBatch(peers)
	if err != nil {
		h.logger.Error("Failed to process peer batch",
			"error", err, "count", len(peers))
	}

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