Run database maintenance under the database lock (closes #94)

Vacuum, Analyze and Checkpoint now take the database lock the batch
writes hold, so maintenance and batch writes wait for each other
instead of meeting at SQLite's write lock, where whichever waited past
the 5-second busy timeout failed with "database is locked". Analyze
runs ANALYZE through PRAGMA optimize, only on tables without statistics
or whose size changed tenfold, reading about 1000 rows of each index:
even limited, ANALYZE counts every index entry, so an hourly ANALYZE of
every table would hold the lock longer as the tables grow. The
analysis_limit=0 in Initialize, which set no limit, is gone.

Model: opus-5-5

internal/database/database.go

@@ -68,6 +68,18 @@ const (
 	sqliteHardHeapLimitBytes = 1610612736
 )
 
+// Settings of the ANALYZE that Analyze runs.
+const (
+	// analysisLimit is about how many rows of each index ANALYZE reads
+	// (PRAGMA analysis_limit); SQLite estimates the rest.
+	analysisLimit = 1000
+	// optimizeAllTables is the PRAGMA optimize mask that runs ANALYZE (0x2)
+	// on every table that has an index without statistics or whose size has
+	// changed tenfold since it was last analyzed, not only on the tables this
+	// connection has queried (0x10000).
+	optimizeAllTables = 0x10002
+)
+
 // Bounds on the work GetRandomPrefixesByLength does to pick its sample.
 const (
 	// maxPrefixesToShuffle is the most prefixes of one mask length that are
@@ -182,7 +194,6 @@ func (d *Database) Initialize() error {
 	// enough for the whole pool.
 	pragmas := []string{
 		"PRAGMA journal_mode=WAL", // Write-Ahead Logging
-		"PRAGMA analysis_limit=0", // Disable automatic ANALYZE
 		fmt.Sprintf("PRAGMA soft_heap_limit=%d", sqliteSoftHeapLimitBytes),
 		fmt.Sprintf("PRAGMA hard_heap_limit=%d", sqliteHardHeapLimitBytes),
 	}
@@ -2151,14 +2162,18 @@ func (d *Database) GetIPInfoContext(
 	return info, nil
 }
 
-// Vacuum runs incremental vacuum to reclaim unused pages without blocking
-// writes. It frees up to the specified number of pages per call (0 = all
-// freeable pages).
+// Vacuum runs incremental vacuum to return up to 1000 unused pages to the
+// filesystem. Like Analyze and Checkpoint, it holds the database lock, so it
+// and the batch writes wait for each other instead of failing on SQLite's
+// write lock.
 func (d *Database) Vacuum(ctx context.Context) error {
 	// Free up to 1000 pages per call (~4MB with default 4KB page size)
-	// This keeps each vacuum operation quick and non-blocking
+	// This keeps each vacuum operation quick
 	const pagesToFree = 1000
 
+	d.lock("Vacuum")
+	defer d.unlock()
+
 	rows, err := d.db.QueryContext(ctx,
 		fmt.Sprintf("PRAGMA incremental_vacuum(%d)", pagesToFree))
 	if err != nil {
@@ -2180,9 +2195,20 @@ func (d *Database) Vacuum(ctx context.Context) error {
 	return nil
 }
 
-// Analyze runs the SQLite ANALYZE command to update query planner statistics.
+// Analyze updates the query planner statistics. It runs ANALYZE, through
+// PRAGMA optimize, only on the tables that have an index without statistics or
+// whose size has changed tenfold since they were last analyzed: even with
+// analysisLimit, ANALYZE counts every entry of each index, so it takes longer
+// as the tables grow. It holds the database lock.
 func (d *Database) Analyze(ctx context.Context) error {
-	_, err := d.db.ExecContext(ctx, "ANALYZE")
+	d.lock("Analyze")
+	defer d.unlock()
+
+	// analysis_limit is a setting of the connection, so it is set in the same
+	// call as PRAGMA optimize, which runs both statements on one connection.
+	_, err := d.db.ExecContext(ctx, fmt.Sprintf(
+		"PRAGMA analysis_limit=%d; PRAGMA optimize=%#x",
+		analysisLimit, optimizeAllTables))
 	if err != nil {
 		return fmt.Errorf("failed to analyze database: %w", err)
 	}
@@ -2192,8 +2218,11 @@ func (d *Database) Analyze(ctx context.Context) error {
 
 // Checkpoint runs a WAL checkpoint to transfer data from the WAL to the main database.
 // Uses TRUNCATE mode which blocks writers briefly but ensures complete checkpoint,
-// keeping the WAL small for fast read performance.
+// keeping the WAL small for fast read performance. It holds the database lock.
 func (d *Database) Checkpoint(ctx context.Context) error {
+	d.lock("Checkpoint")
+	defer d.unlock()
+
 	_, err := d.db.ExecContext(ctx, "PRAGMA wal_checkpoint(TRUNCATE)")
 	if err != nil {
 		return fmt.Errorf("failed to checkpoint WAL: %w", err)
@@ -2214,21 +2243,26 @@ func (d *Database) Ping(ctx context.Context) error {
 	return nil
 }
 
+// The debug messages lock logs, with the operation's name, when it starts to
+// wait for the database mutex and when it holds it.
+const (
+	lockWaitMessage = "Acquiring database lock"
+	lockHeldMessage = "Database lock acquired"
+)
+
 // lock acquires the database mutex and logs debug information
 func (d *Database) lock(operation string) {
 	// Get caller information
 	_, file, line, _ := runtime.Caller(1)
 	caller := fmt.Sprintf("%s:%d", filepath.Base(file), line)
 
-	d.logger.Debug("Acquiring database lock",
-		"operation", operation, "caller", caller)
+	d.logger.Debug(lockWaitMessage, "operation", operation, "caller", caller)
 
 	d.mu.Lock()
 	d.lockedAt = time.Now()
 	d.lockedBy = fmt.Sprintf("%s (%s)", operation, caller)
 
-	d.logger.Debug("Database lock acquired",
-		"operation", operation, "caller", caller)
+	d.logger.Debug(lockHeldMessage, "operation", operation, "caller", caller)
 }
 
 // unlock releases the database mutex and logs debug information including

internal/database/database_internal_test.go

@@ -5,6 +5,7 @@ import (
 	"database/sql"
 	"errors"
 	"fmt"
+	"log/slog"
 	"net/netip"
 	"slices"
 	"sync"
@@ -47,6 +48,23 @@ const (
 	asnSecondBand = 100
 )
 
+// Parameters for the tests of maintenance and the database lock.
+const (
+	// lockMessageBuffer is how many lock messages a test database's logger
+	// holds before logging waits for the test to read them.
+	lockMessageBuffer = 64
+	// operationKey is the log attribute that names the operation taking the
+	// lock.
+	operationKey = "operation"
+	// testBatchWrite is the operation name the tests take the lock under.
+	testBatchWrite = "test batch write"
+	// lockWaitTimeout is the context timeout of a maintenance call that the
+	// tests let run out while it waits for the lock.
+	lockWaitTimeout = time.Millisecond
+	// lockTestASN is the AS the batch write in the lock tests stores.
+	lockTestASN = 64500
+)
+
 // Parameters for the prefix-length sample tests.
 const (
 	// sampleLimit is how many prefixes the tests ask the sample for.
@@ -474,14 +492,10 @@ func TestVacuumReturnsFreePages(t *testing.T) {
 	}
 }
 
-// TestBatchWriteDuringCheckpoint reproduces issue #25. A batch write reads
-// (SELECT) before it writes (INSERT/UPDATE). Under the default deferred locking
-// the transaction begins as a reader and, when the maintainer's WAL checkpoint
-// holds the write lock, its upgrade to writer fails immediately with "database
-// is locked" without honouring busy_timeout, dropping the batch. With
-// _txlock=immediate the transaction takes the write lock at BEGIN and waits, so
-// no batch is dropped. The checkpoint runs without the Database mutex, exactly
-// as the background maintainer does in production.
+// TestBatchWriteDuringCheckpoint runs checkpoints back to back while batch
+// writes run and checks that every batch write succeeds: a checkpoint holds the
+// database lock as a batch write does, and checkpoints must not keep batch
+// writes from getting it.
 func TestBatchWriteDuringCheckpoint(t *testing.T) {
 	t.Parallel()
 
@@ -530,6 +544,210 @@ func TestBatchWriteDuringCheckpoint(t *testing.T) {
 	wg.Wait()
 }
 
+// lockMessage is a debug message the database lock logged, and the name of the
+// operation that logged it.
+type lockMessage struct {
+	text      string
+	operation string
+}
+
+// lockMessages is a log handler that passes on each message the database lock
+// logs. The lock tests read it to know that an operation waits for the lock or
+// holds it, instead of guessing how long that takes.
+type lockMessages chan lockMessage
+
+// Enabled lets every message through, so the lock's debug messages arrive.
+func (m lockMessages) Enabled(context.Context, slog.Level) bool { return true }
+
+// Handle passes on a message that names an operation and drops the others.
+func (m lockMessages) Handle(_ context.Context, record slog.Record) error {
+	record.Attrs(func(attr slog.Attr) bool {
+		if attr.Key != operationKey {
+			return true
+		}
+
+		m <- lockMessage{text: record.Message, operation: attr.Value.String()}
+
+		return false
+	})
+
+	return nil
+}
+
+// WithAttrs returns m: the tests need no attributes added.
+func (m lockMessages) WithAttrs([]slog.Attr) slog.Handler { return m }
+
+// WithGroup returns m: the tests need no groups.
+func (m lockMessages) WithGroup(string) slog.Handler { return m }
+
+// newDatabaseWithLockMessages opens a new database whose logger passes the lock's
+// messages to the returned channel.
+func newDatabaseWithLockMessages(t *testing.T) (*Database, lockMessages) {
+	t.Helper()
+
+	messages := make(lockMessages, lockMessageBuffer)
+	cfg := &config.Config{StateDir: t.TempDir()}
+
+	db, err := New(cfg, &logger.Logger{Logger: slog.New(messages)})
+	if err != nil {
+		t.Fatalf("failed to create database: %v", err)
+	}
+
+	t.Cleanup(func() { _ = db.Close() })
+
+	return db, messages
+}
+
+// waitForLockMessage waits until operation logs text, and fails the test if
+// returned delivers first: the call running operation returned without logging
+// text.
+func waitForLockMessage(
+	t *testing.T,
+	messages lockMessages,
+	text, operation string,
+	returned <-chan error,
+) {
+	t.Helper()
+
+	for {
+		select {
+		case message := <-messages:
+			if message.text == text && message.operation == operation {
+				return
+			}
+		case err := <-returned:
+			t.Fatalf("%s returned (error: %v) without logging %q",
+				operation, err, text)
+		}
+	}
+}
+
+// TestVacuumWaitsForBatchWrite holds the database lock as a batch write does and
+// checks that Vacuum waits for it and then succeeds. Vacuum ran without the
+// lock, met the batch writes at SQLite's write lock and failed with "database
+// is locked" (https://git.eeqj.de/sneak/routewatch/issues/94).
+func TestVacuumWaitsForBatchWrite(t *testing.T) {
+	t.Parallel()
+
+	db, messages := newDatabaseWithLockMessages(t)
+
+	db.lock(testBatchWrite)
+
+	vacuumed := make(chan error, 1)
+
+	go func() { vacuumed <- db.Vacuum(context.Background()) }()
+
+	waitForLockMessage(t, messages, lockWaitMessage, "Vacuum", vacuumed)
+
+	db.unlock()
+
+	err := <-vacuumed
+	if err != nil {
+		t.Fatalf("Vacuum after the batch write: %v", err)
+	}
+}
+
+// TestBatchWriteWaitsForAnalyze keeps Analyze waiting for SQLite's write lock,
+// which another connection holds, and checks that a batch write started
+// meanwhile waits for Analyze to give up the database lock, that both then
+// succeed, and that Analyze stored statistics. ANALYZE ran without the database
+// lock, and a batch write that waited longer than the busy timeout for it failed
+// with "database is locked".
+func TestBatchWriteWaitsForAnalyze(t *testing.T) {
+	t.Parallel()
+
+	db, messages := newDatabaseWithLockMessages(t)
+	ctx := context.Background()
+	asns := map[int]time.Time{lockTestASN: time.Now().UTC()}
+
+	// A stored AS gives ANALYZE a table to analyze and statistics to write.
+	err := db.GetOrCreateASNBatch(asns)
+	if err != nil {
+		t.Fatalf("failed to store an AS: %v", err)
+	}
+
+	// With _txlock=immediate this transaction takes SQLite's write lock at
+	// BEGIN, so ANALYZE has to wait for it while Analyze holds the database lock.
+	writer, err := db.db.BeginTx(ctx, nil)
+	if err != nil {
+		t.Fatalf("failed to begin the transaction ANALYZE waits for: %v", err)
+	}
+
+	defer func() { _ = writer.Rollback() }()
+
+	analyzed := make(chan error, 1)
+
+	go func() { analyzed <- db.Analyze(ctx) }()
+
+	waitForLockMessage(t, messages, lockHeldMessage, "Analyze", analyzed)
+
+	written := make(chan error, 1)
+
+	go func() { written <- db.GetOrCreateASNBatch(asns) }()
+
+	waitForLockMessage(t, messages, lockWaitMessage, "GetOrCreateASNBatch", written)
+
+	err = writer.Rollback()
+	if err != nil {
+		t.Fatalf("failed to end the transaction ANALYZE waits for: %v", err)
+	}
+
+	err = <-analyzed
+	if err != nil {
+		t.Errorf("Analyze: %v", err)
+	}
+
+	err = <-written
+	if err != nil {
+		t.Errorf("batch write after Analyze: %v", err)
+	}
+
+	var statistics int
+
+	err = db.db.QueryRowContext(ctx,
+		"SELECT count(*) FROM sqlite_stat1 WHERE tbl = 'asns'").Scan(&statistics)
+	if err != nil {
+		t.Fatalf("failed to read the statistics: %v", err)
+	}
+
+	if statistics == 0 {
+		t.Error("Analyze stored no statistics for the asns table")
+	}
+}
+
+// TestCheckpointLockWaitEndsWithContext holds the database lock while the
+// context of a waiting Checkpoint times out, and checks that Checkpoint then
+// returns the context's error and leaves the lock free.
+func TestCheckpointLockWaitEndsWithContext(t *testing.T) {
+	t.Parallel()
+
+	db, messages := newDatabaseWithLockMessages(t)
+
+	db.lock(testBatchWrite)
+
+	ctx, cancel := context.WithTimeout(context.Background(), lockWaitTimeout)
+	defer cancel()
+
+	checkpointed := make(chan error, 1)
+
+	go func() { checkpointed <- db.Checkpoint(ctx) }()
+
+	waitForLockMessage(t, messages, lockWaitMessage, "Checkpoint", checkpointed)
+	<-ctx.Done()
+	db.unlock()
+
+	err := <-checkpointed
+	if !errors.Is(err, context.DeadlineExceeded) {
+		t.Fatalf("Checkpoint = %v, want %v", err, context.DeadlineExceeded)
+	}
+
+	if !db.mu.TryLock() {
+		t.Fatal("Checkpoint kept the database lock after its context ended")
+	}
+
+	db.mu.Unlock()
+}
+
 // TestGetRandomPrefixesByLengthListsPrefixOnce stores one prefix of each
 // family learned from two peers in the same second and checks that the
 // prefix-length sample lists each of them once.

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