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/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) {