Keep one entry per AS in the AS handler's batch (closes #217)

HandleMessage added one entry to the batch for every AS of every path,
and waited for the write in progress once the batch held 300,000. In a
burst of the feed its queue then filled, and the streamer dropped
messages.

The batch is now a map from AS number to the earliest and the latest
time of its updates. It holds one entry per AS, so it no longer grows
with the feed's rate. HandleMessage never waits for a write.
asnBatchSize now counts ASes. Stop logs the entries of the batch being
written and of the next batch under "batch_entries". An AS in both
counts twice.

Model: opus-5-5

internal/routewatch/ashandler.go

@@ -18,43 +18,34 @@ const (
 	// (asnBatchTimeout), so a queue smaller than asnBatchSize is fine.
 	asHandlerQueueSize = 20000
 
-	// asnBatchSize is the number of ASN operations to batch together
+	// asnBatchSize is the number of ASes in a full batch, which wakes the
+	// flush goroutine to write it.
 	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.
+// 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
 
-	// batchMaxSize is asnBatchMaxSize. Only HandleMessage reads it, so a test
-	// may lower it before it hands the handler a message.
-	batchMaxSize int
-
-	// 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 batchMaxSize.
-	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
+	// 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.
@@ -62,16 +53,10 @@ type ASHandler struct {
 	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.
+// 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)
 }
@@ -83,15 +68,13 @@ func newASHandler(
 	db database.Store, logger *logger.Logger, batchTimeout time.Duration,
 ) *ASHandler {
 	h := &ASHandler{
-		db:           db,
-		logger:       logger,
-		batchMaxSize: asnBatchMaxSize,
-		batch:        make([]asnOp, 0, asnBatchSize),
-		batchFullCh:  make(chan struct{}, 1),
-		stopCh:       make(chan context.Context, 1),
-		done:         make(chan struct{}),
+		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{}),
 	}
-	h.batchTaken.L = &h.mu
 
 	// Start the flush timer goroutine
 	go flushLoop(batchTimeout, h.batchFullCh, h.stopCh, h.done, h.flushBatch)
@@ -115,48 +98,25 @@ func (h *ASHandler) QueueCapacity() int {
 	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.
+// 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) {
-	// 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) >= h.batchMaxSize {
-		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
+	// 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 {
-			h.batch = append(h.batch, asnOp{
-				number:    asn,
-				timestamp: timestamp,
-			})
+			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 growing until the write ends
-	// or the batch reaches batchMaxSize.
+	// waits in batchFullCh and the batch keeps taking updates until the write
+	// ends.
 	if len(h.batch) >= asnBatchSize {
 		select {
 		case h.batchFullCh <- struct{}{}:
@@ -165,10 +125,11 @@ func (h *ASHandler) HandleMessage(msg *ristypes.RISMessage) {
 	}
 }
 
-// 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.
+// 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
 
@@ -183,14 +144,14 @@ func (h *ASHandler) Stop(ctx context.Context) {
 
 	if unwritten > 0 {
 		h.logger.Error("Stop deadline passed before the AS handler wrote "+
-			"its updates", "updates", unwritten)
+			"its batches", "batch_entries", unwritten)
 	}
 }
 
-// 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 {
+// 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()
 
@@ -200,8 +161,7 @@ func (h *ASHandler) takeBatch() []asnOp {
 
 	batch := h.batch
 	h.writing = len(batch)
-	h.batch = make([]asnOp, 0, asnBatchSize)
-	h.batchTaken.Broadcast()
+	h.batch = make(map[int]database.SeenTimes)
 
 	return batch
 }
@@ -214,21 +174,11 @@ func (h *ASHandler) flushBatch() {
 		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)
+	err := h.db.GetOrCreateASNBatch(batch)
 	if err != nil {
 		h.logger.Error("Failed to process ASN batch",
-			"error", err, "count", len(asns))
+			"error", err, "count", len(batch))
 	}
 
 	h.mu.Lock()