diff --git a/README.md b/README.md index 05ca2a5..dfaa156 100644 --- a/README.md +++ b/README.md @@ -191,7 +191,11 @@ routewatch/ batches for performance - **Backpressure**: Probabilistic message dropping when queues exceed 50% capacity -- **Graceful Shutdown**: 60-second timeout, flushes all pending batches +- **Graceful Shutdown**: on SIGTERM or SIGINT the streamer stops, the handlers + write the batches they hold, the HTTP server stops and the database is + checkpointed and closed, all within 8 seconds, and the daemon exits with + status 0. What is not written by then is lost, and the handler logs how many + updates (see "Running under upaas"). - **Reconnection**: Any end of the stream other than shutdown, including RIS Live ending the response normally, marks the feed disconnected and reconnects with exponential backoff (5s-320s). The backoff resets after a stream that @@ -345,6 +349,19 @@ What the [upaas](https://git.eeqj.de/sneak/upaas) app needs: `/.well-known/healthcheck.json` on the container port. upaas reads the container's health 60 seconds after a deploy and fails the deploy unless it is `healthy`. +- Stop grace: the container needs 9 seconds between SIGTERM and SIGKILL. upaas + gives 10 (`stopTimeoutSeconds` in its `internal/docker/client.go`), as + `docker stop` does by default. The daemon allows 8 seconds for the handlers to + write what they hold and for the requests in progress to finish, then stops + the HTTP server, closes the database and exits with status 0. A longer grace + changes nothing. What the handlers have not written after 8 seconds is lost, + and each logs how many updates: when the prefix handler has fallen behind the + feed, its batch holds up to 250,000 updates, which take about 20 seconds to + write, so a stop then loses them, withdrawals and peer downs included. A route + withdrawn in them stays in the live route table until its peer announces or + withdraws the prefix again or its session goes down. When the 8 seconds run + out the database is not closed, and SQLite rolls back a write the exit cut off + when the database next opens. ## Development diff --git a/cmd/streamdumper/main.go b/cmd/streamdumper/main.go index 95a0722..57ff529 100644 --- a/cmd/streamdumper/main.go +++ b/cmd/streamdumper/main.go @@ -35,7 +35,7 @@ func main() { if err != nil { log.Fatal("Failed to start streamer:", err) } - defer s.Stop() + defer s.Stop(context.Background()) // Set up context with cancellation ctx, cancel := context.WithCancel(context.Background()) diff --git a/internal/routewatch/app.go b/internal/routewatch/app.go index 3aeb1e3..a4d1d5b 100644 --- a/internal/routewatch/app.go +++ b/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{ diff --git a/internal/routewatch/app_integration_internal_test.go b/internal/routewatch/app_integration_internal_test.go index 8f094ca..1cb5f07 100644 --- a/internal/routewatch/app_integration_internal_test.go +++ b/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,175 @@ 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. 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 + ) + + store := newBlockingStore() + defer store.releaseWrites() + + 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, store, 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) + } +} + +// closeOrderStore is a mockStore that records, in order, the AS writes, +// checkpoints and closes it is given. +type closeOrderStore struct { + *mockStore + + 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.mockStore.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{mockStore: 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) + } +} diff --git a/internal/routewatch/ashandler.go b/internal/routewatch/ashandler.go index 1484165..3991c80 100644 --- a/internal/routewatch/ashandler.go +++ b/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() } diff --git a/internal/routewatch/asnfetcher.go b/internal/routewatch/asnfetcher.go index 5dbf1eb..46a1765 100644 --- a/internal/routewatch/asnfetcher.go +++ b/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() diff --git a/internal/routewatch/batchwrite_internal_test.go b/internal/routewatch/batchwrite_internal_test.go index 82ebeb0..40a2808 100644 --- a/internal/routewatch/batchwrite_internal_test.go +++ b/internal/routewatch/batchwrite_internal_test.go @@ -1,6 +1,7 @@ package routewatch import ( + "context" "slices" "sync" "testing" @@ -151,14 +152,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 +179,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 +218,105 @@ 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. +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" + + store := newBlockingStore() + h := NewPrefixHandler(store, logger.New()) + fill := slices.Repeat([]string{testV4Prefix}, prefixBatchSize) + stopWhileAWriteIsBlocked(t, store, h, h.done, + []*ristypes.RISMessage{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) + } +} + +func TestASHandlerSkipsItsLastWriteAfterTheStopDeadline(t *testing.T) { + t.Parallel() + + store := newBlockingStore() + h := NewASHandler(store, logger.New()) + 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) + } +} + +func TestPeerHandlerSkipsItsLastWriteAfterTheStopDeadline(t *testing.T) { + t.Parallel() + + const otherPeerIP = "2001:db8:ffff::2" + + store := newBlockingStore() + h := NewPeerHandler(store, logger.New()) + 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) + } +} + // prefixAnnouncement returns a message in which testPeerIP announces prefixes // with origin AS testASNA. func prefixAnnouncement(prefixes []string) *ristypes.RISMessage { diff --git a/internal/routewatch/cli.go b/internal/routewatch/cli.go index ea7d475..1bf6907 100644 --- a/internal/routewatch/cli.go +++ b/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 diff --git a/internal/routewatch/dbmaintainer.go b/internal/routewatch/dbmaintainer.go index 6ecdf0b..abd075c 100644 --- a/internal/routewatch/dbmaintainer.go +++ b/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) diff --git a/internal/routewatch/peerhandler.go b/internal/routewatch/peerhandler.go index dcb1bc4..7854695 100644 --- a/internal/routewatch/peerhandler.go +++ b/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() } diff --git a/internal/routewatch/peeringhandler.go b/internal/routewatch/peeringhandler.go index d30d06b..f304082 100644 --- a/internal/routewatch/peeringhandler.go +++ b/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), diff --git a/internal/routewatch/peeringhandler_internal_test.go b/internal/routewatch/peeringhandler_internal_test.go index 13c5d66..b11737a 100644 --- a/internal/routewatch/peeringhandler_internal_test.go +++ b/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) diff --git a/internal/routewatch/prefixhandler.go b/internal/routewatch/prefixhandler.go index 4467e2a..8465c35 100644 --- a/internal/routewatch/prefixhandler.go +++ b/internal/routewatch/prefixhandler.go @@ -1,6 +1,7 @@ package routewatch import ( + "context" "net" "net/netip" "slices" @@ -50,8 +51,8 @@ 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 + // Batching. mu guards batch, held, peersDown and writing 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 @@ -66,13 +67,18 @@ 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 // 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 +101,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 +208,28 @@ 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, so 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 := len(h.peersDown) + h.mu.Unlock() + + if unwritten > 0 || peersDown > 0 { + h.logger.Error("Stop deadline passed before the prefix handler wrote "+ + "its updates", "updates", unwritten, "peers_down", peersDown) + } } // peerSessionDown notes that the peer's BGP session went down, so that the next @@ -230,7 +254,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 +267,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 +280,8 @@ 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. func (h *PrefixHandler) takeBatch() ([]prefixUpdate, []string) { h.mu.Lock() defer h.mu.Unlock() @@ -268,6 +295,7 @@ func (h *PrefixHandler) takeBatch() ([]prefixUpdate, []string) { batch = slices.Concat(h.held, h.batch) } + h.writing = len(batch) h.held = nil h.batch = make([]prefixUpdate, 0, prefixBatchSize) h.peersDown = nil @@ -354,6 +382,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, @@ -439,8 +471,10 @@ 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. h.held = kept + h.writing = len(toWrite) for _, peer := range failedPeers { if !slices.Contains(h.peersDown, peer) { diff --git a/internal/server/server.go b/internal/server/server.go index d709b7a..fd3a9ae 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -103,7 +103,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 @@ -111,7 +113,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. diff --git a/internal/streamer/streamer.go b/internal/streamer/streamer.go index 698768f..2df25f5 100644 --- a/internal/streamer/streamer.go +++ b/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 diff --git a/internal/streamer/streamer_internal_test.go b/internal/streamer/streamer_internal_test.go index 308044a..df7fc82 100644 --- a/internal/streamer/streamer_internal_test.go +++ b/internal/streamer/streamer_internal_test.go @@ -377,7 +377,7 @@ func TestStopQueuesTheHeldStateMessages(t *testing.T) { stopped := make(chan struct{}) go func() { - s.Stop() + s.Stop(t.Context()) close(stopped) }() @@ -412,6 +412,48 @@ 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. +func TestStopReturnsByItsDeadline(t *testing.T) { + t.Parallel() + + const ( + deadline = 100 * time.Millisecond + margin = 2 * time.Second + ) + + s := newStreamerWithFullStateQueue("") + s.handlers[0].heldState = []*ristypes.RISMessage{ + {Type: ristypes.MessageTypeState}, + } + // 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) + } +} + // 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 +471,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 +485,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 +718,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: