From 28df9f8d49e0d0194f3bb06de3854d5391f56c73 Mon Sep 17 00:00:00 2001 From: clawbot <35+clawbot@noreply.example.org> Date: Mon, 19 Oct 2026 13:08:48 +0200 Subject: [PATCH] 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/routewatch/app.go | 8 +- .../app_integration_internal_test.go | 125 ++++++++++++++++++ internal/routewatch/peeringhandler.go | 16 ++- .../peeringhandler_internal_test.go | 16 +++ internal/streamer/streamer.go | 44 ++++-- internal/streamer/streamer_internal_test.go | 52 +++++++- 6 files changed, 235 insertions(+), 26 deletions(-) diff --git a/internal/routewatch/app.go b/internal/routewatch/app.go index 9e79f11..3aeb1e3 100644 --- a/internal/routewatch/app.go +++ b/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 diff --git a/internal/routewatch/app_integration_internal_test.go b/internal/routewatch/app_integration_internal_test.go index 13bb7aa..bdab0a4 100644 --- a/internal/routewatch/app_integration_internal_test.go +++ b/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) + } + } +} diff --git a/internal/routewatch/peeringhandler.go b/internal/routewatch/peeringhandler.go index 65ce19a..d30d06b 100644 --- a/internal/routewatch/peeringhandler.go +++ b/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() diff --git a/internal/routewatch/peeringhandler_internal_test.go b/internal/routewatch/peeringhandler_internal_test.go index 32da7e0..13c5d66 100644 --- a/internal/routewatch/peeringhandler_internal_test.go +++ b/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. diff --git a/internal/streamer/streamer.go b/internal/streamer/streamer.go index e54412c..698768f 100644 --- a/internal/streamer/streamer.go +++ b/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. diff --git a/internal/streamer/streamer_internal_test.go b/internal/streamer/streamer_internal_test.go index e56ff91..308044a 100644 --- a/internal/streamer/streamer_internal_test.go +++ b/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) {