Drain the handler queues before stopping the handlers on shutdown (closes #69)

Shutdown stopped the handlers before the streamer, so updates still in the
handler queues, up to 20,000 per queue, were not written, and a lost
withdrawal left its route in the table after a restart.

Shutdown now stops the streamer first. Streamer.Stop queues the STATE
messages it holds for a full queue, closes the queues and returns once
every worker has handed its handler everything on its queue; then each
handler writes what it holds before its Stop returns. PeeringHandler.Stop
now waits for its loop's final write instead of running a second one
beside it.

Model: opus-5-5

internal/routewatch/app.go

@@ -149,7 +149,10 @@ func (rw *RouteWatch) Shutdown(ctx context.Context) {
 	rw.shutdown = true
 	rw.mu.Unlock()
 
-	// Stop batched handlers first to flush remaining batches
+	// Stop the streamer first: it returns once the handlers have been handed
+	// every message it queued. Then each handler writes what it holds.
+	rw.streamer.Stop()
+
 	if rw.dbHandler != nil {
 		rw.logger.Info("Flushing database handler")
 		rw.dbHandler.Stop()
@@ -180,9 +183,6 @@ func (rw *RouteWatch) Shutdown(ctx context.Context) {
 		rw.dbMaintainer.Stop()
 	}
 
-	// Stop services
-	rw.streamer.Stop()
-
 	// Stop HTTP server with a timeout of its own, however much of ctx is left
 	const serverStopTimeout = 5 * time.Second
 

internal/routewatch/app_integration_internal_test.go

@@ -3,6 +3,9 @@ package routewatch
 import (
 	"context"
 	"fmt"
+	"io"
+	"net/http"
+	"net/http/httptest"
 	"strings"
 	"sync"
 	"testing"
@@ -12,6 +15,7 @@ import (
 	"git.eeqj.de/sneak/routewatch/internal/database"
 	"git.eeqj.de/sneak/routewatch/internal/logger"
 	"git.eeqj.de/sneak/routewatch/internal/metrics"
+	"git.eeqj.de/sneak/routewatch/internal/ristypes"
 	"git.eeqj.de/sneak/routewatch/internal/server"
 	"git.eeqj.de/sneak/routewatch/internal/streamer"
 	"github.com/google/uuid"
@@ -529,3 +533,124 @@ func TestRouteWatchLiveFeed(t *testing.T) {
 
 	t.Logf("Recorded %d AS peering relationships in 5 seconds", stats.Peerings)
 }
+
+// gatedPrefixHandler is a PrefixHandler whose worker in the streamer waits on
+// the first message until release is closed, so the messages after it stay on
+// the streamer's queue.
+type gatedPrefixHandler struct {
+	*PrefixHandler
+
+	release chan struct{}
+}
+
+func (h gatedPrefixHandler) HandleMessage(msg *ristypes.RISMessage) {
+	<-h.release
+	h.PrefixHandler.HandleMessage(msg)
+}
+
+// announcementFeed returns the lines RIS Live would send for one UPDATE per
+// prefix, each announcing that prefix from testPeerIP with origin testASNA,
+// followed by a KEEPALIVE. The streamer hands a KEEPALIVE to no handler, so
+// once a raw handler has seen it, every update before it is queued.
+func announcementFeed(prefixes []string) string {
+	var feed strings.Builder
+
+	for _, prefix := range prefixes {
+		fmt.Fprintf(&feed, `{"type":"ris_message","data":{"type":"UPDATE",`+
+			`"peer":"%s","peer_asn":"%d","timestamp":1700000000,"path":[%d],`+
+			`"announcements":[{"next_hop":"%s","prefixes":["%s"]}]}}`+"\n",
+			testPeerIP, testASNA, testASNA, testPeerIP, prefix)
+	}
+
+	fmt.Fprintf(&feed, `{"type":"ris_message","data":{"type":"KEEPALIVE",`+
+		`"peer":"%s","peer_asn":"%d","timestamp":1700000000}}`+"\n",
+		testPeerIP, testASNA)
+
+	return feed.String()
+}
+
+// TestShutdownWritesTheUpdatesQueuedInTheStreamer has a local server send
+// announcements of three prefixes and keep the response open. The prefix
+// handler's worker waits on the first one until the server sees the streamer
+// close the connection, which only Shutdown does, so all three are still in
+// the streamer when Shutdown starts. Each prefix must be in the database when
+// Shutdown returns.
+func TestShutdownWritesTheUpdatesQueuedInTheStreamer(t *testing.T) {
+	t.Parallel()
+
+	const maxWait = 10 * time.Second
+
+	prefixes := []string{"192.0.2.0/24", "198.51.100.0/24", "203.0.113.0/24"}
+	feed := announcementFeed(prefixes)
+
+	release := make(chan struct{})
+
+	srv := httptest.NewServer(http.HandlerFunc(
+		func(w http.ResponseWriter, r *http.Request) {
+			_, _ = io.WriteString(w, feed)
+			_ = http.NewResponseController(w).Flush()
+
+			<-r.Context().Done()
+			close(release)
+		}))
+	defer srv.Close()
+
+	store := newMockStore()
+	log := logger.New()
+	cfg := &config.Config{}
+	s := streamer.New(log, metrics.New())
+	s.SetURL(srv.URL)
+
+	rw := New(Dependencies{
+		DB:       store,
+		Streamer: s,
+		Server:   server.New(store, s, log, cfg),
+		Logger:   log,
+		Config:   cfg,
+	})
+	rw.prefixHandler = NewPrefixHandler(store, log)
+	s.RegisterHandler(gatedPrefixHandler{
+		PrefixHandler: rw.prefixHandler, release: release,
+	})
+
+	linesSeen := make(chan struct{}, len(prefixes)+1)
+
+	s.RegisterRawHandler(func(string) { linesSeen <- struct{}{} })
+
+	err := s.Start(t.Context())
+	if err != nil {
+		t.Fatalf("Start returned error: %v", err)
+	}
+
+	for range cap(linesSeen) {
+		select {
+		case <-linesSeen:
+		case <-time.After(maxWait):
+			t.Fatalf("the streamer did not read %d lines within %s",
+				cap(linesSeen), maxWait)
+		}
+	}
+
+	shutDown := make(chan struct{})
+
+	go func() {
+		rw.Shutdown(t.Context())
+		close(shutDown)
+	}()
+
+	select {
+	case <-shutDown:
+	case <-time.After(maxWait):
+		t.Fatalf("Shutdown did not return within %s", maxWait)
+	}
+
+	store.mu.Lock()
+	defer store.mu.Unlock()
+
+	for _, prefix := range prefixes {
+		if _, ok := store.Prefixes[prefix]; !ok {
+			t.Errorf("%s was queued in the streamer when Shutdown started "+
+				"but is not in the database", prefix)
+		}
+	}
+}

internal/routewatch/peeringhandler.go

@@ -47,6 +47,7 @@ type PeeringHandler struct {
 	droppedPaths int                  // paths dropped because the map was full
 
 	stopCh chan struct{}
+	wg     sync.WaitGroup
 }
 
 // NewPeeringHandler creates and initializes a new PeeringHandler with the
@@ -62,6 +63,8 @@ func NewPeeringHandler(db database.Store, logger *logger.Logger) *PeeringHandler
 	}
 
 	// Start the periodic processing goroutine
+	h.wg.Add(1)
+
 	go h.processLoop()
 
 	return h
@@ -130,17 +133,18 @@ func (h *PeeringHandler) ProcessPeeringsNow() {
 }
 
 // Stop gracefully shuts down the handler by signaling the background
-// goroutine to stop and performing a final synchronous processing of
-// any remaining AS paths. This ensures no peering data is lost during
-// shutdown.
+// goroutine to stop and waiting while it processes the remaining AS paths,
+// so no peering data is lost during shutdown.
 func (h *PeeringHandler) Stop() {
 	close(h.stopCh)
-	// Process any remaining peerings synchronously
-	h.processPeerings()
+	h.wg.Wait()
 }
 
-// processLoop runs periodically to process AS paths into peerings
+// processLoop runs periodically to process AS paths into peerings, and once
+// more on Stop.
 func (h *PeeringHandler) processLoop() {
+	defer h.wg.Done()
+
 	ticker := time.NewTicker(peeringProcessInterval)
 	defer ticker.Stop()
 

internal/routewatch/peeringhandler_internal_test.go

@@ -188,6 +188,22 @@ func TestPeeringHandlerProcessesEachRunAndEmpties(t *testing.T) {
 	}
 }
 
+// TestPeeringHandlerStopWritesTheRemainingPaths hands a running handler one
+// path and stops it. The peering must be in the database when Stop returns.
+func TestPeeringHandlerStopWritesTheRemainingPaths(t *testing.T) {
+	t.Parallel()
+
+	store := &recordingStore{mockStore: newMockStore()}
+	h := NewPeeringHandler(store, logger.New())
+
+	handle(h, time.Now().UTC(), testASNA, testASNB)
+	h.Stop()
+
+	if got := store.callCount(); got != 1 {
+		t.Errorf("want 1 RecordPeering call when Stop returns, got %d", got)
+	}
+}
+
 // TestPeeringHandlerCapDropsAndCounts verifies that a full map drops new paths
 // and counts them, while a path already tracked is refreshed rather than
 // dropped.