routewatch/internal/routewatch/ashandler.go

package routewatch

import (
	"context"
	"slices"
	"sync"
	"time"

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

const (
	// asHandlerQueueSize is the queue capacity for ASN 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
	// (asnBatchTimeout), so a queue smaller than asnBatchSize is fine.
	asHandlerQueueSize = 20000

	// asnBatchSize is the number of ASN operations to batch together
	asnBatchSize = 30000

	// asnBatchMaxSize is the most ASN operations 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.
	asnBatchMaxSize = 300000

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

// ASHandler processes Autonomous System Number (ASN) information extracted from
// BGP UPDATE messages. It uses batched database operations to efficiently store
// ASN data, collecting operations into batches that are flushed either when the
// batch reaches a size threshold or after a timeout period.
type ASHandler struct {
	db     database.Store
	logger *logger.Logger

	// Batching. mu guards batch 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 batch is at asnBatchMaxSize.
	mu    sync.Mutex
	batch []asnOp
	// writing is how many operations the flush goroutine took from batch
	// and has not finished writing.
	writing 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 asnOp struct {
	number    int
	timestamp time.Time
}

// NewASHandler creates and returns a new ASHandler instance. It initializes
// the batching system and starts a background goroutine that periodically
// flushes accumulated ASN operations to the database. The caller must call
// Stop when finished to ensure all pending operations are flushed and the
// background goroutine is terminated.
func NewASHandler(db database.Store, logger *logger.Logger) *ASHandler {
	h := &ASHandler{
		db:          db,
		logger:      logger,
		batch:       make([]asnOp, 0, asnBatchSize),
		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
}

// WantsMessage reports whether this handler should process messages of the
// given type. ASHandler only processes "UPDATE" messages, as these contain
// the AS path information needed to track autonomous systems.
func (h *ASHandler) WantsMessage(messageType string) bool {
	// We only care about UPDATE messages for the database
	return messageType == ristypes.MessageTypeUpdate
}

// QueueCapacity returns the recommended message queue size for this handler.
// ASHandler uses a large queue capacity to accommodate high-volume BGP streams,
// as the batching mechanism allows efficient processing of accumulated messages.
func (h *ASHandler) QueueCapacity() int {
	// Batching allows us to use a larger queue
	return asHandlerQueueSize
}

// HandleMessage processes a RIS Live BGP message by extracting all ASNs from
// the AS path and queuing them for batch insertion into the database. The
// origin ASN and all transit ASNs are recorded with their associated
// timestamps. When the batch reaches the configured size threshold, the flush
// goroutine is woken to write it.
func (h *ASHandler) HandleMessage(msg *ristypes.RISMessage) {
	// Use the pre-parsed timestamp
	timestamp := msg.ParsedTimestamp

	originASN := msg.Path.Origin()

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

	// Queue origin ASN operation
	if originASN > 0 {
		h.batch = append(h.batch, asnOp{
			number:    originASN,
			timestamp: timestamp,
		})
	}

	// Also track all ASNs in the path, the members of an AS set included
	for _, asn := range slices.Concat(msg.Path.Sequence, msg.Path.Set) {
		if asn > 0 {
			h.batch = append(h.batch, asnOp{
				number:    asn,
				timestamp: timestamp,
			})
		}
	}

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

// Stop stops the flush goroutine, which writes the operations 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 operations are not written.
func (h *ASHandler) Stop(ctx context.Context) {
	h.stopCh <- ctx

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

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

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

// 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 *ASHandler) flushLoop() {
	defer close(h.done)

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

	for {
		select {
		case <-ticker.C:
			if time.Since(h.lastFlush) >= asnBatchTimeout {
				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 operations and leaves an empty batch in their
// place for HandleMessage to fill. They count in writing until flushBatch has
// written them.
func (h *ASHandler) takeBatch() []asnOp {
	h.mu.Lock()
	defer h.mu.Unlock()

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

	batch := h.batch
	h.writing = len(batch)
	h.batch = make([]asnOp, 0, asnBatchSize)
	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 *ASHandler) flushBatch() {
	batch := h.takeBatch()
	if len(batch) == 0 {
		return
	}

	// The earliest and the latest time the batch saw each AS. The batch holds
	// several updates for an AS, from many peers, so a new AS is stored first
	// seen at the earliest of them and last seen at the latest.
	asns := make(map[int]database.SeenTimes)
	for _, op := range batch {
		seen := asns[op.number]
		seen.Add(op.timestamp)
		asns[op.number] = seen
	}

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

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

	h.lastFlush = time.Now()
}