Bound each shutdown stage by the stop deadline (closes #100)

Shutdown passes its context to every stage: the streamer, the four
handlers, the WHOIS fetcher, the database maintainer and the HTTP server
each stop waiting when it ends. A handler then skips its last write and
logs how many updates it did not write, and the HTTP server closes the
connections still open. The database is checkpointed and closed last,
and left open once the deadline has passed, since a write may still
hold it. The stop hook gives Shutdown 8 seconds and fx 9, inside the 10
seconds upaas and docker stop allow, so a stop signal exits 0.

Model: opus-5-5

internal/routewatch/prefixhandler.go

@@ -1,6 +1,7 @@
 package routewatch
 
 import (
+	"context"
 	"net"
 	"net/netip"
 	"slices"
@@ -50,9 +51,9 @@ type PrefixHandler struct {
 	logger  *logger.Logger
 	metrics *metrics.Tracker
 
-	// Batching. mu guards batch, held and peersDown only: the flush goroutine
-	// takes them under mu and writes them after releasing mu, so
-	// HandleMessage waits for a database write only while batch is at
+	// Batching. mu guards batch, held, peersDown, writing and deleting only:
+	// the flush goroutine takes them under mu and writes them after releasing
+	// mu, so HandleMessage waits for a database write only while batch is at
 	// prefixBatchMaxSize.
 	mu    sync.Mutex
 	batch []prefixUpdate
@@ -66,13 +67,21 @@ type PrefixHandler struct {
 	// since the batch was taken, and those whose route delete failed in the
 	// last flush.
 	peersDown []string
+	// writing is how many updates the flush goroutine took and has not
+	// finished writing.
+	writing int
+	// deleting is how many peers from peersDown the flush goroutine took and
+	// has not yet deleted the routes of.
+	deleting 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      chan struct{}
-	wg          sync.WaitGroup
+	// 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 prefixUpdate struct {
@@ -95,12 +104,12 @@ func NewPrefixHandler(db database.Store, logger *logger.Logger) *PrefixHandler {
 		batch:       make([]prefixUpdate, 0, prefixBatchSize),
 		batchFullCh: make(chan struct{}, 1),
 		lastFlush:   time.Now(),
-		stopCh:      make(chan struct{}),
+		stopCh:      make(chan context.Context, 1),
+		done:        make(chan struct{}),
 	}
 	h.batchTaken.L = &h.mu
 
 	// Start the flush timer goroutine
-	h.wg.Add(1)
 	go h.flushLoop()
 
 	return h
@@ -202,10 +211,39 @@ func (h *PrefixHandler) HandleMessage(msg *ristypes.RISMessage) {
 	}
 }
 
-// Stop gracefully stops the handler and flushes remaining batches
-func (h *PrefixHandler) Stop() {
-	close(h.stopCh)
-	h.wg.Wait()
+// Stop stops the flush goroutine, which writes the updates and peer downs 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. A peer whose route delete fails in
+// the last write is not retried either. Stop logs how many updates are not
+// written and how many peer downs have not had their routes deleted.
+func (h *PrefixHandler) Stop(ctx context.Context) {
+	h.stopCh <- ctx
+
+	select {
+	case <-h.done:
+	case <-ctx.Done():
+	}
+
+	h.mu.Lock()
+	unwritten := h.writing + len(h.batch) + len(h.held)
+	peersDown := h.deleting + len(h.peersDown)
+	h.mu.Unlock()
+
+	if unwritten == 0 && peersDown == 0 {
+		return
+	}
+
+	if ctx.Err() != nil {
+		h.logger.Error("Stop deadline passed before the prefix handler wrote "+
+			"its updates", "updates", unwritten, "peers_down", peersDown)
+
+		return
+	}
+
+	h.logger.Error("The prefix handler stopped without deleting the routes "+
+		"of peers whose session went down", "updates", unwritten,
+		"peers_down", peersDown)
 }
 
 // peerSessionDown notes that the peer's BGP session went down, so that the next
@@ -230,7 +268,7 @@ func (h *PrefixHandler) peerSessionDown(peer string) {
 // 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 *PrefixHandler) flushLoop() {
-	defer h.wg.Done()
+	defer close(h.done)
 
 	ticker := time.NewTicker(prefixBatchTimeout)
 	defer ticker.Stop()
@@ -243,9 +281,11 @@ func (h *PrefixHandler) flushLoop() {
 			}
 		case <-h.batchFullCh:
 			h.flushBatch()
-		case <-h.stopCh:
-			// Final flush
-			h.flushBatch()
+		case ctx := <-h.stopCh:
+			// Final flush, unless the stop deadline has passed
+			if ctx.Err() == nil {
+				h.flushBatch()
+			}
 
 			return
 		}
@@ -254,7 +294,9 @@ func (h *PrefixHandler) flushLoop() {
 
 // takeBatch returns the pending updates, the held ones first, and the peers in
 // peersDown, and leaves an empty batch, no held updates and no peers in their
-// place for HandleMessage to fill.
+// place for HandleMessage to fill. The updates count in writing until
+// flushBatch has written them, and the peers in deleting until their routes
+// are deleted.
 func (h *PrefixHandler) takeBatch() ([]prefixUpdate, []string) {
 	h.mu.Lock()
 	defer h.mu.Unlock()
@@ -268,6 +310,8 @@ func (h *PrefixHandler) takeBatch() ([]prefixUpdate, []string) {
 		batch = slices.Concat(h.held, h.batch)
 	}
 
+	h.writing = len(batch)
+	h.deleting = len(peersDown)
 	h.held = nil
 	h.batch = make([]prefixUpdate, 0, prefixBatchSize)
 	h.peersDown = nil
@@ -354,6 +398,10 @@ func (h *PrefixHandler) flushBatch() {
 
 	successCount := h.writeBatch(routesToUpsert, routesToDelete, prefixesToUpdate)
 
+	h.mu.Lock()
+	h.writing = 0
+	h.mu.Unlock()
+
 	elapsed := time.Since(startTime)
 	h.logger.Debug("Flushed prefix batch",
 		"batch_size", batchSize,
@@ -369,7 +417,8 @@ func (h *PrefixHandler) flushBatch() {
 // deleteLiveRoutesOfPeers deletes all live routes of each of peers, the peers
 // whose BGP session went down, and returns those whose delete failed. Most
 // downs come from peers that have no routes, so only a delete that removed
-// routes is logged.
+// routes is logged. A peer whose delete succeeded no longer counts in deleting;
+// a failed one counts until keepForNextFlush puts it back in peersDown.
 func (h *PrefixHandler) deleteLiveRoutesOfPeers(peers []string) []string {
 	var failed []string
 
@@ -386,6 +435,10 @@ func (h *PrefixHandler) deleteLiveRoutesOfPeers(peers []string) []string {
 			continue
 		}
 
+		h.mu.Lock()
+		h.deleting--
+		h.mu.Unlock()
+
 		if deleted > 0 {
 			h.logger.Info(
 				"Deleted the live routes of a peer whose session went down",
@@ -439,8 +492,12 @@ func (h *PrefixHandler) keepForNextFlush(
 		kept = slices.Clone(kept[len(kept)-prefixBatchSize:])
 	}
 
-	// takeBatch left held empty, and only this flush fills it.
+	// takeBatch left held empty, and only this flush fills it. The kept
+	// updates now count in held, not in writing, and the failed peers in
+	// peersDown, not in deleting.
 	h.held = kept
+	h.writing = len(toWrite)
+	h.deleting = 0
 
 	for _, peer := range failedPeers {
 		if !slices.Contains(h.peersDown, peer) {