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 ASes in a full batch, which wakes the
	// flush goroutine to write it.
	asnBatchSize = 30000

	// 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 records every AS in the AS paths of BGP UPDATE messages. Its batch
// holds the earliest and the latest time of each AS's updates. The flush
// goroutine writes the batch to the database when it holds asnBatchSize ASes,
// and every asnBatchTimeout.
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
	// never waits for a database write.
	mu sync.Mutex
	// batch holds the earliest and the latest time of the updates for each AS
	// since the flush goroutine last took it. It holds one entry per AS,
	// however many updates arrive.
	batch map[int]database.SeenTimes
	// writing is how many ASes the flush goroutine took from batch and has
	// not finished writing.
	writing     int
	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
}

// NewASHandler creates and returns a new ASHandler instance. It initializes
// the batching system and starts a background goroutine that periodically
// writes the collected ASes to the database. The caller must call Stop when
// finished, to write the pending ASes and end the background goroutine.
func NewASHandler(db database.Store, logger *logger.Logger) *ASHandler {
	return newASHandler(db, logger, asnBatchTimeout)
}

// newASHandler is NewASHandler with the flush timer's interval, batchTimeout,
// in place of asnBatchTimeout. A test that the timer must not disturb gives an
// interval longer than the test runs.
func newASHandler(
	db database.Store, logger *logger.Logger, batchTimeout time.Duration,
) *ASHandler {
	h := &ASHandler{
		db:          db,
		logger:      logger,
		batch:       make(map[int]database.SeenTimes),
		batchFullCh: make(chan struct{}, 1),
		stopCh:      make(chan context.Context, 1),
		done:        make(chan struct{}),
	}

	// Start the flush timer goroutine
	go flushLoop(batchTimeout, h.batchFullCh, h.stopCh, h.done, h.flushBatch)

	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 adds the time of msg to the batch entry of each AS in its AS
// path, the origin AS and the members of an AS set included. When the batch
// holds asnBatchSize ASes, HandleMessage wakes the flush goroutine to write it.
func (h *ASHandler) HandleMessage(msg *ristypes.RISMessage) {
	h.mu.Lock()
	defer h.mu.Unlock()

	// The origin AS is the last AS of the sequence, so the loop adds it too.
	for _, asn := range slices.Concat(msg.Path.Sequence, msg.Path.Set) {
		if asn > 0 {
			seen := h.batch[asn]
			seen.Add(msg.ParsedTimestamp)
			h.batch[asn] = seen
		}
	}

	// A full batch wakes the flush goroutine. If it is writing, the wake-up
	// waits in batchFullCh and the batch keeps taking updates until the write
	// ends.
	if len(h.batch) >= asnBatchSize {
		select {
		case h.batchFullCh <- struct{}{}:
		default: // already woken
		}
	}
}

// Stop stops the flush goroutine and waits for it until ctx ends. The flush
// goroutine writes the batch before it returns, unless ctx has ended. A write
// still running when the process exits is lost. Stop logs how many entries are
// not written, in the batch being written and in the next batch. An AS in both
// batches counts twice.
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 batches", "batch_entries", unwritten)
	}
}

// takeBatch returns the pending ASes 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() map[int]database.SeenTimes {
	h.mu.Lock()
	defer h.mu.Unlock()

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

	batch := h.batch
	h.writing = len(batch)
	h.batch = make(map[int]database.SeenTimes)

	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
	}

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

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