Bound each shutdown stage by the stop deadline (closes #100)

Shutdown passes its context to every stage: the streamer, the four
handlers, the WHOIS fetcher, the database maintainer and the HTTP server
each stop waiting when it ends. A handler then skips its last write and
logs how many updates it did not write, and the HTTP server closes the
connections still open. The database is checkpointed and closed last,
and left open once the deadline has passed, since a write may still
hold it. The stop hook gives Shutdown 8 seconds and fx 9, inside the 10
seconds upaas and docker stop allow, so a stop signal exits 0.

Model: opus-5-5

internal/routewatch/app.go

@@ -137,7 +137,12 @@ func (rw *RouteWatch) Run(ctx context.Context) error {
 	return nil
 }
 
-// Shutdown performs graceful shutdown of all services
+// Shutdown stops all services in order: the streamer, the handlers, which
+// write what they hold, the WHOIS fetcher, the database maintainer and the
+// HTTP server, and last the database, which it checkpoints and closes. Each
+// stage waits at most until ctx ends, so Shutdown returns soon after: a
+// handler logs what it could not write by then, and the HTTP server closes
+// the connections of requests still running.
 func (rw *RouteWatch) Shutdown(ctx context.Context) {
 	rw.mu.Lock()
 	if rw.shutdown {
@@ -150,51 +155,48 @@ func (rw *RouteWatch) Shutdown(ctx context.Context) {
 	rw.mu.Unlock()
 
 	// 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()
+	// every message it queued, or when ctx ends. Then each handler writes
+	// what it holds.
+	rw.streamer.Stop(ctx)
 
 	if rw.dbHandler != nil {
 		rw.logger.Info("Flushing database handler")
-		rw.dbHandler.Stop()
+		rw.dbHandler.Stop(ctx)
 	}
 
 	if rw.peerHandler != nil {
 		rw.logger.Info("Flushing peer handler")
-		rw.peerHandler.Stop()
+		rw.peerHandler.Stop(ctx)
 	}
 
 	if rw.prefixHandler != nil {
 		rw.logger.Info("Flushing prefix handler")
-		rw.prefixHandler.Stop()
+		rw.prefixHandler.Stop(ctx)
 	}
 
 	if rw.peeringHandler != nil {
 		rw.logger.Info("Flushing peering handler")
-		rw.peeringHandler.Stop()
+		rw.peeringHandler.Stop(ctx)
 	}
 
 	// Stop ASN WHOIS fetcher
 	if rw.asnFetcher != nil {
-		rw.asnFetcher.Stop()
+		rw.asnFetcher.Stop(ctx)
 	}
 
 	// Stop database maintainer
 	if rw.dbMaintainer != nil {
-		rw.dbMaintainer.Stop()
+		rw.dbMaintainer.Stop(ctx)
 	}
 
-	// Stop HTTP server with a timeout of its own, however much of ctx is left
-	const serverStopTimeout = 5 * time.Second
-
-	stopCtx, cancel := context.WithTimeout(
-		context.WithoutCancel(ctx), serverStopTimeout)
-	defer cancel()
-
-	err := rw.server.Stop(stopCtx)
+	err := rw.server.Stop(ctx)
 	if err != nil {
-		rw.logger.Error("Failed to stop HTTP server gracefully", "error", err)
+		rw.logger.Error("Stopped the HTTP server before every request finished",
+			"error", err)
 	}
 
+	rw.closeDatabase(ctx)
+
 	// Log final metrics
 	metrics := rw.streamer.GetMetrics()
 	rw.logger.Info("Final metrics",
@@ -206,6 +208,27 @@ func (rw *RouteWatch) Shutdown(ctx context.Context) {
 	)
 }
 
+// closeDatabase checkpoints and closes the database, unless ctx has ended: a
+// write may then still be running, and both would wait for it. SQLite rolls
+// back a write cut off by the exit when the database next opens.
+func (rw *RouteWatch) closeDatabase(ctx context.Context) {
+	if ctx.Err() != nil {
+		rw.logger.Error("Stop deadline passed; the database is not closed")
+
+		return
+	}
+
+	err := rw.db.Checkpoint(ctx)
+	if err != nil {
+		rw.logger.Error("Failed to checkpoint the database", "error", err)
+	}
+
+	err = rw.db.Close()
+	if err != nil {
+		rw.logger.Error("Failed to close the database", "error", err)
+	}
+}
+
 // getModule provides all fx dependencies
 func getModule() []fx.Option {
 	return []fx.Option{

internal/routewatch/app_integration_internal_test.go

@@ -4,8 +4,11 @@ import (
 	"context"
 	"fmt"
 	"io"
+	"net"
 	"net/http"
 	"net/http/httptest"
+	"slices"
+	"strconv"
 	"strings"
 	"sync"
 	"testing"
@@ -643,3 +646,187 @@ func TestShutdownWritesTheUpdatesQueuedInTheStreamer(t *testing.T) {
 		}
 	}
 }
+
+// freePort returns a TCP port that was free a moment ago, for a server the
+// test starts.
+func freePort(t *testing.T) int {
+	t.Helper()
+
+	var listenConfig net.ListenConfig
+
+	listener, err := listenConfig.Listen(t.Context(), "tcp", ":0")
+	if err != nil {
+		t.Fatalf("failed to find a free port: %v", err)
+	}
+	defer func() { _ = listener.Close() }()
+
+	addr, ok := listener.Addr().(*net.TCPAddr)
+	if !ok {
+		t.Fatalf("listener address %v is not a TCP address", listener.Addr())
+	}
+
+	return addr.Port
+}
+
+// TestShutdownReturnsByItsDeadlineWhileAWriteIsBlocked starts the HTTP server
+// and has the AS handler start a write that the store holds until the test
+// ends, past Shutdown's deadline. Shutdown must return within its deadline
+// plus a small margin, with the HTTP server stopped, and must neither
+// checkpoint nor close the database: the real one would wait for the write.
+// Before the fix Shutdown waited for the write.
+func TestShutdownReturnsByItsDeadlineWhileAWriteIsBlocked(t *testing.T) {
+	t.Parallel()
+
+	const (
+		deadline = 100 * time.Millisecond
+		// margin is how long Shutdown may take after its deadline, to stop
+		// the HTTP server.
+		margin = 2 * time.Second
+	)
+
+	blocking := newBlockingStore()
+	defer blocking.releaseWrites()
+
+	store := &closeOrderStore{Store: blocking}
+	log := logger.New()
+	port := freePort(t)
+	cfg := &config.Config{Port: port}
+	s := streamer.New(log, metrics.New())
+
+	rw := New(Dependencies{
+		DB:       store,
+		Streamer: s,
+		Server:   server.New(store, s, log, cfg),
+		Logger:   log,
+		Config:   cfg,
+	})
+
+	err := rw.server.Start(t.Context())
+	if err != nil {
+		t.Fatalf("failed to start the HTTP server: %v", err)
+	}
+
+	rw.dbHandler = NewASHandler(store, log)
+	fill := ristypes.ASPath{Sequence: slices.Repeat([]int{testASNA}, asnBatchSize)}
+	startBlockedWrite(t, blocking, rw.dbHandler,
+		[]*ristypes.RISMessage{{Path: fill}})
+
+	ctx, cancel := context.WithTimeout(t.Context(), deadline)
+	defer cancel()
+
+	shutDown := make(chan struct{})
+
+	go func() {
+		rw.Shutdown(ctx)
+		close(shutDown)
+	}()
+
+	select {
+	case <-shutDown:
+	case <-time.After(deadline + margin):
+		t.Fatalf("Shutdown did not return within %s of its deadline while "+
+			"a write was blocked", margin)
+	}
+
+	var dialer net.Dialer
+
+	addr := net.JoinHostPort("127.0.0.1", strconv.Itoa(port))
+
+	conn, err := dialer.DialContext(t.Context(), "tcp", addr)
+	if err == nil {
+		_ = conn.Close()
+
+		t.Errorf("the HTTP server still accepts connections on %s after "+
+			"Shutdown returned", addr)
+	}
+
+	store.callsMu.Lock()
+	defer store.callsMu.Unlock()
+
+	if slices.Contains(store.calls, callCheckpoint) ||
+		slices.Contains(store.calls, callClose) {
+		t.Errorf("Shutdown made the calls %q past its deadline, want no %s "+
+			"and no %s", store.calls, callCheckpoint, callClose)
+	}
+}
+
+// closeOrderStore is a store that records, in order, the AS writes,
+// checkpoints and closes it is given. It passes the AS writes on to Store.
+type closeOrderStore struct {
+	database.Store
+
+	callsMu sync.Mutex
+	calls   []string
+}
+
+// The calls closeOrderStore records.
+const (
+	callWrite      = "AS write"
+	callCheckpoint = "checkpoint"
+	callClose      = "close"
+)
+
+func (s *closeOrderStore) GetOrCreateASNBatch(asns map[int]time.Time) error {
+	s.record(callWrite)
+
+	return s.Store.GetOrCreateASNBatch(asns)
+}
+
+func (s *closeOrderStore) Checkpoint(_ context.Context) error {
+	s.record(callCheckpoint)
+
+	return nil
+}
+
+func (s *closeOrderStore) Close() error {
+	s.record(callClose)
+
+	return nil
+}
+
+func (s *closeOrderStore) record(call string) {
+	s.callsMu.Lock()
+	defer s.callsMu.Unlock()
+
+	s.calls = append(s.calls, call)
+}
+
+// TestShutdownCheckpointsAndClosesTheDatabaseLast hands the AS handler one
+// message and shuts down with time to spare. The handler's last write must
+// come first, then a checkpoint, then the close. Before the fix Shutdown
+// neither checkpointed nor closed the database.
+func TestShutdownCheckpointsAndClosesTheDatabaseLast(t *testing.T) {
+	t.Parallel()
+
+	const deadline = 10 * time.Second
+
+	store := &closeOrderStore{Store: newMockStore()}
+	log := logger.New()
+	cfg := &config.Config{}
+	s := streamer.New(log, metrics.New())
+
+	rw := New(Dependencies{
+		DB:       store,
+		Streamer: s,
+		Server:   server.New(store, s, log, cfg),
+		Logger:   log,
+		Config:   cfg,
+	})
+	rw.dbHandler = NewASHandler(store, log)
+	rw.dbHandler.HandleMessage(&ristypes.RISMessage{
+		Path: ristypes.ASPath{Sequence: []int{testASNA}},
+	})
+
+	ctx, cancel := context.WithTimeout(t.Context(), deadline)
+	defer cancel()
+
+	rw.Shutdown(ctx)
+
+	store.callsMu.Lock()
+	defer store.callsMu.Unlock()
+
+	want := []string{callWrite, callCheckpoint, callClose}
+	if !slices.Equal(store.calls, want) {
+		t.Errorf("Shutdown made the calls %q, want %q", store.calls, want)
+	}
+}

internal/routewatch/ashandler.go

@@ -1,6 +1,7 @@
 package routewatch
 
 import (
+	"context"
 	"slices"
 	"sync"
 	"time"
@@ -39,18 +40,23 @@ type ASHandler struct {
 	db     database.Store
 	logger *logger.Logger
 
-	// Batching. mu guards batch only: the flush goroutine takes the batch
-	// under mu and writes it after releasing mu, so HandleMessage waits for
-	// a database write only while batch is at asnBatchMaxSize.
+	// Batching. mu guards batch and writing only: the flush goroutine takes
+	// the batch under mu and writes it after releasing mu, so HandleMessage
+	// waits for a database write only while batch is at asnBatchMaxSize.
 	mu    sync.Mutex
 	batch []asnOp
+	// writing is how many operations the flush goroutine took from batch
+	// and has not finished writing.
+	writing int
 	// batchTaken is signalled when the flush goroutine takes batch; its L is
 	// &mu.
 	batchTaken  sync.Cond
 	batchFullCh chan struct{} // wakes the flush goroutine when batch is full
 	lastFlush   time.Time     // used only by the flush goroutine
-	stopCh      chan struct{}
-	wg          sync.WaitGroup
+	// stopCh hands Stop's context to the flush goroutine. It holds one, so
+	// Stop does not wait for a write in progress to hand it over.
+	stopCh chan context.Context
+	done   chan struct{} // closed when the flush goroutine returns
 }
 
 type asnOp struct {
@@ -70,12 +76,12 @@ func NewASHandler(db database.Store, logger *logger.Logger) *ASHandler {
 		batch:       make([]asnOp, 0, asnBatchSize),
 		batchFullCh: make(chan struct{}, 1),
 		lastFlush:   time.Now(),
-		stopCh:      make(chan struct{}),
+		stopCh:      make(chan context.Context, 1),
+		done:        make(chan struct{}),
 	}
 	h.batchTaken.L = &h.mu
 
 	// Start the flush timer goroutine
-	h.wg.Add(1)
 	go h.flushLoop()
 
 	return h
@@ -147,20 +153,32 @@ func (h *ASHandler) HandleMessage(msg *ristypes.RISMessage) {
 	}
 }
 
-// Stop gracefully shuts down the ASHandler by signaling the background flush
-// goroutine to terminate and waiting for it to complete. Any pending ASN
-// operations in the current batch are flushed to the database before Stop
-// returns. This method should be called during application shutdown to ensure
-// no data is lost.
-func (h *ASHandler) Stop() {
-	close(h.stopCh)
-	h.wg.Wait()
+// Stop stops the flush goroutine, which writes the operations the handler holds
+// before it returns, and waits for it until ctx ends. Once ctx has ended the
+// flush goroutine skips that last write, and a write still running is lost
+// when the process exits, so Stop logs how many operations are not written.
+func (h *ASHandler) Stop(ctx context.Context) {
+	h.stopCh <- ctx
+
+	select {
+	case <-h.done:
+	case <-ctx.Done():
+	}
+
+	h.mu.Lock()
+	unwritten := h.writing + len(h.batch)
+	h.mu.Unlock()
+
+	if unwritten > 0 {
+		h.logger.Error("Stop deadline passed before the AS handler wrote "+
+			"its updates", "updates", unwritten)
+	}
 }
 
 // flushLoop runs in a goroutine, the only one that writes the batch: it
 // flushes when the batch is full, periodically, and once more on Stop.
 func (h *ASHandler) flushLoop() {
-	defer h.wg.Done()
+	defer close(h.done)
 
 	ticker := time.NewTicker(asnBatchTimeout)
 	defer ticker.Stop()
@@ -173,9 +191,11 @@ func (h *ASHandler) flushLoop() {
 			}
 		case <-h.batchFullCh:
 			h.flushBatch()
-		case <-h.stopCh:
-			// Final flush
-			h.flushBatch()
+		case ctx := <-h.stopCh:
+			// Final flush, unless the stop deadline has passed
+			if ctx.Err() == nil {
+				h.flushBatch()
+			}
 
 			return
 		}
@@ -183,7 +203,8 @@ func (h *ASHandler) flushLoop() {
 }
 
 // takeBatch returns the pending operations and leaves an empty batch in their
-// place for HandleMessage to fill.
+// place for HandleMessage to fill. They count in writing until flushBatch has
+// written them.
 func (h *ASHandler) takeBatch() []asnOp {
 	h.mu.Lock()
 	defer h.mu.Unlock()
@@ -193,6 +214,7 @@ func (h *ASHandler) takeBatch() []asnOp {
 	}
 
 	batch := h.batch
+	h.writing = len(batch)
 	h.batch = make([]asnOp, 0, asnBatchSize)
 	h.batchTaken.Broadcast()
 
@@ -222,5 +244,9 @@ func (h *ASHandler) flushBatch() {
 			"error", err, "count", len(asnMap))
 	}
 
+	h.mu.Lock()
+	h.writing = 0
+	h.mu.Unlock()
+
 	h.lastFlush = time.Now()
 }

internal/routewatch/asnfetcher.go

@@ -43,7 +43,7 @@ type ASNFetcher struct {
 	logger         *slog.Logger
 	immediateQueue chan int
 	stopCh         chan struct{}
-	wg             sync.WaitGroup
+	done           chan struct{} // closed when the fetcher goroutine returns
 
 	// fetchMu ensures only one fetch runs at a time
 	fetchMu sync.Mutex
@@ -67,6 +67,7 @@ func NewASNFetcher(db database.Store, logger *slog.Logger) *ASNFetcher {
 		logger:          logger.With("component", "asn_fetcher"),
 		immediateQueue:  make(chan int, immediateQueueSize),
 		stopCh:          make(chan struct{}),
+		done:            make(chan struct{}),
 		currentInterval: baseInterval,
 		successTimes:    make([]time.Time, 0),
 		errorTimes:      make([]time.Time, 0),
@@ -76,8 +77,6 @@ func NewASNFetcher(db database.Store, logger *slog.Logger) *ASNFetcher {
 // Start begins the background ASN fetcher goroutine. It runs until Stop:
 // the end of ctx does not stop it.
 func (f *ASNFetcher) Start(ctx context.Context) {
-	f.wg.Add(1)
-
 	go f.run(context.WithoutCancel(ctx))
 
 	f.logger.Info("ASN fetcher started",
@@ -87,11 +86,17 @@ func (f *ASNFetcher) Start(ctx context.Context) {
 	)
 }
 
-// Stop gracefully shuts down the fetcher.
-func (f *ASNFetcher) Stop() {
+// Stop stops the fetcher and waits, until ctx ends, for a lookup in progress to
+// finish.
+func (f *ASNFetcher) Stop(ctx context.Context) {
 	close(f.stopCh)
-	f.wg.Wait()
-	f.logger.Info("ASN fetcher stopped")
+
+	select {
+	case <-f.done:
+		f.logger.Info("ASN fetcher stopped")
+	case <-ctx.Done():
+		f.logger.Warn("Stop deadline passed before the ASN fetcher stopped")
+	}
 }
 
 // QueueImmediate queues an ASN for immediate WHOIS lookup.
@@ -197,7 +202,7 @@ func (f *ASNFetcher) recordFailure() {
 
 // run is the main background loop.
 func (f *ASNFetcher) run(ctx context.Context) {
-	defer f.wg.Done()
+	defer close(f.done)
 
 	timer := time.NewTimer(f.getInterval())
 	defer timer.Stop()

internal/routewatch/asnfetcher_internal_test.go

@@ -155,3 +155,51 @@ func TestFetchAndUpdateMarksAnASWithoutReferral(t *testing.T) {
 			baseInterval)
 	}
 }
+
+// TestASNFetcherStopReturnsByItsDeadline starts a fetcher, has it look up an
+// AS whose WHOIS server does not answer until the test ends, and stops it with
+// a short deadline while that lookup is in progress. Stop must return all the
+// same instead of waiting for the lookup.
+func TestASNFetcherStopReturnsByItsDeadline(t *testing.T) {
+	t.Parallel()
+
+	const stopDeadline = 100 * time.Millisecond
+
+	lookupStarted := make(chan struct{})
+	release := make(chan struct{})
+
+	f := NewASNFetcher(newMockStore(), logger.New().Logger)
+	f.whoisClient.DialContext = func(
+		context.Context, string, string,
+	) (net.Conn, error) {
+		close(lookupStarted)
+		<-release
+
+		return nil, syscall.ECONNREFUSED
+	}
+
+	f.Start(t.Context())
+	f.QueueImmediate(testASNA)
+
+	select {
+	case <-lookupStarted:
+	case <-time.After(handleMessageTimeout):
+		t.Fatal("the fetcher did not start the lookup")
+	}
+
+	ctx, cancel := context.WithTimeout(t.Context(), stopDeadline)
+	defer cancel()
+
+	if !returnsInTime(func() { f.Stop(ctx) }) {
+		t.Fatalf("Stop with a deadline of %s did not return while a lookup "+
+			"was in progress", stopDeadline)
+	}
+
+	close(release)
+
+	select {
+	case <-f.done:
+	case <-time.After(handleMessageTimeout):
+		t.Fatal("the fetcher goroutine did not return after the lookup ended")
+	}
+}

internal/routewatch/batchwrite_internal_test.go

@@ -1,6 +1,10 @@
 package routewatch
 
 import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"log/slog"
 	"slices"
 	"sync"
 	"testing"
@@ -22,8 +26,9 @@ const (
 )
 
 // blockingStore is a mockStore whose batch writes for the AS, peer and prefix
-// handlers wait until the test calls releaseWrites. writeStarted is closed when
-// the first of them starts. It records what each write was given.
+// handlers, delete of a peer's routes and peering write wait until the test
+// calls releaseWrites. writeStarted is closed when the first of them starts. It
+// records what each batch write was given.
 type blockingStore struct {
 	*mockStore
 
@@ -85,6 +90,18 @@ func (s *blockingStore) UpdatePeerBatch(peers map[string]database.PeerUpdate) er
 	return nil
 }
 
+func (s *blockingStore) DeleteLiveRoutesOfPeer(peerIP string) (int, error) {
+	s.waitForRelease()
+
+	return s.mockStore.DeleteLiveRoutesOfPeer(peerIP)
+}
+
+func (s *blockingStore) RecordPeering(asA, asB int, ts time.Time) error {
+	s.waitForRelease()
+
+	return s.mockStore.RecordPeering(asA, asB, ts)
+}
+
 func (s *blockingStore) waitForRelease() {
 	s.startOnce.Do(func() { close(s.writeStarted) })
 	<-s.release
@@ -96,6 +113,40 @@ func (s *blockingStore) releaseWrites() {
 	s.releaseOnce.Do(func() { close(s.release) })
 }
 
+// stopLog holds what a handler's Stop logs it did not write.
+//
+//nolint:tagliatelle // field names are the keys the handlers log
+type stopLog struct {
+	Msg       string `json:"msg"`
+	Updates   int    `json:"updates"`
+	PeersDown int    `json:"peers_down"`
+	Paths     int    `json:"paths"`
+}
+
+// newErrorLogger returns a logger that writes only errors, as JSON lines, to
+// logs.
+func newErrorLogger(logs *bytes.Buffer) *logger.Logger {
+	return &logger.Logger{Logger: slog.New(slog.NewJSONHandler(logs,
+		&slog.HandlerOptions{Level: slog.LevelError}))}
+}
+
+// lastLogged decodes the last line in logs, which is the one Stop logs: a
+// handler logs nothing after it.
+func lastLogged(t *testing.T, logs *bytes.Buffer) stopLog {
+	t.Helper()
+
+	lines := bytes.Split(bytes.TrimSpace(logs.Bytes()), []byte("\n"))
+
+	var logged stopLog
+
+	err := json.Unmarshal(lines[len(lines)-1], &logged)
+	if err != nil {
+		t.Fatalf("want the error Stop logs, got %q: %v", logs.String(), err)
+	}
+
+	return logged
+}
+
 // returnsInTime reports whether f returns within handleMessageTimeout.
 func returnsInTime(f func()) bool {
 	done := make(chan struct{})
@@ -151,14 +202,14 @@ func handleWhileAWriteIsBlocked(
 	store *blockingStore,
 	h interface {
 		HandleMessage(msg *ristypes.RISMessage)
-		Stop()
+		Stop(ctx context.Context)
 	},
 	fill []*ristypes.RISMessage,
 	extra *ristypes.RISMessage,
 ) {
 	t.Helper()
 
-	defer h.Stop()
+	defer h.Stop(t.Context())
 	defer store.releaseWrites()
 
 	startBlockedWrite(t, store, h, fill)
@@ -178,14 +229,14 @@ func handleAtMaxBatchSize(
 	store *blockingStore,
 	h interface {
 		HandleMessage(msg *ristypes.RISMessage)
-		Stop()
+		Stop(ctx context.Context)
 	},
 	fill, toMax []*ristypes.RISMessage,
 	extra *ristypes.RISMessage,
 ) {
 	t.Helper()
 
-	defer h.Stop()
+	defer h.Stop(t.Context())
 	defer store.releaseWrites()
 
 	startBlockedWrite(t, store, h, fill)
@@ -217,6 +268,173 @@ func handleAtMaxBatchSize(
 	}
 }
 
+// stopWhileAWriteIsBlocked starts a blocked write with the messages in fill,
+// hands h the message extra, and stops h with a deadline that passes while
+// store holds the write. Stop must return all the same. Then it lets the write
+// finish and waits until done, closed when h's flush goroutine returns. Past
+// the deadline that goroutine skips its last write, so extra is not written,
+// and Stop logs as not written both extra and what the held write took.
+func stopWhileAWriteIsBlocked(
+	t *testing.T,
+	store *blockingStore,
+	h interface {
+		HandleMessage(msg *ristypes.RISMessage)
+		Stop(ctx context.Context)
+	},
+	done <-chan struct{},
+	fill []*ristypes.RISMessage,
+	extra *ristypes.RISMessage,
+) {
+	t.Helper()
+
+	const stopDeadline = 100 * time.Millisecond
+
+	defer store.releaseWrites()
+
+	startBlockedWrite(t, store, h, fill)
+	h.HandleMessage(extra)
+
+	ctx, cancel := context.WithTimeout(t.Context(), stopDeadline)
+	defer cancel()
+
+	if !returnsInTime(func() { h.Stop(ctx) }) {
+		t.Fatalf("Stop with a deadline of %s did not return while a write "+
+			"was blocked", stopDeadline)
+	}
+
+	store.releaseWrites()
+
+	select {
+	case <-done:
+	case <-time.After(handleMessageTimeout):
+		t.Fatal("the flush goroutine did not return after the write finished")
+	}
+}
+
+func TestPrefixHandlerSkipsItsLastWriteAfterTheStopDeadline(t *testing.T) {
+	t.Parallel()
+
+	const extraPrefix = "198.51.100.0/24"
+
+	var logs bytes.Buffer
+
+	store := newBlockingStore()
+	h := NewPrefixHandler(store, newErrorLogger(&logs))
+	fill := slices.Repeat([]string{testV4Prefix}, prefixBatchSize)
+	// The peer's session goes down first, so the write held is the delete of
+	// its routes, and Stop must log that peer down too.
+	stopWhileAWriteIsBlocked(t, store, h, h.done,
+		[]*ristypes.RISMessage{sessionDown(), prefixAnnouncement(fill)},
+		prefixAnnouncement([]string{extraPrefix}))
+
+	store.mu.Lock()
+	defer store.mu.Unlock()
+
+	if store.prefixes[extraPrefix] {
+		t.Errorf("%s was written after the stop deadline", extraPrefix)
+	}
+
+	logged := lastLogged(t, &logs)
+	if logged.Updates != prefixBatchSize+1 || logged.PeersDown != 1 {
+		t.Errorf("Stop logged %d updates and %d peer downs not written, "+
+			"want %d and 1", logged.Updates, logged.PeersDown, prefixBatchSize+1)
+	}
+}
+
+func TestASHandlerSkipsItsLastWriteAfterTheStopDeadline(t *testing.T) {
+	t.Parallel()
+
+	// Each message gives an update for its origin AS and one for each AS
+	// in its path: fill asnBatchSize+1, extra 2.
+	const wantUpdates = asnBatchSize + 1 + 2
+
+	var logs bytes.Buffer
+
+	store := newBlockingStore()
+	h := NewASHandler(store, newErrorLogger(&logs))
+	fill := ristypes.ASPath{Sequence: slices.Repeat([]int{testASNA}, asnBatchSize)}
+	stopWhileAWriteIsBlocked(t, store, h, h.done,
+		[]*ristypes.RISMessage{{Path: fill}},
+		&ristypes.RISMessage{Path: ristypes.ASPath{Sequence: []int{testASNB}}})
+
+	store.mu.Lock()
+	defer store.mu.Unlock()
+
+	if store.asns[testASNB] {
+		t.Errorf("AS%d was written after the stop deadline", testASNB)
+	}
+
+	if logged := lastLogged(t, &logs); logged.Updates != wantUpdates {
+		t.Errorf("Stop logged %d updates not written, want %d",
+			logged.Updates, wantUpdates)
+	}
+}
+
+func TestPeerHandlerSkipsItsLastWriteAfterTheStopDeadline(t *testing.T) {
+	t.Parallel()
+
+	const otherPeerIP = "2001:db8:ffff::2"
+
+	var logs bytes.Buffer
+
+	store := newBlockingStore()
+	h := NewPeerHandler(store, newErrorLogger(&logs))
+	stopWhileAWriteIsBlocked(t, store, h, h.done,
+		slices.Repeat([]*ristypes.RISMessage{{Peer: testPeerIP}}, peerBatchSize),
+		&ristypes.RISMessage{Peer: otherPeerIP})
+
+	store.mu.Lock()
+	defer store.mu.Unlock()
+
+	if store.peers[otherPeerIP] {
+		t.Errorf("peer %s was written after the stop deadline", otherPeerIP)
+	}
+
+	if logged := lastLogged(t, &logs); logged.Updates != peerBatchSize+1 {
+		t.Errorf("Stop logged %d updates not written, want %d",
+			logged.Updates, peerBatchSize+1)
+	}
+}
+
+// TestPeeringHandlerStopCountsTheWriteInProgress hands a running peering
+// handler one AS path and stops it with a deadline that passes while store
+// holds Stop's last write, of that path's peering. Stop must log that path as
+// not written.
+func TestPeeringHandlerStopCountsTheWriteInProgress(t *testing.T) {
+	t.Parallel()
+
+	const stopDeadline = 100 * time.Millisecond
+
+	var logs bytes.Buffer
+
+	store := newBlockingStore()
+	defer store.releaseWrites()
+
+	h := NewPeeringHandler(store, newErrorLogger(&logs))
+	handle(h, time.Now().UTC(), testASNA, testASNB)
+
+	ctx, cancel := context.WithTimeout(t.Context(), stopDeadline)
+	defer cancel()
+
+	if !returnsInTime(func() { h.Stop(ctx) }) {
+		t.Fatalf("Stop with a deadline of %s did not return while a write "+
+			"was blocked", stopDeadline)
+	}
+
+	store.releaseWrites()
+
+	select {
+	case <-h.done:
+	case <-time.After(handleMessageTimeout):
+		t.Fatal("the processing goroutine did not return after the write " +
+			"finished")
+	}
+
+	if logged := lastLogged(t, &logs); logged.Paths != 1 {
+		t.Errorf("Stop logged %d paths not written, want 1", logged.Paths)
+	}
+}
+
 // prefixAnnouncement returns a message in which testPeerIP announces prefixes
 // with origin AS testASNA.
 func prefixAnnouncement(prefixes []string) *ristypes.RISMessage {

internal/routewatch/cli.go

@@ -14,8 +14,14 @@ import (
 )
 
 const (
-	// shutdownTimeout is the maximum time allowed for graceful shutdown
-	shutdownTimeout = 60 * time.Second
+	// shutdownTimeout is how long Shutdown has after a stop signal. upaas,
+	// like docker stop by default, kills the container 10 seconds after the
+	// stop signal (README "Running under upaas").
+	shutdownTimeout = 8 * time.Second
+	// stopTimeout is fx's limit on the stop hook, past which fx exits with
+	// status 1: shutdownTimeout and time for Shutdown to close the HTTP
+	// server after its deadline.
+	stopTimeout = 9 * time.Second
 	// debugInterval is how often to log debug stats
 	debugInterval = 60 * time.Second
 	// bytesPerMB is bytes per megabyte
@@ -63,7 +69,7 @@ func logDebugStats(logger *logger.Logger) {
 func CLIEntry() {
 	app := fx.New(
 		fx.Options(getModule()...),
-		fx.StopTimeout(shutdownTimeout), // Allow 60 seconds for graceful shutdown
+		fx.StopTimeout(stopTimeout),
 		fx.Invoke(func(
 			lc fx.Lifecycle,
 			rw *RouteWatch,
@@ -106,6 +112,10 @@ func CLIEntry() {
 				},
 				OnStop: func(ctx context.Context) error {
 					logger.Info("Shutting down RouteWatch")
+
+					ctx, cancel := context.WithTimeout(ctx, shutdownTimeout)
+					defer cancel()
+
 					rw.Shutdown(ctx)
 
 					return nil

internal/routewatch/dbmaintainer.go

@@ -46,7 +46,7 @@ type DBMaintainer struct {
 	db     database.Store
 	logger *slog.Logger
 	stopCh chan struct{}
-	wg     sync.WaitGroup
+	done   chan struct{} // closed when the maintenance goroutine returns
 
 	// Stats tracking
 	statsMu             sync.Mutex
@@ -67,14 +67,13 @@ func NewDBMaintainer(db database.Store, logger *slog.Logger) *DBMaintainer {
 		db:     db,
 		logger: logger.With("component", "db_maintainer"),
 		stopCh: make(chan struct{}),
+		done:   make(chan struct{}),
 	}
 }
 
 // Start begins the background maintenance goroutine. It runs until Stop:
 // the end of ctx does not stop it.
 func (m *DBMaintainer) Start(ctx context.Context) {
-	m.wg.Add(1)
-
 	go m.run(context.WithoutCancel(ctx))
 
 	m.logger.Info("Database maintainer started",
@@ -84,16 +83,23 @@ func (m *DBMaintainer) Start(ctx context.Context) {
 	)
 }
 
-// Stop gracefully shuts down the maintainer.
-func (m *DBMaintainer) Stop() {
+// Stop stops the maintainer and waits, until ctx ends, for a maintenance task
+// in progress to finish.
+func (m *DBMaintainer) Stop(ctx context.Context) {
 	close(m.stopCh)
-	m.wg.Wait()
-	m.logger.Info("Database maintainer stopped")
+
+	select {
+	case <-m.done:
+		m.logger.Info("Database maintainer stopped")
+	case <-ctx.Done():
+		m.logger.Warn("Stop deadline passed before the database maintainer " +
+			"stopped")
+	}
 }
 
 // run is the main background loop.
 func (m *DBMaintainer) run(ctx context.Context) {
-	defer m.wg.Done()
+	defer close(m.done)
 
 	// Use different timers for each task
 	checkpointTimer := time.NewTimer(checkpointInterval)

internal/routewatch/dbmaintainer_internal_test.go

@@ -6,6 +6,7 @@ import (
 	"encoding/json"
 	"log/slog"
 	"testing"
+	"time"
 
 	"git.eeqj.de/sneak/routewatch/internal/database"
 )
@@ -52,3 +53,68 @@ func TestIncompleteCheckpointNotLoggedAsCompleted(t *testing.T) {
 			m.checkpointCount)
 	}
 }
+
+// heldCheckpointStore is a store whose checkpoint closes started and then
+// waits until release is closed.
+type heldCheckpointStore struct {
+	*mockStore
+
+	started chan struct{}
+	release chan struct{}
+}
+
+// Checkpoint waits until release is closed.
+func (s heldCheckpointStore) Checkpoint(context.Context) error {
+	close(s.started)
+	<-s.release
+
+	return nil
+}
+
+// TestDBMaintainerStopReturnsByItsDeadline stops the maintainer with a short
+// deadline while its goroutine runs a checkpoint that the store holds until
+// the test ends. Stop must return all the same instead of waiting for the
+// checkpoint.
+func TestDBMaintainerStopReturnsByItsDeadline(t *testing.T) {
+	t.Parallel()
+
+	const stopDeadline = 100 * time.Millisecond
+
+	store := heldCheckpointStore{
+		mockStore: newMockStore(),
+		started:   make(chan struct{}),
+		release:   make(chan struct{}),
+	}
+	m := NewDBMaintainer(store, slog.Default())
+
+	// Start would run the first checkpoint only after checkpointInterval.
+	// This goroutine stands for the one Start runs and checkpoints at once.
+	go func() {
+		defer close(m.done)
+
+		m.runCheckpoint(t.Context())
+	}()
+
+	select {
+	case <-store.started:
+	case <-time.After(handleMessageTimeout):
+		t.Fatal("the maintainer did not start the checkpoint")
+	}
+
+	ctx, cancel := context.WithTimeout(t.Context(), stopDeadline)
+	defer cancel()
+
+	if !returnsInTime(func() { m.Stop(ctx) }) {
+		t.Fatalf("Stop with a deadline of %s did not return while a "+
+			"checkpoint was in progress", stopDeadline)
+	}
+
+	close(store.release)
+
+	select {
+	case <-m.done:
+	case <-time.After(handleMessageTimeout):
+		t.Fatal("the maintenance goroutine did not return after the " +
+			"checkpoint ended")
+	}
+}

internal/routewatch/peerhandler.go

@@ -5,6 +5,7 @@ package routewatch
 // peer state in the database.
 
 import (
+	"context"
 	"strconv"
 	"sync"
 	"time"
@@ -42,18 +43,24 @@ type PeerHandler struct {
 	db     database.Store
 	logger *logger.Logger
 
-	// Batching. mu guards peerBatch only: the flush goroutine takes the batch
-	// under mu and writes it after releasing mu, so HandleMessage waits for
-	// a database write only while peerBatch is at peerBatchMaxSize.
+	// Batching. mu guards peerBatch and writing only: the flush goroutine
+	// takes the batch under mu and writes it after releasing mu, so
+	// HandleMessage waits for a database write only while peerBatch is at
+	// peerBatchMaxSize.
 	mu        sync.Mutex
 	peerBatch []peerUpdate
+	// writing is how many updates the flush goroutine took from peerBatch
+	// and has not finished writing.
+	writing int
 	// batchTaken is signalled when the flush goroutine takes peerBatch; its L
 	// is &mu.
 	batchTaken  sync.Cond
 	batchFullCh chan struct{} // wakes the flush goroutine when peerBatch is full
 	lastFlush   time.Time     // used only by the flush goroutine
-	stopCh      chan struct{}
-	wg          sync.WaitGroup
+	// stopCh hands Stop's context to the flush goroutine. It holds one, so
+	// Stop does not wait for a write in progress to hand it over.
+	stopCh chan context.Context
+	done   chan struct{} // closed when the flush goroutine returns
 }
 
 type peerUpdate struct {
@@ -74,12 +81,12 @@ func NewPeerHandler(db database.Store, logger *logger.Logger) *PeerHandler {
 		peerBatch:   make([]peerUpdate, 0, peerBatchSize),
 		batchFullCh: make(chan struct{}, 1),
 		lastFlush:   time.Now(),
-		stopCh:      make(chan struct{}),
+		stopCh:      make(chan context.Context, 1),
+		done:        make(chan struct{}),
 	}
 	h.batchTaken.L = &h.mu
 
 	// Start the flush timer goroutine
-	h.wg.Add(1)
 	go h.flushLoop()
 
 	return h
@@ -144,16 +151,32 @@ func (h *PeerHandler) HandleMessage(msg *ristypes.RISMessage) {
 	}
 }
 
-// Stop gracefully stops the handler and flushes remaining batches
-func (h *PeerHandler) Stop() {
-	close(h.stopCh)
-	h.wg.Wait()
+// Stop stops the flush goroutine, which writes the updates the handler holds
+// before it returns, and waits for it until ctx ends. Once ctx has ended the
+// flush goroutine skips that last write, and a write still running is lost
+// when the process exits, so Stop logs how many updates are not written.
+func (h *PeerHandler) Stop(ctx context.Context) {
+	h.stopCh <- ctx
+
+	select {
+	case <-h.done:
+	case <-ctx.Done():
+	}
+
+	h.mu.Lock()
+	unwritten := h.writing + len(h.peerBatch)
+	h.mu.Unlock()
+
+	if unwritten > 0 {
+		h.logger.Error("Stop deadline passed before the peer handler wrote "+
+			"its updates", "updates", unwritten)
+	}
 }
 
 // flushLoop runs in a goroutine, the only one that writes the batch: it
 // flushes when the batch is full, periodically, and once more on Stop.
 func (h *PeerHandler) flushLoop() {
-	defer h.wg.Done()
+	defer close(h.done)
 
 	ticker := time.NewTicker(peerBatchTimeout)
 	defer ticker.Stop()
@@ -166,9 +189,11 @@ func (h *PeerHandler) flushLoop() {
 			}
 		case <-h.batchFullCh:
 			h.flushBatch()
-		case <-h.stopCh:
-			// Final flush
-			h.flushBatch()
+		case ctx := <-h.stopCh:
+			// Final flush, unless the stop deadline has passed
+			if ctx.Err() == nil {
+				h.flushBatch()
+			}
 
 			return
 		}
@@ -176,7 +201,8 @@ func (h *PeerHandler) flushLoop() {
 }
 
 // takeBatch returns the pending updates and leaves an empty batch in their
-// place for HandleMessage to fill.
+// place for HandleMessage to fill. They count in writing until flushBatch has
+// written them.
 func (h *PeerHandler) takeBatch() []peerUpdate {
 	h.mu.Lock()
 	defer h.mu.Unlock()
@@ -186,6 +212,7 @@ func (h *PeerHandler) takeBatch() []peerUpdate {
 	}
 
 	batch := h.peerBatch
+	h.writing = len(batch)
 	h.peerBatch = make([]peerUpdate, 0, peerBatchSize)
 	h.batchTaken.Broadcast()
 
@@ -227,5 +254,9 @@ func (h *PeerHandler) flushBatch() {
 			"error", err, "count", len(dbPeerMap))
 	}
 
+	h.mu.Lock()
+	h.writing = 0
+	h.mu.Unlock()
+
 	h.lastFlush = time.Now()
 }

internal/routewatch/peeringhandler.go

@@ -1,6 +1,7 @@
 package routewatch
 
 import (
+	"context"
 	"encoding/json"
 	"sync"
 	"time"
@@ -45,9 +46,14 @@ type PeeringHandler struct {
 	mu           sync.Mutex
 	asPaths      map[string]time.Time // key is JSON-encoded AS path
 	droppedPaths int                  // paths dropped because the map was full
+	// writing is how many paths processPeerings took from asPaths and has
+	// not finished writing the peerings of.
+	writing int
 
-	stopCh chan struct{}
-	wg     sync.WaitGroup
+	// stopCh hands Stop's context to the processing goroutine. It holds one,
+	// so Stop does not wait for a write in progress to hand it over.
+	stopCh chan context.Context
+	done   chan struct{} // closed when the processing goroutine returns
 }
 
 // NewPeeringHandler creates and initializes a new PeeringHandler with the
@@ -59,12 +65,11 @@ func NewPeeringHandler(db database.Store, logger *logger.Logger) *PeeringHandler
 		db:      db,
 		logger:  logger,
 		asPaths: make(map[string]time.Time),
-		stopCh:  make(chan struct{}),
+		stopCh:  make(chan context.Context, 1),
+		done:    make(chan struct{}),
 	}
 
 	// Start the periodic processing goroutine
-	h.wg.Add(1)
-
 	go h.processLoop()
 
 	return h
@@ -132,18 +137,33 @@ func (h *PeeringHandler) ProcessPeeringsNow() {
 	h.processPeerings()
 }
 
-// Stop gracefully shuts down the handler by signaling the background
-// 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)
-	h.wg.Wait()
+// Stop stops the processing goroutine, which writes the peerings of the AS
+// paths the handler holds before it returns, and waits for it until ctx ends.
+// Once ctx has ended the processing goroutine skips that last write, and a
+// write still running is lost when the process exits, so Stop logs how many
+// paths' peerings are not written.
+func (h *PeeringHandler) Stop(ctx context.Context) {
+	h.stopCh <- ctx
+
+	select {
+	case <-h.done:
+	case <-ctx.Done():
+	}
+
+	h.mu.Lock()
+	unwritten := h.writing + len(h.asPaths)
+	h.mu.Unlock()
+
+	if unwritten > 0 {
+		h.logger.Error("Stop deadline passed before the peering handler wrote "+
+			"its peerings", "paths", unwritten)
+	}
 }
 
 // processLoop runs periodically to process AS paths into peerings, and once
 // more on Stop.
 func (h *PeeringHandler) processLoop() {
-	defer h.wg.Done()
+	defer close(h.done)
 
 	ticker := time.NewTicker(peeringProcessInterval)
 	defer ticker.Stop()
@@ -152,9 +172,11 @@ func (h *PeeringHandler) processLoop() {
 		select {
 		case <-ticker.C:
 			h.processPeerings()
-		case <-h.stopCh:
-			// Final processing
-			h.processPeerings()
+		case ctx := <-h.stopCh:
+			// Final processing, unless the stop deadline has passed
+			if ctx.Err() == nil {
+				h.processPeerings()
+			}
 
 			return
 		}
@@ -169,6 +191,7 @@ func (h *PeeringHandler) processPeerings() {
 	h.mu.Lock()
 	paths := h.asPaths
 	h.asPaths = make(map[string]time.Time)
+	h.writing = len(paths)
 	dropped := h.droppedPaths
 	h.mu.Unlock()
 
@@ -234,6 +257,10 @@ func (h *PeeringHandler) processPeerings() {
 		}
 	}
 
+	h.mu.Lock()
+	h.writing = 0
+	h.mu.Unlock()
+
 	h.logger.Info("Processed AS peerings",
 		"paths", len(paths),
 		"unique_peerings", len(peerings),

internal/routewatch/peeringhandler_internal_test.go

@@ -53,7 +53,6 @@ func newTestHandler(db database.Store) *PeeringHandler {
 		db:      db,
 		logger:  logger.New(),
 		asPaths: make(map[string]time.Time),
-		stopCh:  make(chan struct{}),
 	}
 }
 
@@ -197,7 +196,7 @@ func TestPeeringHandlerStopWritesTheRemainingPaths(t *testing.T) {
 	h := NewPeeringHandler(store, logger.New())
 
 	handle(h, time.Now().UTC(), testASNA, testASNB)
-	h.Stop()
+	h.Stop(t.Context())
 
 	if got := store.callCount(); got != 1 {
 		t.Errorf("want 1 RecordPeering call when Stop returns, got %d", got)

internal/routewatch/prefixhandler.go

@@ -1,6 +1,7 @@
 package routewatch
 
 import (
+	"context"
 	"net"
 	"net/netip"
 	"slices"
@@ -50,9 +51,9 @@ type PrefixHandler struct {
 	logger  *logger.Logger
 	metrics *metrics.Tracker
 
-	// Batching. mu guards batch, held and peersDown only: the flush goroutine
-	// takes them under mu and writes them after releasing mu, so
-	// HandleMessage waits for a database write only while batch is at
+	// Batching. mu guards batch, held, peersDown, writing and deleting only:
+	// the flush goroutine takes them under mu and writes them after releasing
+	// mu, so HandleMessage waits for a database write only while batch is at
 	// prefixBatchMaxSize.
 	mu    sync.Mutex
 	batch []prefixUpdate
@@ -66,13 +67,21 @@ type PrefixHandler struct {
 	// since the batch was taken, and those whose route delete failed in the
 	// last flush.
 	peersDown []string
+	// writing is how many updates the flush goroutine took and has not
+	// finished writing.
+	writing int
+	// deleting is how many peers from peersDown the flush goroutine took and
+	// has not yet deleted the routes of.
+	deleting int
 	// batchTaken is signalled when the flush goroutine takes batch; its L is
 	// &mu.
 	batchTaken  sync.Cond
 	batchFullCh chan struct{} // wakes the flush goroutine when batch is full
 	lastFlush   time.Time     // used only by the flush goroutine
-	stopCh      chan struct{}
-	wg          sync.WaitGroup
+	// stopCh hands Stop's context to the flush goroutine. It holds one, so
+	// Stop does not wait for a write in progress to hand it over.
+	stopCh chan context.Context
+	done   chan struct{} // closed when the flush goroutine returns
 }
 
 type prefixUpdate struct {
@@ -95,12 +104,12 @@ func NewPrefixHandler(db database.Store, logger *logger.Logger) *PrefixHandler {
 		batch:       make([]prefixUpdate, 0, prefixBatchSize),
 		batchFullCh: make(chan struct{}, 1),
 		lastFlush:   time.Now(),
-		stopCh:      make(chan struct{}),
+		stopCh:      make(chan context.Context, 1),
+		done:        make(chan struct{}),
 	}
 	h.batchTaken.L = &h.mu
 
 	// Start the flush timer goroutine
-	h.wg.Add(1)
 	go h.flushLoop()
 
 	return h
@@ -202,10 +211,39 @@ func (h *PrefixHandler) HandleMessage(msg *ristypes.RISMessage) {
 	}
 }
 
-// Stop gracefully stops the handler and flushes remaining batches
-func (h *PrefixHandler) Stop() {
-	close(h.stopCh)
-	h.wg.Wait()
+// Stop stops the flush goroutine, which writes the updates and peer downs the
+// handler holds before it returns, and waits for it until ctx ends. Once ctx
+// has ended the flush goroutine skips that last write, and a write still
+// running is lost when the process exits. A peer whose route delete fails in
+// the last write is not retried either. Stop logs how many updates are not
+// written and how many peer downs have not had their routes deleted.
+func (h *PrefixHandler) Stop(ctx context.Context) {
+	h.stopCh <- ctx
+
+	select {
+	case <-h.done:
+	case <-ctx.Done():
+	}
+
+	h.mu.Lock()
+	unwritten := h.writing + len(h.batch) + len(h.held)
+	peersDown := h.deleting + len(h.peersDown)
+	h.mu.Unlock()
+
+	if unwritten == 0 && peersDown == 0 {
+		return
+	}
+
+	if ctx.Err() != nil {
+		h.logger.Error("Stop deadline passed before the prefix handler wrote "+
+			"its updates", "updates", unwritten, "peers_down", peersDown)
+
+		return
+	}
+
+	h.logger.Error("The prefix handler stopped without deleting the routes "+
+		"of peers whose session went down", "updates", unwritten,
+		"peers_down", peersDown)
 }
 
 // peerSessionDown notes that the peer's BGP session went down, so that the next
@@ -230,7 +268,7 @@ func (h *PrefixHandler) peerSessionDown(peer string) {
 // flushLoop runs in a goroutine, the only one that writes the batch: it
 // flushes when the batch is full, periodically, and once more on Stop.
 func (h *PrefixHandler) flushLoop() {
-	defer h.wg.Done()
+	defer close(h.done)
 
 	ticker := time.NewTicker(prefixBatchTimeout)
 	defer ticker.Stop()
@@ -243,9 +281,11 @@ func (h *PrefixHandler) flushLoop() {
 			}
 		case <-h.batchFullCh:
 			h.flushBatch()
-		case <-h.stopCh:
-			// Final flush
-			h.flushBatch()
+		case ctx := <-h.stopCh:
+			// Final flush, unless the stop deadline has passed
+			if ctx.Err() == nil {
+				h.flushBatch()
+			}
 
 			return
 		}
@@ -254,7 +294,9 @@ func (h *PrefixHandler) flushLoop() {
 
 // takeBatch returns the pending updates, the held ones first, and the peers in
 // peersDown, and leaves an empty batch, no held updates and no peers in their
-// place for HandleMessage to fill.
+// place for HandleMessage to fill. The updates count in writing until
+// flushBatch has written them, and the peers in deleting until their routes
+// are deleted.
 func (h *PrefixHandler) takeBatch() ([]prefixUpdate, []string) {
 	h.mu.Lock()
 	defer h.mu.Unlock()
@@ -268,6 +310,8 @@ func (h *PrefixHandler) takeBatch() ([]prefixUpdate, []string) {
 		batch = slices.Concat(h.held, h.batch)
 	}
 
+	h.writing = len(batch)
+	h.deleting = len(peersDown)
 	h.held = nil
 	h.batch = make([]prefixUpdate, 0, prefixBatchSize)
 	h.peersDown = nil
@@ -354,6 +398,10 @@ func (h *PrefixHandler) flushBatch() {
 
 	successCount := h.writeBatch(routesToUpsert, routesToDelete, prefixesToUpdate)
 
+	h.mu.Lock()
+	h.writing = 0
+	h.mu.Unlock()
+
 	elapsed := time.Since(startTime)
 	h.logger.Debug("Flushed prefix batch",
 		"batch_size", batchSize,
@@ -369,7 +417,8 @@ func (h *PrefixHandler) flushBatch() {
 // deleteLiveRoutesOfPeers deletes all live routes of each of peers, the peers
 // whose BGP session went down, and returns those whose delete failed. Most
 // downs come from peers that have no routes, so only a delete that removed
-// routes is logged.
+// routes is logged. A peer whose delete succeeded no longer counts in deleting;
+// a failed one counts until keepForNextFlush puts it back in peersDown.
 func (h *PrefixHandler) deleteLiveRoutesOfPeers(peers []string) []string {
 	var failed []string
 
@@ -386,6 +435,10 @@ func (h *PrefixHandler) deleteLiveRoutesOfPeers(peers []string) []string {
 			continue
 		}
 
+		h.mu.Lock()
+		h.deleting--
+		h.mu.Unlock()
+
 		if deleted > 0 {
 			h.logger.Info(
 				"Deleted the live routes of a peer whose session went down",
@@ -439,8 +492,12 @@ func (h *PrefixHandler) keepForNextFlush(
 		kept = slices.Clone(kept[len(kept)-prefixBatchSize:])
 	}
 
-	// takeBatch left held empty, and only this flush fills it.
+	// takeBatch left held empty, and only this flush fills it. The kept
+	// updates now count in held, not in writing, and the failed peers in
+	// peersDown, not in deleting.
 	h.held = kept
+	h.writing = len(toWrite)
+	h.deleting = 0
 
 	for _, peer := range failedPeers {
 		if !slices.Contains(h.peersDown, peer) {

internal/routewatch/prefixhandler_internal_test.go

@@ -1,11 +1,13 @@
 package routewatch
 
 import (
+	"bytes"
 	"context"
 	"errors"
 	"fmt"
 	"net/netip"
 	"slices"
+	"strings"
 	"sync"
 	"testing"
 	"time"
@@ -766,6 +768,53 @@ func TestPrefixHandlerKeepsAtMostABatchOfUpdatesWhileADeleteKeepsFailing(
 	}
 }
 
+// TestPrefixHandlerStopLogsAFailedDeleteWithoutTheDeadline stops a prefix
+// handler, with no deadline, after a peer's session went down and the peer
+// announced a prefix, into a database whose every delete of a peer's routes
+// fails. Stop must log the peer down and the announcement as not written, and
+// must not say that the stop deadline passed.
+func TestPrefixHandlerStopLogsAFailedDeleteWithoutTheDeadline(t *testing.T) {
+	t.Parallel()
+
+	var logs bytes.Buffer
+
+	h := NewPrefixHandler(peerDeleteAlwaysFailsStore{Database: newTestDatabase(t)},
+		newErrorLogger(&logs))
+	h.HandleMessage(sessionDown())
+	h.HandleMessage(peerAnnouncement(testPeerIP, testV4Prefix))
+	h.Stop(t.Context())
+
+	logged := lastLogged(t, &logs)
+	if strings.Contains(logged.Msg, "deadline") {
+		t.Errorf("Stop logged %q although its deadline had not passed",
+			logged.Msg)
+	}
+
+	if logged.Updates != 1 || logged.PeersDown != 1 {
+		t.Errorf("Stop logged %d updates and %d peer downs not written, "+
+			"want 1 and 1", logged.Updates, logged.PeersDown)
+	}
+}
+
+// TestPrefixHandlerStopLogsNothingAfterADeleteThatSucceeds stops a prefix
+// handler, with no deadline, after a peer's session went down, into a database
+// whose delete of the peer's routes succeeds. Stop must log nothing: a peer
+// whose routes are deleted no longer counts as a peer down left undone.
+func TestPrefixHandlerStopLogsNothingAfterADeleteThatSucceeds(t *testing.T) {
+	t.Parallel()
+
+	var logs bytes.Buffer
+
+	h := NewPrefixHandler(newTestDatabase(t), newErrorLogger(&logs))
+	h.HandleMessage(sessionDown())
+	h.Stop(t.Context())
+
+	if logs.Len() != 0 {
+		t.Errorf("Stop logged %q after the routes of the peer whose session "+
+			"went down were deleted", logs.String())
+	}
+}
+
 // newTestDatabase opens a database in a new temporary directory. It is closed
 // when the test ends.
 func newTestDatabase(t *testing.T) *database.Database {

internal/server/server.go

@@ -109,7 +109,9 @@ func (s *Server) Start(ctx context.Context) error {
 	return nil
 }
 
-// Stop gracefully stops the HTTP server
+// Stop stops the HTTP server. It waits, until ctx ends, for the requests in
+// progress to finish, and then closes the connections of those still running
+// and returns ctx's error.
 func (s *Server) Stop(ctx context.Context) error {
 	if s.srv == nil {
 		return nil
@@ -117,7 +119,12 @@ func (s *Server) Stop(ctx context.Context) error {
 
 	s.logger.Info("Stopping HTTP server")
 
-	return s.srv.Shutdown(ctx)
+	err := s.srv.Shutdown(ctx)
+	if err != nil {
+		return errors.Join(err, s.srv.Close())
+	}
+
+	return nil
 }
 
 // SetASNFetcher sets the ASN WHOIS fetcher for on-demand lookups.

internal/server/server_internal_test.go

@@ -0,0 +1,92 @@
+package server
+
+import (
+	"context"
+	"net"
+	"net/http"
+	"testing"
+	"time"
+
+	"git.eeqj.de/sneak/routewatch/internal/logger"
+)
+
+// TestStopClosesTheConnectionOfARequestStillRunning serves a request whose
+// handler does not return until the test ends, and stops the server with a
+// deadline that passes while the request runs. The client must then see its
+// connection closed. Before the fix Stop returned at the deadline and left the
+// connection open until the handler returned.
+func TestStopClosesTheConnectionOfARequestStillRunning(t *testing.T) {
+	t.Parallel()
+
+	const (
+		stopDeadline = 100 * time.Millisecond
+		// waitTimeout is how long the test waits for the request to reach
+		// the handler, and then for the client to see its connection closed.
+		waitTimeout = 2 * time.Second
+		// readHeaderTimeout bounds the test server's read of the request
+		// header, as every http.Server must.
+		readHeaderTimeout = time.Second
+	)
+
+	release := make(chan struct{})
+	defer close(release)
+
+	started := make(chan struct{})
+
+	var listenConfig net.ListenConfig
+
+	listener, err := listenConfig.Listen(t.Context(), "tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatalf("failed to open the test listener: %v", err)
+	}
+
+	s := &Server{logger: logger.New(), srv: &http.Server{
+		Handler: http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
+			close(started)
+			<-release
+		}),
+		ReadHeaderTimeout: readHeaderTimeout,
+	}}
+
+	go func() { _ = s.srv.Serve(listener) }()
+
+	answered := make(chan error, 1)
+
+	go func() {
+		req, err := http.NewRequestWithContext(t.Context(), http.MethodGet,
+			"http://"+listener.Addr().String()+"/", nil)
+		if err != nil {
+			answered <- err
+
+			return
+		}
+
+		resp, err := http.DefaultClient.Do(req)
+		if err == nil {
+			_ = resp.Body.Close()
+		}
+
+		answered <- err
+	}()
+
+	select {
+	case <-started:
+	case <-time.After(waitTimeout):
+		t.Fatal("the request did not reach the handler")
+	}
+
+	ctx, cancel := context.WithTimeout(t.Context(), stopDeadline)
+	defer cancel()
+
+	_ = s.Stop(ctx)
+
+	select {
+	case err := <-answered:
+		if err == nil {
+			t.Error("the request still running at the deadline was answered")
+		}
+	case <-time.After(waitTimeout):
+		t.Fatal("the connection of the request still running at the " +
+			"deadline was open after Stop returned")
+	}
+}

internal/streamer/streamer.go

@@ -253,9 +253,10 @@ func (s *Streamer) Start(ctx context.Context) error {
 
 // 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
+// STATE messages held for it (see queueMessage), or once ctx ends, and then
+// logs how many messages were not handed over. This method is safe to call
 // multiple times.
-func (s *Streamer) Stop() {
+func (s *Streamer) Stop(ctx context.Context) {
 	s.mu.Lock()
 	if s.cancel == nil {
 		// Not started, or already stopped: closing the queues again would panic.
@@ -264,28 +265,56 @@ func (s *Streamer) Stop() {
 		return
 	}
 
-	// Once ctx is cancelled the read loop queues and holds nothing more (see
-	// dispatchMessage).
+	// Once the stream's context is cancelled the read loop queues and holds
+	// nothing more (see dispatchMessage).
 	s.cancel()
 	s.cancel = nil
 
+	notHandedOver := 0
+
 	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
+		// for the worker, which does not take s.mu, to make room, until ctx
+		// ends.
+		for len(info.heldState) > 0 && ctx.Err() == nil {
+			select {
+			case info.queue <- info.heldState[0]:
+				info.heldState = info.heldState[1:]
+			case <-ctx.Done():
+			}
 		}
 
+		notHandedOver += len(info.heldState)
 		info.heldState = nil
 
 		// The worker returns once it has emptied the closed queue.
 		close(info.queue)
 	}
 
+	handlers := s.handlers
 	s.running = false
 	s.mu.Unlock()
 	s.metrics.SetConnected(false)
-	s.workers.Wait()
+
+	workersDone := make(chan struct{})
+
+	go func() {
+		s.workers.Wait()
+		close(workersDone)
+	}()
+
+	select {
+	case <-workersDone:
+	case <-ctx.Done():
+		for _, info := range handlers {
+			notHandedOver += len(info.queue)
+		}
+	}
+
+	if notHandedOver > 0 {
+		s.logger.Error("Stop deadline passed before the handlers took every "+
+			"message", "messages", notHandedOver)
+	}
 }
 
 // IsRunning reports whether the streamer is currently connected and processing

internal/streamer/streamer_internal_test.go

@@ -1,9 +1,12 @@
 package streamer
 
 import (
+	"bytes"
 	"context"
+	"encoding/json"
 	"errors"
 	"io"
+	"log/slog"
 	"net/http"
 	"net/http/httptest"
 	"runtime"
@@ -377,7 +380,7 @@ func TestStopQueuesTheHeldStateMessages(t *testing.T) {
 	stopped := make(chan struct{})
 
 	go func() {
-		s.Stop()
+		s.Stop(t.Context())
 		close(stopped)
 	}()
 
@@ -412,6 +415,68 @@ func TestStopQueuesTheHeldStateMessages(t *testing.T) {
 	}
 }
 
+// TestStopReturnsByItsDeadline stops a streamer whose handler has a full queue
+// and a STATE message held for it, and a worker that stays busy until the test
+// ends, as one does whose handler waits for a database write. Stop must return
+// within its deadline plus a small margin instead of waiting for room on the
+// queue and for the worker, and log as not handed over the held message and
+// every message on the queue.
+func TestStopReturnsByItsDeadline(t *testing.T) {
+	t.Parallel()
+
+	const (
+		deadline = 100 * time.Millisecond
+		margin   = 2 * time.Second
+	)
+
+	var logs bytes.Buffer
+
+	s := newStreamerWithFullStateQueue("")
+	s.logger = &logger.Logger{Logger: slog.New(slog.NewJSONHandler(&logs,
+		&slog.HandlerOptions{Level: slog.LevelError}))}
+	s.handlers[0].heldState = []*ristypes.RISMessage{
+		{Type: ristypes.MessageTypeState},
+	}
+	wantNotHandedOver := 1 + len(s.handlers[0].queue)
+	// Stop cancels the stream Start would have started; there is none here.
+	s.cancel = func() {}
+
+	release := make(chan struct{})
+	defer close(release)
+
+	s.workers.Go(func() { <-release })
+
+	ctx, cancel := context.WithTimeout(t.Context(), deadline)
+	defer cancel()
+
+	stopped := make(chan struct{})
+
+	go func() {
+		s.Stop(ctx)
+		close(stopped)
+	}()
+
+	select {
+	case <-stopped:
+	case <-time.After(deadline + margin):
+		t.Fatalf("Stop did not return within %s of its deadline", margin)
+	}
+
+	var logged struct {
+		Messages int `json:"messages"`
+	}
+
+	err := json.Unmarshal(logs.Bytes(), &logged)
+	if err != nil {
+		t.Fatalf("want the error Stop logs, got %q: %v", logs.String(), err)
+	}
+
+	if logged.Messages != wantNotHandedOver {
+		t.Errorf("Stop logged %d messages not handed over, want %d",
+			logged.Messages, wantNotHandedOver)
+	}
+}
+
 // TestStopBeforeMessageReachesQueues stops the streamer after the read loop
 // has checked for cancellation but before it hands the message to the handler
 // queues. That is the gap a stop from another goroutine can land in, and it
@@ -429,7 +494,7 @@ func TestStopBeforeMessageReachesQueues(t *testing.T) {
 	s := New(logger.New(), metrics.New())
 	s.url = srv.URL
 	s.RegisterHandler(updateHandler{})
-	s.RegisterRawHandler(func(string) { s.Stop() })
+	s.RegisterRawHandler(func(string) { s.Stop(t.Context()) })
 
 	// Start would run the stream in the background, where the test cannot
 	// wait for it. Setting cancel as Start does lets Stop cancel the stream
@@ -443,7 +508,7 @@ func TestStopBeforeMessageReachesQueues(t *testing.T) {
 	}
 
 	// A second Stop must not close the queues again.
-	s.Stop()
+	s.Stop(t.Context())
 }
 
 // TestReconnectWhenStreamEnds has the server end the first response
@@ -676,7 +741,7 @@ func checkReconnect(
 	if err != nil {
 		t.Fatalf("Start returned error: %v", err)
 	}
-	defer s.Stop()
+	defer s.Stop(t.Context())
 
 	select {
 	case <-handler.handled: