diff --git a/README.md b/README.md index 05615d5..838f498 100644 --- a/README.md +++ b/README.md @@ -233,6 +233,15 @@ What happens at each limit: keeps running. - Handler queues: each of the four handler queues holds at most 20,000 messages. When a queue fills, the streamer drops messages instead of blocking. +- Handler batches: while the AS, peer or prefix handler writes a batch to the + database, the updates that arrive collect in its next batch, which stops at + ten times the batch size: 300,000 updates for the AS handler + (`asnBatchMaxSize`, about 10 MiB), 100,000 for the peer handler + (`peerBatchMaxSize`, about 10 MiB) and 250,000 for the prefix handler + (`prefixBatchMaxSize`, about 60 MiB). At that size the handler takes nothing + more off its queue until the write ends, so the queue fills and the streamer + drops messages as above. With the batches being written, the three handlers' + batches take at most about 0.25 GiB of the Go heap's 1.5 GiB. With `DEBUG=routewatch` the daemon logs a `System stats` line every 60 seconds with the goroutine count and Go memory figures. diff --git a/internal/routewatch/ashandler.go b/internal/routewatch/ashandler.go index 3eae6f3..4ff8920 100644 --- a/internal/routewatch/ashandler.go +++ b/internal/routewatch/ashandler.go @@ -19,6 +19,12 @@ const ( // asnBatchSize is the number of ASN operations to batch together asnBatchSize = 30000 + // asnBatchMaxSize is the most ASN operations the batch holds while the + // flush goroutine writes the previous one: ten batches, about 10 MiB + // (README "Memory"). At this size HandleMessage waits until the flush + // goroutine takes the batch. + asnBatchMaxSize = 300000 + // asnBatchTimeout is the maximum time to wait before flushing a batch // DO NOT reduce this timeout - larger batches are more efficient asnBatchTimeout = 2 * time.Second @@ -32,12 +38,16 @@ type ASHandler struct { db database.Store logger *logger.Logger - // Batching - mu sync.Mutex - batch []asnOp - lastFlush time.Time - stopCh chan struct{} - wg sync.WaitGroup + // 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. + mu sync.Mutex + batch []asnOp + batchTaken sync.Cond // signalled when the flush goroutine takes batch; its L is &mu + 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 } type asnOp struct { @@ -52,12 +62,14 @@ type asnOp struct { // background goroutine is terminated. func NewASHandler(db database.Store, logger *logger.Logger) *ASHandler { h := &ASHandler{ - db: db, - logger: logger, - batch: make([]asnOp, 0, asnBatchSize), - lastFlush: time.Now(), - stopCh: make(chan struct{}), + db: db, + logger: logger, + batch: make([]asnOp, 0, asnBatchSize), + batchFullCh: make(chan struct{}, 1), + lastFlush: time.Now(), + stopCh: make(chan struct{}), } + h.batchTaken.L = &h.mu // Start the flush timer goroutine h.wg.Add(1) @@ -85,8 +97,8 @@ func (h *ASHandler) QueueCapacity() int { // HandleMessage processes a RIS Live BGP message by extracting all ASNs from // the AS path and queuing them for batch insertion into the database. The // origin ASN (last element in the path) and all transit ASNs are recorded -// with their associated timestamps. The batch is automatically flushed when -// it reaches the configured size threshold. +// with their associated timestamps. When the batch reaches the configured size +// threshold, the flush goroutine is woken to write it. func (h *ASHandler) HandleMessage(msg *ristypes.RISMessage) { // Use the pre-parsed timestamp timestamp := msg.ParsedTimestamp @@ -100,6 +112,13 @@ func (h *ASHandler) HandleMessage(msg *ristypes.RISMessage) { h.mu.Lock() defer h.mu.Unlock() + // While the batch is at its maximum size, wait until the flush goroutine + // takes it. Meanwhile the handler takes nothing off its queue, which + // fills, and the streamer drops messages for it. + for len(h.batch) >= asnBatchMaxSize { + h.batchTaken.Wait() + } + // Queue origin ASN operation if originASN > 0 { h.batch = append(h.batch, asnOp{ @@ -118,13 +137,19 @@ func (h *ASHandler) HandleMessage(msg *ristypes.RISMessage) { } } - // Check if we need to flush + // A full batch wakes the flush goroutine. If it is writing, the wake-up + // waits in batchFullCh and the batch keeps growing until the write ends + // or the batch reaches asnBatchMaxSize. if len(h.batch) >= asnBatchSize { - h.flushBatchLocked() + select { + case h.batchFullCh <- struct{}{}: + default: // already woken + } } } -// flushLoop runs in a goroutine and periodically flushes batches +// 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() ticker := time.NewTicker(asnBatchTimeout) @@ -133,31 +158,47 @@ func (h *ASHandler) flushLoop() { for { select { case <-ticker.C: - h.mu.Lock() if time.Since(h.lastFlush) >= asnBatchTimeout { - h.flushBatchLocked() + h.flushBatch() } - h.mu.Unlock() + case <-h.batchFullCh: + h.flushBatch() case <-h.stopCh: // Final flush - h.mu.Lock() - h.flushBatchLocked() - h.mu.Unlock() + h.flushBatch() return } } } -// flushBatchLocked flushes the ASN batch to the database (must be called with mutex held) -func (h *ASHandler) flushBatchLocked() { +// takeBatch returns the pending operations and leaves an empty batch in their +// place for HandleMessage to fill. +func (h *ASHandler) takeBatch() []asnOp { + h.mu.Lock() + defer h.mu.Unlock() + if len(h.batch) == 0 { + return nil + } + batch := h.batch + h.batch = make([]asnOp, 0, asnBatchSize) + h.batchTaken.Broadcast() + + return batch +} + +// flushBatch takes the pending batch and writes it to the database, without +// holding the mutex during the write. +func (h *ASHandler) flushBatch() { + batch := h.takeBatch() + if len(batch) == 0 { return } // Process ASNs first (deduped) asnMap := make(map[int]time.Time) - for _, op := range h.batch { + for _, op := range batch { if existing, ok := asnMap[op.number]; !ok || op.timestamp.After(existing) { asnMap[op.number] = op.timestamp } @@ -168,8 +209,6 @@ func (h *ASHandler) flushBatchLocked() { h.logger.Error("Failed to process ASN batch", "error", err, "count", len(asnMap)) } - // Clear batch - h.batch = h.batch[:0] h.lastFlush = time.Now() } diff --git a/internal/routewatch/batchwrite_test.go b/internal/routewatch/batchwrite_test.go new file mode 100644 index 0000000..4551eae --- /dev/null +++ b/internal/routewatch/batchwrite_test.go @@ -0,0 +1,374 @@ +package routewatch + +import ( + "slices" + "sync" + "testing" + "time" + + "git.eeqj.de/sneak/routewatch/internal/database" + "git.eeqj.de/sneak/routewatch/internal/logger" + "git.eeqj.de/sneak/routewatch/internal/ristypes" +) + +const ( + // handleMessageTimeout is how long a test waits for HandleMessage to + // return, or for a write to start, before it fails. + handleMessageTimeout = 2 * time.Second + + // stillWaitingAfter is how long a test watches a HandleMessage that must + // wait. One that does not wait returns well within it. + stillWaitingAfter = 100 * time.Millisecond +) + +// 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. +type blockingStore struct { + *mockStore + + release chan struct{} + releaseOnce sync.Once + writeStarted chan struct{} + startOnce sync.Once + + mu sync.Mutex + prefixes map[string]bool + asns map[int]bool + peers map[string]bool +} + +func newBlockingStore() *blockingStore { + return &blockingStore{ + mockStore: newMockStore(), + release: make(chan struct{}), + writeStarted: make(chan struct{}), + prefixes: make(map[string]bool), + asns: make(map[int]bool), + peers: make(map[string]bool), + } +} + +func (s *blockingStore) waitForRelease() { + s.startOnce.Do(func() { close(s.writeStarted) }) + <-s.release +} + +// releaseWrites lets the held writes, and all later ones, finish. A second call +// does nothing. +func (s *blockingStore) releaseWrites() { + s.releaseOnce.Do(func() { close(s.release) }) +} + +func (s *blockingStore) UpsertLiveRouteBatch(routes []*database.LiveRoute) error { + s.waitForRelease() + s.mu.Lock() + defer s.mu.Unlock() + for _, route := range routes { + s.prefixes[route.Prefix] = true + } + + return nil +} + +func (s *blockingStore) GetOrCreateASNBatch(asns map[int]time.Time) error { + s.waitForRelease() + s.mu.Lock() + defer s.mu.Unlock() + for asn := range asns { + s.asns[asn] = true + } + + return nil +} + +func (s *blockingStore) UpdatePeerBatch(peers map[string]database.PeerUpdate) error { + s.waitForRelease() + s.mu.Lock() + defer s.mu.Unlock() + for peerIP := range peers { + s.peers[peerIP] = true + } + + return nil +} + +// returnsInTime reports whether f returns within handleMessageTimeout. +func returnsInTime(f func()) bool { + done := make(chan struct{}) + go func() { + f() + close(done) + }() + + select { + case <-done: + return true + case <-time.After(handleMessageTimeout): + return false + } +} + +// startBlockedWrite hands h the messages in fill, which fill its batch so that +// the handler writes it to store, and waits until that write starts. store +// holds it until the test calls releaseWrites. +func startBlockedWrite( + t *testing.T, + store *blockingStore, + h interface{ HandleMessage(*ristypes.RISMessage) }, + fill []*ristypes.RISMessage, +) { + t.Helper() + + if !returnsInTime(func() { + for _, msg := range fill { + h.HandleMessage(msg) + } + }) { + t.Fatal("HandleMessage did not return after filling the batch: it waits for the database write") + } + + select { + case <-store.writeStarted: + case <-time.After(handleMessageTimeout): + t.Fatal("the handler did not write its full batch") + } +} + +// handleWhileAWriteIsBlocked starts a blocked write with the messages in fill +// and hands h the message extra while store holds the write. HandleMessage must +// return each time. Then it lets the writes finish and stops h, which writes +// what it holds. +func handleWhileAWriteIsBlocked( + t *testing.T, + store *blockingStore, + h interface { + HandleMessage(*ristypes.RISMessage) + Stop() + }, + fill []*ristypes.RISMessage, + extra *ristypes.RISMessage, +) { + t.Helper() + defer h.Stop() + defer store.releaseWrites() + + startBlockedWrite(t, store, h, fill) + + if !returnsInTime(func() { h.HandleMessage(extra) }) { + t.Fatal("HandleMessage did not return while a database write was blocked") + } +} + +// handleAtMaxBatchSize starts a blocked write with the messages in fill, then, +// while store holds the write, hands h the messages in toMax, which bring its +// next batch to its maximum size, and the message extra. HandleMessage for +// extra must wait until the test lets the write finish, and return after. Then +// it stops h, which writes what it holds. +func handleAtMaxBatchSize( + t *testing.T, + store *blockingStore, + h interface { + HandleMessage(*ristypes.RISMessage) + Stop() + }, + fill, toMax []*ristypes.RISMessage, + extra *ristypes.RISMessage, +) { + t.Helper() + defer h.Stop() + defer store.releaseWrites() + + startBlockedWrite(t, store, h, fill) + for _, msg := range toMax { + h.HandleMessage(msg) + } + + returned := make(chan struct{}) + go func() { + h.HandleMessage(extra) + close(returned) + }() + + select { + case <-returned: + t.Fatal("HandleMessage returned while the batch was at its maximum size and a write was blocked") + case <-time.After(stillWaitingAfter): + } + + store.releaseWrites() + select { + case <-returned: + case <-time.After(handleMessageTimeout): + t.Fatal("HandleMessage did not return after the database write finished") + } +} + +// prefixAnnouncement returns a message in which testPeerIP announces prefixes +// with origin AS testASNA. +func prefixAnnouncement(prefixes []string) *ristypes.RISMessage { + return &ristypes.RISMessage{ + Peer: testPeerIP, + Path: ristypes.ASPath{testASNA}, + Announcements: []ristypes.RISAnnouncement{{NextHop: testPeerIP, Prefixes: prefixes}}, + } +} + +func TestPrefixHandlerHandlesMessagesWhileAWriteIsBlocked(t *testing.T) { + const ( + fillPrefix = "192.0.2.0/24" + extraPrefix = "198.51.100.0/24" + ) + + store := newBlockingStore() + h := NewPrefixHandler(store, logger.New()) + handleWhileAWriteIsBlocked(t, store, h, + []*ristypes.RISMessage{prefixAnnouncement(slices.Repeat([]string{fillPrefix}, prefixBatchSize))}, + prefixAnnouncement([]string{extraPrefix})) + + store.mu.Lock() + defer store.mu.Unlock() + if !store.prefixes[extraPrefix] { + t.Errorf("Stop did not write %s, handled while a write was blocked", extraPrefix) + } +} + +func TestASHandlerHandlesMessagesWhileAWriteIsBlocked(t *testing.T) { + store := newBlockingStore() + h := NewASHandler(store, logger.New()) + handleWhileAWriteIsBlocked(t, store, h, + []*ristypes.RISMessage{{Path: slices.Repeat(ristypes.ASPath{testASNA}, asnBatchSize)}}, + &ristypes.RISMessage{Path: ristypes.ASPath{testASNB}}) + + store.mu.Lock() + defer store.mu.Unlock() + if !store.asns[testASNB] { + t.Errorf("Stop did not write AS%d, handled while a write was blocked", testASNB) + } +} + +func TestPeerHandlerHandlesMessagesWhileAWriteIsBlocked(t *testing.T) { + const otherPeerIP = "2001:db8:ffff::2" + + store := newBlockingStore() + h := NewPeerHandler(store, logger.New()) + handleWhileAWriteIsBlocked(t, store, h, + slices.Repeat([]*ristypes.RISMessage{{Peer: testPeerIP}}, peerBatchSize), + &ristypes.RISMessage{Peer: otherPeerIP}) + + store.mu.Lock() + defer store.mu.Unlock() + if !store.peers[otherPeerIP] { + t.Errorf("Stop did not write peer %s, handled while a write was blocked", otherPeerIP) + } +} + +func TestPrefixHandlerWaitsAtTheMaximumBatchSize(t *testing.T) { + const ( + fillPrefix = "192.0.2.0/24" + extraPrefix = "198.51.100.0/24" + ) + + store := newBlockingStore() + h := NewPrefixHandler(store, logger.New()) + handleAtMaxBatchSize(t, store, h, + []*ristypes.RISMessage{prefixAnnouncement(slices.Repeat([]string{fillPrefix}, prefixBatchSize))}, + []*ristypes.RISMessage{prefixAnnouncement(slices.Repeat([]string{fillPrefix}, prefixBatchMaxSize))}, + prefixAnnouncement([]string{extraPrefix})) + + store.mu.Lock() + defer store.mu.Unlock() + if !store.prefixes[extraPrefix] { + t.Errorf("Stop did not write %s, handled after the batch reached its maximum size", extraPrefix) + } +} + +func TestASHandlerWaitsAtTheMaximumBatchSize(t *testing.T) { + store := newBlockingStore() + h := NewASHandler(store, logger.New()) + handleAtMaxBatchSize(t, store, h, + []*ristypes.RISMessage{{Path: slices.Repeat(ristypes.ASPath{testASNA}, asnBatchSize)}}, + []*ristypes.RISMessage{{Path: slices.Repeat(ristypes.ASPath{testASNA}, asnBatchMaxSize)}}, + &ristypes.RISMessage{Path: ristypes.ASPath{testASNB}}) + + store.mu.Lock() + defer store.mu.Unlock() + if !store.asns[testASNB] { + t.Errorf("Stop did not write AS%d, handled after the batch reached its maximum size", testASNB) + } +} + +func TestPeerHandlerWaitsAtTheMaximumBatchSize(t *testing.T) { + const otherPeerIP = "2001:db8:ffff::2" + + store := newBlockingStore() + h := NewPeerHandler(store, logger.New()) + handleAtMaxBatchSize(t, store, h, + slices.Repeat([]*ristypes.RISMessage{{Peer: testPeerIP}}, peerBatchSize), + slices.Repeat([]*ristypes.RISMessage{{Peer: testPeerIP}}, peerBatchMaxSize), + &ristypes.RISMessage{Peer: otherPeerIP}) + + store.mu.Lock() + defer store.mu.Unlock() + if !store.peers[otherPeerIP] { + t.Errorf("Stop did not write peer %s, handled after the batch reached its maximum size", otherPeerIP) + } +} + +// TestHandlersWakeTheFlushGoroutineWhenTheBatchIsFull hands each handler, built +// without its flush goroutine so that nothing takes a wake-up, one update short +// of a full batch, and a second one exactly a full batch. Only the full batch +// leaves a wake-up in batchFullCh. +func TestHandlersWakeTheFlushGoroutineWhenTheBatchIsFull(t *testing.T) { + tests := []struct { + name string + batchSize int + // handle hands a new handler n updates and returns its batchFullCh. + handle func(n int) chan struct{} + }{ + { + name: "prefix", + batchSize: prefixBatchSize, + handle: func(n int) chan struct{} { + h := &PrefixHandler{batchFullCh: make(chan struct{}, 1)} + h.HandleMessage(prefixAnnouncement(slices.Repeat([]string{"192.0.2.0/24"}, n))) + + return h.batchFullCh + }, + }, + { + name: "AS", + batchSize: asnBatchSize, + handle: func(n int) chan struct{} { + h := &ASHandler{batchFullCh: make(chan struct{}, 1)} + // The origin AS and each AS in the path are one update each. + h.HandleMessage(&ristypes.RISMessage{Path: slices.Repeat(ristypes.ASPath{testASNA}, n-1)}) + + return h.batchFullCh + }, + }, + { + name: "peer", + batchSize: peerBatchSize, + handle: func(n int) chan struct{} { + h := &PeerHandler{batchFullCh: make(chan struct{}, 1)} + for range n { + h.HandleMessage(&ristypes.RISMessage{Peer: testPeerIP}) + } + + return h.batchFullCh + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := len(tt.handle(tt.batchSize - 1)); got != 0 { + t.Errorf("one update short of a full batch left %d wake-ups, want none", got) + } + if got := len(tt.handle(tt.batchSize)); got != 1 { + t.Errorf("a full batch left %d wake-ups, want 1", got) + } + }) + } +} diff --git a/internal/routewatch/peerhandler.go b/internal/routewatch/peerhandler.go index 2f7bf70..9bd1f5d 100644 --- a/internal/routewatch/peerhandler.go +++ b/internal/routewatch/peerhandler.go @@ -23,6 +23,12 @@ const ( // peerBatchSize is the number of peer updates to batch together peerBatchSize = 10000 + // peerBatchMaxSize is the most peer updates the batch holds while the + // flush goroutine writes the previous one: ten batches, about 10 MiB + // (README "Memory"). At this size HandleMessage waits until the flush + // goroutine takes the batch. + peerBatchMaxSize = 100000 + // peerBatchTimeout is the maximum time to wait before flushing a batch peerBatchTimeout = 2 * time.Second ) @@ -35,12 +41,16 @@ type PeerHandler struct { db database.Store logger *logger.Logger - // Batching - mu sync.Mutex - peerBatch []peerUpdate - lastFlush time.Time - stopCh chan struct{} - wg sync.WaitGroup + // 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. + mu sync.Mutex + peerBatch []peerUpdate + batchTaken sync.Cond // signalled when the flush goroutine takes peerBatch; its L is &mu + 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 } type peerUpdate struct { @@ -56,12 +66,14 @@ type peerUpdate struct { // should be stopped by calling Stop when it is no longer needed. func NewPeerHandler(db database.Store, logger *logger.Logger) *PeerHandler { h := &PeerHandler{ - db: db, - logger: logger, - peerBatch: make([]peerUpdate, 0, peerBatchSize), - lastFlush: time.Now(), - stopCh: make(chan struct{}), + db: db, + logger: logger, + peerBatch: make([]peerUpdate, 0, peerBatchSize), + batchFullCh: make(chan struct{}, 1), + lastFlush: time.Now(), + stopCh: make(chan struct{}), } + h.batchTaken.L = &h.mu // Start the flush timer goroutine h.wg.Add(1) @@ -87,8 +99,8 @@ func (h *PeerHandler) QueueCapacity() int { // HandleMessage processes a RIS message to track peer information. // It extracts the peer IP address and ASN from the message and adds -// the update to an internal batch. When the batch reaches peerBatchSize -// or the batch timeout expires, the batch is flushed to the database. +// the update to an internal batch. The flush goroutine writes the batch to the +// database when it reaches peerBatchSize or the batch timeout expires. func (h *PeerHandler) HandleMessage(msg *ristypes.RISMessage) { // Parse peer ASN from string peerASN := 0 @@ -101,6 +113,13 @@ func (h *PeerHandler) HandleMessage(msg *ristypes.RISMessage) { h.mu.Lock() defer h.mu.Unlock() + // While the batch is at its maximum size, wait until the flush goroutine + // takes it. Meanwhile the handler takes nothing off its queue, which + // fills, and the streamer drops messages for it. + for len(h.peerBatch) >= peerBatchMaxSize { + h.batchTaken.Wait() + } + // Add to batch h.peerBatch = append(h.peerBatch, peerUpdate{ peerIP: msg.Peer, @@ -109,13 +128,19 @@ func (h *PeerHandler) HandleMessage(msg *ristypes.RISMessage) { timestamp: msg.ParsedTimestamp, }) - // Check if we need to flush + // A full batch wakes the flush goroutine. If it is writing, the wake-up + // waits in batchFullCh and the batch keeps growing until the write ends + // or the batch reaches peerBatchMaxSize. if len(h.peerBatch) >= peerBatchSize { - h.flushBatchLocked() + select { + case h.batchFullCh <- struct{}{}: + default: // already woken + } } } -// flushLoop runs in a goroutine and periodically flushes batches +// 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() ticker := time.NewTicker(peerBatchTimeout) @@ -124,31 +149,47 @@ func (h *PeerHandler) flushLoop() { for { select { case <-ticker.C: - h.mu.Lock() if time.Since(h.lastFlush) >= peerBatchTimeout { - h.flushBatchLocked() + h.flushBatch() } - h.mu.Unlock() + case <-h.batchFullCh: + h.flushBatch() case <-h.stopCh: // Final flush - h.mu.Lock() - h.flushBatchLocked() - h.mu.Unlock() + h.flushBatch() return } } } -// flushBatchLocked flushes the peer batch to the database (must be called with mutex held) -func (h *PeerHandler) flushBatchLocked() { +// takeBatch returns the pending updates and leaves an empty batch in their +// place for HandleMessage to fill. +func (h *PeerHandler) takeBatch() []peerUpdate { + h.mu.Lock() + defer h.mu.Unlock() + if len(h.peerBatch) == 0 { + return nil + } + batch := h.peerBatch + h.peerBatch = make([]peerUpdate, 0, peerBatchSize) + h.batchTaken.Broadcast() + + return batch +} + +// flushBatch takes the pending batch and writes it to the database, without +// holding the mutex during the write. +func (h *PeerHandler) flushBatch() { + batch := h.takeBatch() + if len(batch) == 0 { return } // Deduplicate by peer IP, keeping the latest update for each peer peerMap := make(map[string]peerUpdate) - for _, update := range h.peerBatch { + for _, update := range batch { if existing, ok := peerMap[update.peerIP]; !ok || update.timestamp.After(existing.timestamp) { peerMap[update.peerIP] = update } @@ -170,8 +211,6 @@ func (h *PeerHandler) flushBatchLocked() { h.logger.Error("Failed to process peer batch", "error", err, "count", len(dbPeerMap)) } - // Clear batch - h.peerBatch = h.peerBatch[:0] h.lastFlush = time.Now() } diff --git a/internal/routewatch/prefixhandler.go b/internal/routewatch/prefixhandler.go index 9833b37..5c76410 100644 --- a/internal/routewatch/prefixhandler.go +++ b/internal/routewatch/prefixhandler.go @@ -24,6 +24,12 @@ const ( // prefixBatchSize is the number of prefix updates to batch together prefixBatchSize = 25000 + // prefixBatchMaxSize is the most prefix updates the batch holds while the + // flush goroutine writes the previous one: ten batches, about 60 MiB + // (README "Memory"). At this size HandleMessage waits until the flush + // goroutine takes the batch. + prefixBatchMaxSize = 250000 + // prefixBatchTimeout is the maximum time to wait before flushing a batch // DO NOT reduce this timeout - larger batches are more efficient prefixBatchTimeout = 1 * time.Second @@ -40,12 +46,16 @@ type PrefixHandler struct { logger *logger.Logger metrics *metrics.Tracker - // Batching - mu sync.Mutex - batch []prefixUpdate - lastFlush time.Time - stopCh chan struct{} - wg sync.WaitGroup + // 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 prefixBatchMaxSize. + mu sync.Mutex + batch []prefixUpdate + batchTaken sync.Cond // signalled when the flush goroutine takes batch; its L is &mu + 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 } type prefixUpdate struct { @@ -60,12 +70,14 @@ type prefixUpdate struct { // NewPrefixHandler creates a new batched prefix tracking handler func NewPrefixHandler(db database.Store, logger *logger.Logger) *PrefixHandler { h := &PrefixHandler{ - db: db, - logger: logger, - batch: make([]prefixUpdate, 0, prefixBatchSize), - lastFlush: time.Now(), - stopCh: make(chan struct{}), + db: db, + logger: logger, + batch: make([]prefixUpdate, 0, prefixBatchSize), + batchFullCh: make(chan struct{}, 1), + lastFlush: time.Now(), + stopCh: make(chan struct{}), } + h.batchTaken.L = &h.mu // Start the flush timer goroutine h.wg.Add(1) @@ -105,6 +117,13 @@ func (h *PrefixHandler) HandleMessage(msg *ristypes.RISMessage) { h.mu.Lock() defer h.mu.Unlock() + // While the batch is at its maximum size, wait until the flush goroutine + // takes it. Meanwhile the handler takes nothing off its queue, which + // fills, and the streamer drops messages for it. + for len(h.batch) >= prefixBatchMaxSize { + h.batchTaken.Wait() + } + // Withdrawals are queued before announcements, because the batch keeps // each peer's last update for a prefix: a message that lists a prefix as // both withdrawn and announced keeps the route (RFC 4271 section 4.3). @@ -140,13 +159,19 @@ func (h *PrefixHandler) HandleMessage(msg *ristypes.RISMessage) { } } - // Check if we need to flush + // A full batch wakes the flush goroutine. If it is writing, the wake-up + // waits in batchFullCh and the batch keeps growing until the write ends + // or the batch reaches prefixBatchMaxSize. if len(h.batch) >= prefixBatchSize { - h.flushBatchLocked() + select { + case h.batchFullCh <- struct{}{}: + default: // already woken + } } } -// flushLoop runs in a goroutine and periodically flushes batches +// 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() ticker := time.NewTicker(prefixBatchTimeout) @@ -155,30 +180,46 @@ func (h *PrefixHandler) flushLoop() { for { select { case <-ticker.C: - h.mu.Lock() if time.Since(h.lastFlush) >= prefixBatchTimeout { - h.flushBatchLocked() + h.flushBatch() } - h.mu.Unlock() + case <-h.batchFullCh: + h.flushBatch() case <-h.stopCh: // Final flush - h.mu.Lock() - h.flushBatchLocked() - h.mu.Unlock() + h.flushBatch() return } } } -// flushBatchLocked flushes the prefix batch to the database (must be called with mutex held) -func (h *PrefixHandler) flushBatchLocked() { +// takeBatch returns the pending updates and leaves an empty batch in their +// place for HandleMessage to fill. +func (h *PrefixHandler) takeBatch() []prefixUpdate { + h.mu.Lock() + defer h.mu.Unlock() + if len(h.batch) == 0 { + return nil + } + batch := h.batch + h.batch = make([]prefixUpdate, 0, prefixBatchSize) + h.batchTaken.Broadcast() + + return batch +} + +// flushBatch takes the pending batch and writes it to the database, without +// holding the mutex during the write. +func (h *PrefixHandler) flushBatch() { + batch := h.takeBatch() + if len(batch) == 0 { return } startTime := time.Now() - batchSize := len(h.batch) + batchSize := len(batch) // For each prefix and peer, keep only the last update the feed delivered. // Updates from different peers for the same prefix are all kept. Arrival @@ -186,7 +227,7 @@ func (h *PrefixHandler) flushBatchLocked() { // updates within one second carry the same one. type prefixAndPeer struct{ prefix, peer string } lastUpdates := make(map[prefixAndPeer]prefixUpdate) - for _, update := range h.batch { + for _, update := range batch { lastUpdates[prefixAndPeer{update.prefix, update.peer}] = update } @@ -260,8 +301,6 @@ func (h *PrefixHandler) flushBatchLocked() { "duration_ms", elapsed.Milliseconds(), ) - // Clear batch - h.batch = h.batch[:0] h.lastFlush = time.Now() } diff --git a/internal/routewatch/prefixhandler_test.go b/internal/routewatch/prefixhandler_test.go index 54c33dd..c1ce235 100644 --- a/internal/routewatch/prefixhandler_test.go +++ b/internal/routewatch/prefixhandler_test.go @@ -31,11 +31,6 @@ func TestPrefixHandlerStoresPrefixesTheIPLookupFinds(t *testing.T) { // Built without NewPrefixHandler's flush timer, so the test flushes itself. h := &PrefixHandler{db: db, logger: logger.New()} - flush := func() { - h.mu.Lock() - defer h.mu.Unlock() - h.flushBatchLocked() - } ts := time.Date(2026, 1, 2, 3, 4, 5, 0, time.UTC) h.HandleMessage(&ristypes.RISMessage{ @@ -47,7 +42,7 @@ func TestPrefixHandlerStoresPrefixesTheIPLookupFinds(t *testing.T) { Prefixes: []string{"2001:db8:1::/48", "192.0.2.0/24"}, }}, }) - flush() + h.flushBatch() for ip, want := range map[string]string{ "2001:db8:1::1": "2001:db8:1::/48", @@ -67,7 +62,7 @@ func TestPrefixHandlerStoresPrefixesTheIPLookupFinds(t *testing.T) { ParsedTimestamp: ts.Add(time.Minute), Withdrawals: []string{"2001:db8:1:0:0:0:0:0/48", "192.0.2.0/24"}, }) - flush() + h.flushBatch() for _, ip := range []string{"2001:db8:1::1", "192.0.2.1"} { if info, err := db.GetIPInfo(ip); !errors.Is(err, database.ErrNoRoute) { @@ -165,9 +160,7 @@ func TestPrefixHandlerAppliesEachPeersLastUpdateInABatch(t *testing.T) { for _, msg := range batch { h.HandleMessage(msg) } - h.mu.Lock() - h.flushBatchLocked() - h.mu.Unlock() + h.flushBatch() } ctx := context.Background() @@ -265,9 +258,7 @@ func TestPrefixHandlerKeepsOneRoutePerPrefixAndPeer(t *testing.T) { h := &PrefixHandler{db: db, logger: logger.New()} for _, msg := range tt.messages { h.HandleMessage(msg) - h.mu.Lock() - h.flushBatchLocked() - h.mu.Unlock() + h.flushBatch() } for _, p := range []string{prefix, otherPrefix} {