routewatch/internal/routewatch/peerhandler.go

package routewatch

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

import (
	"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

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

// PeerHandler tracks BGP peers from all message types using batched operations.
// It maintains a queue of peer updates and periodically flushes them to the database
// in batches to improve performance. The handler deduplicates peer updates within
// each batch, keeping only the most recent update for each peer IP address.
type PeerHandler struct {
	db     database.Store
	logger *logger.Logger

	// Batching. mu guards peerBatch only: the flush goroutine takes the batch
	// under mu and writes it after releasing mu, so HandleMessage never
	// waits for a database write.
	mu          sync.Mutex
	peerBatch   []peerUpdate
	batchFullCh chan struct{} // wakes the flush goroutine when peerBatch is full
	lastFlush   time.Time     // used only by the flush goroutine
	stopCh      chan struct{}
	wg          sync.WaitGroup
}

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 {
	h := &PeerHandler{
		db:          db,
		logger:      logger,
		peerBatch:   make([]peerUpdate, 0, peerBatchSize),
		batchFullCh: make(chan struct{}, 1),
		lastFlush:   time.Now(),
		stopCh:      make(chan struct{}),
	}

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

	return h
}

// WantsMessage returns true for all message types since peer information
// is extracted from every RIS message regardless of type. This satisfies
// the MessageHandler interface.
func (h *PeerHandler) WantsMessage(_ string) bool {
	return true
}

// 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 != "" {
		if asn, err := strconv.Atoi(msg.PeerASN); err == nil {
			peerASN = asn
		}
	}

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

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

// 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 *PeerHandler) flushLoop() {
	defer h.wg.Done()
	ticker := time.NewTicker(peerBatchTimeout)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			if time.Since(h.lastFlush) >= peerBatchTimeout {
				h.flushBatch()
			}
		case <-h.batchFullCh:
			h.flushBatch()
		case <-h.stopCh:
			// Final flush
			h.flushBatch()

			return
		}
	}
}

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

	if len(h.peerBatch) == 0 {
		return nil
	}
	batch := h.peerBatch
	h.peerBatch = make([]peerUpdate, 0, peerBatchSize)

	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
	}

	// Deduplicate by peer IP, keeping the latest update for each peer
	peerMap := make(map[string]peerUpdate)
	for _, update := range batch {
		if existing, ok := peerMap[update.peerIP]; !ok || update.timestamp.After(existing.timestamp) {
			peerMap[update.peerIP] = update
		}
	}

	// Convert to database format
	dbPeerMap := make(map[string]database.PeerUpdate)
	for peerIP, update := range peerMap {
		dbPeerMap[peerIP] = database.PeerUpdate{
			PeerIP:      update.peerIP,
			PeerASN:     update.peerASN,
			MessageType: update.messageType,
			Timestamp:   update.timestamp,
		}
	}

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

	h.lastFlush = time.Now()
}

// Stop gracefully stops the handler and flushes remaining batches
func (h *PeerHandler) Stop() {
	close(h.stopCh)
	h.wg.Wait()
}