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.

internal/streamer/streamer.go

@@ -120,7 +120,7 @@ type handlerInfo struct {
 	queue   chan *ristypes.RISMessage
 	metrics handlerMetrics
 	// heldState holds, in feed order, the STATE messages that found queue
-	// full (see queueMessage). Only the read loop uses it.
+	// full (see queueMessage). Only the read loop uses it, and then Stop.
 	heldState []*ristypes.RISMessage
 }
 
@@ -141,6 +141,7 @@ type Streamer struct {
 	mu           sync.RWMutex
 	cancel       context.CancelFunc
 	running      bool
+	workers      sync.WaitGroup // the handler workers; Stop waits for them
 	metrics      *metrics.Tracker
 	totalDropped atomic.Uint64   // Total dropped messages across all handlers
 	random       *rand.Rand      // Random number generator for backpressure drops
@@ -194,10 +195,18 @@ func (s *Streamer) RegisterHandler(handler MessageHandler) {
 
 	// If we're already running, start a worker for this handler
 	if s.running {
+		s.workers.Add(1)
+
 		go s.runHandlerWorker(info)
 	}
 }
 
+// SetURL makes the streamer read from url instead of RIS Live. Call it before
+// Start. Tests use it to read from a local server.
+func (s *Streamer) SetURL(url string) {
+	s.url = url
+}
+
 // RegisterRawHandler sets a callback to receive raw JSON lines from the stream
 // before they are parsed. Only one raw handler can be registered at a time;
 // subsequent calls will replace the previous handler.
@@ -227,6 +236,8 @@ func (s *Streamer) Start(ctx context.Context) error {
 
 	// Start workers for each handler
 	for _, info := range s.handlers {
+		s.workers.Add(1)
+
 		go s.runHandlerWorker(info)
 	}
 
@@ -240,9 +251,10 @@ func (s *Streamer) Start(ctx context.Context) error {
 	return nil
 }
 
-// Stop halts the streaming connection and shuts down all handler workers.
-// It cancels the streaming context, closes all handler queues, and updates
-// the connection status in metrics. This method is safe to call multiple times.
+// Stop ends the stream and stops the handler workers. It returns once each
+// worker has handed its handler every message left on its queue and then the
+// STATE messages held for it (see queueMessage). This method is safe to call
+// multiple times.
 func (s *Streamer) Stop() {
 	s.mu.Lock()
 	if s.cancel == nil {
@@ -252,16 +264,28 @@ func (s *Streamer) Stop() {
 		return
 	}
 
+	// Once ctx is cancelled the read loop queues and holds nothing more (see
+	// dispatchMessage).
 	s.cancel()
 	s.cancel = nil
-	// Close all handler queues to signal workers to stop
+
 	for _, info := range s.handlers {
+		// The held messages came after everything on the queue. Each waits
+		// for the worker, which does not take s.mu, to make room.
+		for _, msg := range info.heldState {
+			info.queue <- msg
+		}
+
+		info.heldState = nil
+
+		// The worker returns once it has emptied the closed queue.
 		close(info.queue)
 	}
 
 	s.running = false
 	s.mu.Unlock()
 	s.metrics.SetConnected(false)
+	s.workers.Wait()
 }
 
 // IsRunning reports whether the streamer is currently connected and processing
@@ -354,6 +378,8 @@ func (s *Streamer) GetDroppedMessages() uint64 {
 
 // runHandlerWorker processes messages for a specific handler
 func (s *Streamer) runHandlerWorker(info *handlerInfo) {
+	defer s.workers.Done()
+
 	for msg := range info.queue {
 		start := time.Now()
 
@@ -824,16 +850,16 @@ func (s *Streamer) dispatchMessage(
 }
 
 // queueMessage puts msg on the handler's queue and reports false if it dropped
-// msg instead. It never waits for room, so a full queue holds up neither the
-// read loop nor Stop.
+// msg instead. It never waits for room, so a full queue does not hold up the
+// read loop.
 //
 // A message is dropped while the queue is filling up (shouldDropForBackpressure)
 // or full, except a STATE message: one whose state is down is the only notice
 // that all of a peer's routes are gone. A STATE message skips the random drop,
 // and one that finds the queue full is held in info.heldState. Held messages
 // go into the queue, oldest first, before anything else for the handler, as
-// soon as there is room; while any is still held, the handler's other messages
-// are dropped, so none of them overtakes a held down.
+// soon as there is room, or when Stop is called; while any is still held, the
+// handler's other messages are dropped, so none of them overtakes a held down.
 //
 // The held list stays small: a STATE message is a few hundred bytes and the
 // feed sends about ten a second, so a queue full for a minute holds about 600.

internal/streamer/streamer_internal_test.go

@@ -255,6 +255,24 @@ func TestHeldStateMessagesKeepTheirPlaceInTheQueue(t *testing.T) {
 	}
 }
 
+// takeAll takes the messages off queue, as a handler's worker would, and sends
+// them on the returned channel once queue is closed.
+func takeAll(queue chan *ristypes.RISMessage) chan []*ristypes.RISMessage {
+	taken := make(chan []*ristypes.RISMessage, 1)
+
+	go func() {
+		var msgs []*ristypes.RISMessage
+
+		for msg := range queue {
+			msgs = append(msgs, msg)
+		}
+
+		taken <- msgs
+	}()
+
+	return taken
+}
+
 // newStreamerWithFullStateQueue returns a streamer reading from url with one
 // handler, a stateHandler whose queue is full.
 func newStreamerWithFullStateQueue(url string) *Streamer {
@@ -306,20 +324,23 @@ func TestOtherHandlersKeepReceivingWhileAStateMessageIsHeld(t *testing.T) {
 	}
 }
 
-// TestStopReturnsAtOnceWithStateMessagesHeld has the server send two STATE
-// messages and hold the response open while the queue of the handler that
-// wants them is full. Stop must return at once and end the stream.
-func TestStopReturnsAtOnceWithStateMessagesHeld(t *testing.T) {
+// TestStopQueuesTheHeldStateMessages has the server send two STATE messages
+// and an UPDATE and hold the response open while the queue of the handler that
+// wants STATE messages is full, so both are held. Stop must put both on the
+// queue, after the messages already on it, before it closes the queue, and
+// must end the stream.
+func TestStopQueuesTheHeldStateMessages(t *testing.T) {
 	t.Parallel()
 
 	const (
-		lines   = 2
+		held    = 2
+		lines   = held + 1
 		maxWait = 5 * time.Second
 	)
 
 	srv := httptest.NewServer(http.HandlerFunc(
 		func(w http.ResponseWriter, r *http.Request) {
-			_, _ = io.WriteString(w, stateDownLine+stateDownLine)
+			_, _ = io.WriteString(w, stateDownLine+stateDownLine+updateLine)
 
 			flush(w)
 			<-r.Context().Done()
@@ -328,7 +349,7 @@ func TestStopReturnsAtOnceWithStateMessagesHeld(t *testing.T) {
 
 	s := newStreamerWithFullStateQueue(srv.URL)
 	// The raw handler sees each line before the read loop hands it to the
-	// queues, so once it has seen the second line, the first is held.
+	// queues, so once it has seen the UPDATE, both STATE messages are held.
 	linesSeen := make(chan struct{}, lines)
 
 	s.RegisterRawHandler(func(string) { linesSeen <- struct{}{} })
@@ -350,6 +371,9 @@ func TestStopReturnsAtOnceWithStateMessagesHeld(t *testing.T) {
 		}
 	}
 
+	// No worker runs here: the test takes from the full queue in its place.
+	queue := s.handlers[0].queue
+	taken := takeAll(queue)
 	stopped := make(chan struct{})
 
 	go func() {
@@ -364,6 +388,20 @@ func TestStopReturnsAtOnceWithStateMessagesHeld(t *testing.T) {
 			maxWait)
 	}
 
+	// Stop has closed the queue, so takeAll is done or about to be.
+	msgs := <-taken
+	if len(msgs) != cap(queue)+held {
+		t.Fatalf("%d messages came off the queue, want the %d on it and the "+
+			"%d held", len(msgs), cap(queue), held)
+	}
+
+	for _, msg := range msgs[cap(queue):] {
+		if msg.Peer != stateDownPeer || msg.State != ristypes.StateDown {
+			t.Errorf("message after those on the queue has peer %q, state %q; "+
+				"want the down from %s", msg.Peer, msg.State, stateDownPeer)
+		}
+	}
+
 	select {
 	case err := <-streamErr:
 		if !errors.Is(err, context.Canceled) {