Write each handler batch from its flush goroutine, outside the mutex (closes #65)

The AS, peer and prefix handlers wrote their batch to the database while
holding the mutex HandleMessage needs, and HandleMessage wrote a full batch
itself. For as long as a write took, the handler took nothing off its queue,
and the streamer dropped updates once the queue was half full.

Now only the flush goroutine writes: it takes the batch under the mutex,
releases it, then writes. A full batch wakes the flush goroutine instead.
Stop still writes what the handler holds.

While a write runs, the next batch has no size limit; the README memory
section says so.

Model: opus-5-5

README.md

@@ -230,6 +230,10 @@ 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 has no size
+  limit. A database that stays slower than the feed grows memory instead of
+  dropping updates.
 
 With `DEBUG=routewatch` the daemon logs a `System stats` line every 60 seconds
 with the goroutine count and Go memory figures.

internal/routewatch/ashandler.go

@@ -32,12 +32,15 @@ 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 never
+	// waits for a database write.
+	mu          sync.Mutex
+	batch       []asnOp
+	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,11 +55,12 @@ 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{}),
 	}
 
 	// Start the flush timer goroutine
@@ -85,8 +89,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
@@ -118,13 +122,18 @@ 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.
 	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 +142,46 @@ 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)
+
+	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 +192,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()
 }
 

internal/routewatch/batchwrite_test.go

@@ -0,0 +1,192 @@
+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"
+)
+
+// handleMessageTimeout is how long a test waits for HandleMessage to return,
+// or for a write to start, before it fails.
+const handleMessageTimeout = 2 * time.Second
+
+// blockingStore is a mockStore whose batch writes for the AS, peer and prefix
+// handlers wait until the test closes release. writeStarted is closed when the
+// first of them starts. It records what each write was given.
+type blockingStore struct {
+	*mockStore
+
+	release      chan struct{}
+	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
+}
+
+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
+	}
+}
+
+// handleWhileAWriteIsBlocked hands h the messages in fill, which fill its
+// batch so that the handler writes it to store, waits until the write starts,
+// 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 close(store.release)
+
+	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")
+	}
+
+	if !returnsInTime(func() { h.HandleMessage(extra) }) {
+		t.Fatal("HandleMessage did not return while a database write was blocked")
+	}
+}
+
+func TestPrefixHandlerHandlesMessagesWhileAWriteIsBlocked(t *testing.T) {
+	const (
+		fillPrefix  = "192.0.2.0/24"
+		extraPrefix = "198.51.100.0/24"
+	)
+
+	announce := func(prefixes []string) *ristypes.RISMessage {
+		return &ristypes.RISMessage{
+			Peer:          testPeerIP,
+			Path:          ristypes.ASPath{testASNA},
+			Announcements: []ristypes.RISAnnouncement{{NextHop: testPeerIP, Prefixes: prefixes}},
+		}
+	}
+
+	store := newBlockingStore()
+	h := NewPrefixHandler(store, logger.New())
+	handleWhileAWriteIsBlocked(t, store, h,
+		[]*ristypes.RISMessage{announce(slices.Repeat([]string{fillPrefix}, prefixBatchSize))},
+		announce([]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)
+	}
+}

internal/routewatch/peerhandler.go

@@ -35,12 +35,15 @@ 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 never
+	// waits for a database write.
+	mu          sync.Mutex
+	peerBatch   []peerUpdate
+	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,11 +59,12 @@ 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{}),
 	}
 
 	// Start the flush timer goroutine
@@ -87,8 +91,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
@@ -109,13 +113,18 @@ 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.
 	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 +133,46 @@ 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)
+
+	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 +194,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()
 }
 

internal/routewatch/prefixhandler.go

@@ -40,12 +40,15 @@ 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 never
+	// waits for a database write.
+	mu          sync.Mutex
+	batch       []prefixUpdate
+	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,11 +63,12 @@ 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{}),
 	}
 
 	// Start the flush timer goroutine
@@ -140,13 +144,18 @@ 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.
 	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 +164,45 @@ 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)
+
+	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 +210,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
 	}
 
@@ -261,8 +285,6 @@ func (h *PrefixHandler) flushBatchLocked() {
 		"duration_ms", elapsed.Milliseconds(),
 	)
 
-	// Clear batch
-	h.batch = h.batch[:0]
 	h.lastFlush = time.Now()
 }
 

internal/routewatch/prefixhandler_test.go

