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

Shutdown stopped the handlers before the streamer, so the updates still
on the handler queues reached handlers that no longer wrote them. It now
stops the streamer first. Streamer.Stop puts the STATE messages it holds
for a full queue on that queue, closes the queues and returns once every
worker has handed its handler everything on its queue; then each handler
writes what it holds. PeeringHandler.Stop now waits for its loop's final
write instead of racing it.

Streamer.SetURL lets the shutdown test read from a local server.

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
+// goroutine to stop and waiting while it processes the remaining AS paths,
-// any remaining AS paths. This ensures no peering data is lost during
+// so no peering data is lost during shutdown.
-// shutdown.
 func (h *PeeringHandler) Stop() {
 	close(h.stopCh)
-	// Process any remaining peerings synchronously
+	h.wg.Wait()
-	h.processPeerings()
 }
 
-// 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.
+// Stop ends the stream and stops the handler workers. It returns once each
-// It cancels the streaming context, closes all handler queues, and updates
+// worker has handed its handler every message left on its queue and then the
-// the connection status in metrics. This method is safe to call multiple times.
+// 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
+// msg instead. It never waits for room, so a full queue does not hold up the
-// read loop nor Stop.
+// 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
+// soon as there is room, or when Stop is called; while any is still held, the
-// are dropped, so none of them overtakes a held down.
+// 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
+// TestStopQueuesTheHeldStateMessages has the server send two STATE messages
-// messages and hold the response open while the queue of the handler that
+// and an UPDATE 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.
+// wants STATE messages is full, so both are held. Stop must put both on the
-func TestStopReturnsAtOnceWithStateMessagesHeld(t *testing.T) {
+// 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) {