@@ -30,11 +30,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{
@@ -46,7 +41,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",
@@ -66,7 +61,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) {
@@ -164,9 +159,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()

internal/server/handlers.go

@@ -201,16 +201,12 @@ func (s *Server) handleStatusJSON() http.HandlerFunc {
 		select {
 		case <-ctx.Done():
 			s.logger.Error("Database stats timeout in status.json")
-			writeJSONError(w, http.StatusGatewayTimeout, requestTimeoutMessage)
+			writeJSONError(w, http.StatusRequestTimeout, "Database timeout")
 
 			return
 		case err := <-errChan:
 			s.logger.Error("Failed to get database stats", "error", err)
-			if deadlinePassed(ctx) {
-				writeJSONError(w, http.StatusGatewayTimeout, requestTimeoutMessage)
-			} else {
-				writeJSONError(w, http.StatusInternalServerError, err.Error())
-			}
+			writeJSONError(w, http.StatusInternalServerError, err.Error())
 
 			return
 		case dbStats = <-statsChan:
@@ -390,7 +386,7 @@ func (s *Server) handleStats() http.HandlerFunc {
 		// Check if context is already cancelled
 		select {
 		case <-ctx.Done():
-			writeJSONError(w, http.StatusGatewayTimeout, requestTimeoutMessage)
+			http.Error(w, "Request timeout", http.StatusRequestTimeout)
 
 			return
 		default:
@@ -419,16 +415,11 @@ func (s *Server) handleStats() http.HandlerFunc {
 		select {
 		case <-ctx.Done():
 			s.logger.Error("Database stats timeout")
-			writeJSONError(w, http.StatusGatewayTimeout, requestTimeoutMessage)
-
+			// Don't write response here - timeout middleware already handles it
 			return
 		case err := <-errChan:
 			s.logger.Error("Failed to get database stats", "error", err)
-			if deadlinePassed(ctx) {
-				writeJSONError(w, http.StatusGatewayTimeout, requestTimeoutMessage)
-			} else {
-				writeJSONError(w, http.StatusInternalServerError, err.Error())
-			}
+			writeJSONError(w, http.StatusInternalServerError, err.Error())
 
 			return
 		case dbStats = <-statsChan:
@@ -631,13 +622,6 @@ func (s *Server) handleIPInfo() http.HandlerFunc {
 			}
 		}
 
-		// The time limit passed while resolving host names.
-		if deadlinePassed(ctx) {
-			writeJSONError(w, http.StatusGatewayTimeout, requestTimeoutMessage)
-
-			return
-		}
-
 		if len(ipsToLookup) == 0 {
 			writeJSONError(w, http.StatusBadRequest, "No valid IPs or hostnames provided")
 
@@ -680,13 +664,6 @@ func (s *Server) handleIPInfo() http.HandlerFunc {
 			}
 		}
 
-		// The time limit passed while looking up the IPs.
-		if deadlinePassed(ctx) {
-			writeJSONError(w, http.StatusGatewayTimeout, requestTimeoutMessage)
-
-			return
-		}
-
 		// Return response (even if no results, include errors)
 		if len(response.Results) == 0 && len(response.Errors) > 0 {
 			writeJSONError(w, http.StatusNotFound, "No routes found: "+response.Errors[0])
@@ -726,12 +703,9 @@ func (s *Server) handleASDetailJSON() http.HandlerFunc {
 
 		asInfo, prefixes, err := s.db.GetASDetailsContext(r.Context(), asn)
 		if err != nil {
-			switch {
-			case errors.Is(err, database.ErrNoRoute):
+			if errors.Is(err, database.ErrNoRoute) {
 				writeJSONError(w, http.StatusNotFound, err.Error())
-			case deadlinePassed(r.Context()):
-				writeJSONError(w, http.StatusGatewayTimeout, requestTimeoutMessage)
-			default:
+			} else {
 				writeJSONError(w, http.StatusInternalServerError, err.Error())
 			}
 
@@ -787,12 +761,9 @@ func (s *Server) handlePrefixDetailJSON() http.HandlerFunc {
 
 		routes, err := s.db.GetPrefixDetailsContext(r.Context(), prefix)
 		if err != nil {
-			switch {
-			case errors.Is(err, database.ErrNoRoute):
+			if errors.Is(err, database.ErrNoRoute) {
 				writeJSONError(w, http.StatusNotFound, err.Error())
-			case deadlinePassed(r.Context()):
-				writeJSONError(w, http.StatusGatewayTimeout, requestTimeoutMessage)
-			default:
+			} else {
 				writeJSONError(w, http.StatusInternalServerError, err.Error())
 			}
 
@@ -832,12 +803,9 @@ func (s *Server) handleASDetail() http.HandlerFunc {
 
 		asInfo, prefixes, err := s.db.GetASDetailsContext(r.Context(), asn)
 		if err != nil {
-			switch {
-			case errors.Is(err, database.ErrNoRoute):
+			if errors.Is(err, database.ErrNoRoute) {
 				http.Error(w, "AS not found", http.StatusNotFound)
-			case deadlinePassed(r.Context()):
-				http.Error(w, requestTimeoutMessage, http.StatusGatewayTimeout)
-			default:
+			} else {
 				s.logger.Error("Failed to get AS details", "error", err)
 				http.Error(w, "Internal server error", http.StatusInternalServerError)
 			}
@@ -964,12 +932,9 @@ func (s *Server) handlePrefixDetail() http.HandlerFunc {
 
 		routes, err := s.db.GetPrefixDetailsContext(r.Context(), prefix)
 		if err != nil {
-			switch {
-			case errors.Is(err, database.ErrNoRoute):
+			if errors.Is(err, database.ErrNoRoute) {
 				http.Error(w, "Prefix not found", http.StatusNotFound)
-			case deadlinePassed(r.Context()):
-				http.Error(w, requestTimeoutMessage, http.StatusGatewayTimeout)
-			default:
+			} else {
 				s.logger.Error("Failed to get prefix details", "error", err)
 				http.Error(w, "Internal server error", http.StatusInternalServerError)
 			}
@@ -1123,12 +1088,8 @@ func (s *Server) handlePrefixLength() http.HandlerFunc {
 		const maxPrefixes = 500
 		prefixes, err := s.db.GetRandomPrefixesByLengthContext(r.Context(), maskLength, ipVersion, maxPrefixes)
 		if err != nil {
-			if deadlinePassed(r.Context()) {
-				http.Error(w, requestTimeoutMessage, http.StatusGatewayTimeout)
-			} else {
-				s.logger.Error("Failed to get prefixes by length", "error", err)
-				http.Error(w, "Internal server error", http.StatusInternalServerError)
-			}
+			s.logger.Error("Failed to get prefixes by length", "error", err)
+			http.Error(w, "Internal server error", http.StatusInternalServerError)
 
 			return
 		}
@@ -1218,12 +1179,8 @@ func (s *Server) handlePrefixLength6() http.HandlerFunc {
 		const maxPrefixes = 500
 		prefixes, err := s.db.GetRandomPrefixesByLengthContext(r.Context(), maskLength, ipVersion, maxPrefixes)
 		if err != nil {
-			if deadlinePassed(r.Context()) {
-				http.Error(w, requestTimeoutMessage, http.StatusGatewayTimeout)
-			} else {
-				s.logger.Error("Failed to get prefixes by length", "error", err)
-				http.Error(w, "Internal server error", http.StatusInternalServerError)
-			}
+			s.logger.Error("Failed to get prefixes by length", "error", err)
+			http.Error(w, "Internal server error", http.StatusInternalServerError)
 
 			return
 		}

internal/server/handlers_test.go

@@ -200,179 +200,6 @@ func TestPrefixDetailHandlersParsePrefix(t *testing.T) {
 	}
 }
 
-// shortRequestDeadline is the deadline the tests below give a request. The
-// server's 30-second limit and the stats handlers' 4-second limit inherit it,
-// so the tests do not wait for either.
-const shortRequestDeadline = 50 * time.Millisecond
-
-// TestRequestPastTimeLimitAnswers504 sends a request through the server's
-// router to a page and to a JSON route whose handlers wait for the time limit
-// to pass and return without answering. The page answers 504; the JSON route
-// answers 504 with the JSON error body. Before the fix both answered 408,
-// written by the time limit middleware from its own goroutine while the
-// handler could still be writing.
-func TestRequestPastTimeLimitAnswers504(t *testing.T) {
-	s := New(nil, streamer.New(logger.New(), metrics.New()), logger.New(), &config.Config{})
-
-	waitForTimeLimit := func(_ http.ResponseWriter, r *http.Request) {
-		<-r.Context().Done()
-	}
-	// Registered the way setupRoutes registers pages and JSON routes.
-	s.router.Get("/slow-page", waitForTimeLimit)
-	s.router.Get("/slow.json", JSONValidationMiddleware(http.HandlerFunc(waitForTimeLimit)).ServeHTTP)
-
-	tests := []struct {
-		path         string
-		wantJSONBody bool
-	}{
-		{path: "/slow-page"},
-		{path: "/slow.json", wantJSONBody: true},
-	}
-
-	for _, tt := range tests {
-		ctx, cancel := context.WithTimeout(context.Background(), shortRequestDeadline)
-		rec := httptest.NewRecorder()
-		s.router.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, tt.path, nil).WithContext(ctx))
-		cancel()
-
-		if rec.Code != http.StatusGatewayTimeout {
-			t.Errorf("%s: status %d, want %d; body %s", tt.path, rec.Code, http.StatusGatewayTimeout, rec.Body)
-		}
-		if tt.wantJSONBody {
-			checkJSONErrorCode(t, tt.path, rec, http.StatusGatewayTimeout)
-		}
-	}
-}
-
-// TestStatsHandlersAnswer504WhenDatabaseReadRunsOut calls each stats handler
-// with a store that does not answer before the handler's 4-second limit runs
-// out, and checks that the handler itself answers 504 with the JSON error
-// body. Before the fix status.json answered 408, and /api/v1/stats wrote
-// nothing, which the JSON check on its route turned into 500.
-func TestStatsHandlersAnswer504WhenDatabaseReadRunsOut(t *testing.T) {
-	release := make(chan struct{})
-	defer close(release)
-	db := blockingStatsDB{release: release}
-	s := New(db, streamer.New(logger.New(), metrics.New()), logger.New(), &config.Config{})
-
-	handlers := map[string]http.HandlerFunc{
-		"status.json": s.handleStatusJSON(),
-		"stats":       s.handleStats(),
-	}
-
-	for name, handler := range handlers {
-		ctx, cancel := context.WithTimeout(context.Background(), shortRequestDeadline)
-		rec := httptest.NewRecorder()
-		handler(rec, httptest.NewRequest(http.MethodGet, "/", nil).WithContext(ctx))
-		cancel()
-
-		if rec.Code != http.StatusGatewayTimeout {
-			t.Errorf("%s: status %d, want %d; body %s", name, rec.Code, http.StatusGatewayTimeout, rec.Body)
-		}
-		checkJSONErrorCode(t, name, rec, http.StatusGatewayTimeout)
-	}
-}
-
-// slowDB embeds database.Store (left nil) and overrides the reads the routes
-// in TestRoutesAnswer504WhenDatabaseReadIsCutOff make. Each read waits until
-// its context ends and returns the context's error, as a slow database read
-// cut off by a time limit does.
-type slowDB struct {
-	database.Store
-}
-
-func (slowDB) GetStatsContext(ctx context.Context) (database.Stats, error) {
-	<-ctx.Done()
-
-	return database.Stats{}, ctx.Err()
-}
-
-func (slowDB) GetIPInfoContext(ctx context.Context, _ string) (*database.IPInfo, error) {
-	<-ctx.Done()
-
-	return nil, ctx.Err()
-}
-
-func (slowDB) GetASDetailsContext(ctx context.Context, _ int) (*database.ASN, []database.LiveRoute, error) {
-	<-ctx.Done()
-
-	return nil, nil, ctx.Err()
-}
-
-func (slowDB) GetPrefixDetailsContext(ctx context.Context, _ string) ([]database.LiveRoute, error) {
-	<-ctx.Done()
-
-	return nil, ctx.Err()
-}
-
-func (slowDB) GetRandomPrefixesByLengthContext(ctx context.Context, _, _, _ int) ([]database.LiveRoute, error) {
-	<-ctx.Done()
-
-	return nil, ctx.Err()
-}
-
-// TestRoutesAnswer504WhenDatabaseReadIsCutOff sends a request through the
-// server's router for each route that reads the database, with a store whose
-// reads fail when the request's time limit passes. Each route answers 504, the
-// JSON routes with the JSON error body. Before the fix they answered 408 or,
-// when the handler wrote first, its own error: 500, or 404 on the IP lookups.
-// The health check is left out: it answers 503 when its own 2-second ping runs
-// out.
-func TestRoutesAnswer504WhenDatabaseReadIsCutOff(t *testing.T) {
-	s := New(slowDB{}, streamer.New(logger.New(), metrics.New()), logger.New(), &config.Config{})
-
-	tests := []struct {
-		path         string
-		wantJSONBody bool
-	}{
-		{path: "/as/64500"},
-		{path: "/prefix/192.0.2.0/24"},
-		{path: "/prefixlength/24"},
-		{path: "/prefixlength6/48"},
-		{path: "/status.json", wantJSONBody: true},
-		{path: "/ip/192.0.2.1", wantJSONBody: true},
-		{path: "/api/v1/stats", wantJSONBody: true},
-		{path: "/api/v1/ip/192.0.2.1", wantJSONBody: true},
-		{path: "/api/v1/as/64500", wantJSONBody: true},
-		{path: "/api/v1/prefix/192.0.2.0/24", wantJSONBody: true},
-	}
-
-	for _, tt := range tests {
-		ctx, cancel := context.WithTimeout(context.Background(), shortRequestDeadline)
-		rec := httptest.NewRecorder()
-		s.router.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, tt.path, nil).WithContext(ctx))
-		cancel()
-
-		if rec.Code != http.StatusGatewayTimeout {
-			t.Errorf("%s: status %d, want %d; body %s", tt.path, rec.Code, http.StatusGatewayTimeout, rec.Body)
-		}
-		if tt.wantJSONBody {
-			checkJSONErrorCode(t, tt.path, rec, http.StatusGatewayTimeout)
-		}
-	}
-}
-
-// checkJSONErrorCode checks that rec holds the JSON error body with the given
-// code.
-func checkJSONErrorCode(t *testing.T, name string, rec *httptest.ResponseRecorder, want int) {
-	t.Helper()
-
-	var body struct {
-		Status string `json:"status"`
-		Error  struct {
-			Code int `json:"code"`
-		} `json:"error"`
-	}
-	if err := json.Unmarshal(rec.Body.Bytes(), &body); err != nil {
-		t.Errorf("%s: decoding the answer %q: %v", name, rec.Body, err)
-
-		return
-	}
-	if body.Status != "error" || body.Error.Code != want {
-		t.Errorf("%s: answer %s, want status \"error\" and error code %d", name, rec.Body, want)
-	}
-}
-
 // settledGoroutineCount lets transient goroutines finish, then reports the
 // current count.
 func settledGoroutineCount() int {

internal/server/middleware.go

@@ -4,7 +4,6 @@ import (
 	"bytes"
 	"context"
 	"encoding/json"
-	"errors"
 	"fmt"
 	"log/slog"
 	"net/http"
@@ -110,24 +109,142 @@ func JSONResponseMiddleware(next http.Handler) http.Handler {
 	})
 }
 
-// requestTimeoutMessage is the error message sent with 504 Gateway Timeout.
-const requestTimeoutMessage = "Request timeout"
+// timeoutWriter wraps ResponseWriter to prevent concurrent writes after timeout
+type timeoutWriter struct {
+	http.ResponseWriter
+	mu      sync.Mutex
+	written bool
+	header  http.Header // cached header to prevent concurrent access
+}
 
-// deadlinePassed reports whether ctx ended because its deadline passed. When
-// a handler's database read or other lookup fails, it uses this to tell a
-// request cut off by its time limit, which answers 504, from other failures.
-func deadlinePassed(ctx context.Context) bool {
-	return errors.Is(ctx.Err(), context.DeadlineExceeded)
+func (tw *timeoutWriter) Write(b []byte) (int, error) {
+	tw.mu.Lock()
+	defer tw.mu.Unlock()
+
+	if tw.written {
+		return 0, nil // Discard writes after timeout
+	}
+
+	return tw.ResponseWriter.Write(b)
+}
+
+func (tw *timeoutWriter) WriteHeader(statusCode int) {
+	tw.mu.Lock()
+	defer tw.mu.Unlock()
+
+	if tw.written {
+		return // Discard writes after timeout
+	}
+
+	tw.ResponseWriter.WriteHeader(statusCode)
+}
+
+func (tw *timeoutWriter) Header() http.Header {
+	tw.mu.Lock()
+	defer tw.mu.Unlock()
+
+	if tw.written {
+		// Return a copy to prevent modifications after timeout
+		if tw.header == nil {
+			tw.header = make(http.Header)
+		}
+
+		return tw.header
+	}
+
+	return tw.ResponseWriter.Header()
+}
+
+func (tw *timeoutWriter) markWritten() {
+	tw.mu.Lock()
+	defer tw.mu.Unlock()
+	tw.written = true
+}
+
+// TimeoutMiddleware creates an HTTP middleware that enforces a request timeout.
+// If the handler does not complete within the specified duration, the middleware
+// returns a JSON error response with HTTP status 408 (Request Timeout).
+//
+// The timeout parameter specifies the maximum duration allowed for request processing.
+// The returned middleware handles panics from the wrapped handler by re-panicking
+// after cleanup, and prevents concurrent writes to the response after timeout occurs.
+func TimeoutMiddleware(timeout time.Duration, logger *slog.Logger) func(http.Handler) http.Handler {
+	return func(next http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			startTime := time.Now()
+
+			ctx, cancel := context.WithTimeout(r.Context(), timeout)
+			defer cancel()
+
+			tw := &timeoutWriter{
+				ResponseWriter: w,
+				header:         make(http.Header),
+			}
+
+			done := make(chan struct{})
+			panicChan := make(chan interface{}, 1)
+
+			go func() {
+				defer func() {
+					if p := recover(); p != nil {
+						panicChan <- p
+					}
+				}()
+
+				next.ServeHTTP(tw, r.WithContext(ctx))
+				close(done)
+			}()
+
+			select {
+			case p := <-panicChan:
+				panic(p)
+			case <-done:
+				return
+			case <-ctx.Done():
+				tw.markWritten() // Prevent the handler from writing after timeout
+				execTime := time.Since(startTime)
+
+				// Log the timeout as a warning
+				logger.Warn("Request timeout",
+					"method", r.Method,
+					"path", r.URL.Path,
+					"duration_ms", execTime.Milliseconds(),
+					"remote_addr", r.RemoteAddr,
+				)
+
+				// Write directly to the underlying writer since we've marked tw as written
+				// This is safe because markWritten() prevents the handler from writing
+				tw.mu.Lock()
+				w.Header().Set("Content-Type", "application/json")
+				w.WriteHeader(http.StatusRequestTimeout)
+				tw.mu.Unlock()
+
+				response := map[string]interface{}{
+					"status": "error",
+					"error": map[string]interface{}{
+						"msg":  "Request timeout",
+						"code": http.StatusRequestTimeout,
+					},
+					"@meta": map[string]interface{}{
+						"time_zone":      "UTC",
+						"api_version":    1,
+						"execution_time": fmt.Sprintf("%d ms", execTime.Milliseconds()),
+					},
+				}
+
+				_ = json.NewEncoder(w).Encode(response)
+			}
+		})
+	}
 }
 
 // JSONValidationMiddleware is an HTTP middleware that validates JSON API responses.
 // It ensures that responses with Content-Type "application/json" contain valid JSON.
 //
 // If a response is not valid JSON or is empty when JSON is expected, the middleware
-// returns a properly formatted JSON error response. If the request's time limit
-// has passed, it answers 504 with the error message "Request timeout". For other
-// errors, it returns "Internal server error" with status 500 if the original
-// status was 200.
+// returns a properly formatted JSON error response. For timeout errors (status 408),
+// the error message will be "Request timeout". For other errors, it returns
+// "Internal server error" with status 500 if the original status was 200.
 //
 // Non-JSON responses are passed through unchanged.
 func JSONValidationMiddleware(next http.Handler) http.Handler {
@@ -176,14 +293,11 @@ func JSONValidationMiddleware(next http.Handler) http.Handler {
 		if statusCode == http.StatusOK {
 			statusCode = http.StatusInternalServerError
 		}
-		if deadlinePassed(r.Context()) {
-			statusCode = http.StatusGatewayTimeout
-		}
 		w.WriteHeader(statusCode)
 
 		errorMsg := "Internal server error"
-		if statusCode == http.StatusGatewayTimeout {
-			errorMsg = requestTimeoutMessage
+		if statusCode == http.StatusRequestTimeout {
+			errorMsg = "Request timeout"
 		}
 
 		response := map[string]interface{}{

internal/server/routes.go

@@ -16,12 +16,9 @@ func (s *Server) setupRoutes() {
 	r.Use(middleware.RealIP)
 	r.Use(RequestLoggerMiddleware(s.logger.Logger)) // Structured request logging
 	r.Use(middleware.Recoverer)
-	r.Use(JSONResponseMiddleware)
-	// Timeout must come after JSONResponseMiddleware, which sends its answer
-	// when the handler returns and would send a 200 ahead of Timeout's 504 for
-	// a handler that wrote nothing.
 	const requestTimeout = 30 * time.Second // Increased from 8s for slow queries
-	r.Use(middleware.Timeout(requestTimeout))
+	r.Use(TimeoutMiddleware(requestTimeout, s.logger.Logger))
+	r.Use(JSONResponseMiddleware)
 
 	// Routes
 	r.Get("/", s.handleIndex())