Answer 504, not 408, when a request outlasts its time limit (closes #70)

The 30-second limit is now chi's middleware.Timeout, which answers 504
and never writes from a second goroutine. It runs after
JSONResponseMiddleware, so a handler that wrote nothing is not sent a
200 first.

One rule covers the handlers: right after each database read or host
name lookup, before looking at what came back, error or not, a handler
asks deadlinePassed whether the time limit has passed and, if so,
answers 504, with the JSON error body on JSON routes. A cut-off read can
return an error, or part of its rows or none without one, so only this
check tells it from a complete read. JSONValidationMiddleware answers
504 with the JSON error body when the limit has passed and the handler
wrote nothing.

Model: opus-5-5

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.

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()
 }
 

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)
+			}
+		})
+	}
+}

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()
 }
 

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()
 }
 

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} {

internal/server/handlers.go

@@ -198,19 +198,24 @@ func (s *Server) handleStatusJSON() http.HandlerFunc {
 		}()
 
 		var dbStats database.Stats
+		var err error
 		select {
 		case <-ctx.Done():
+			err = ctx.Err()
+		case err = <-errChan:
+		case dbStats = <-statsChan:
+		}
+		if deadlinePassed(ctx) {
 			s.logger.Error("Database stats timeout in status.json")
-			writeJSONError(w, http.StatusRequestTimeout, "Database timeout")
+			writeJSONError(w, http.StatusGatewayTimeout, requestTimeoutMessage)
 
 			return
-		case err := <-errChan:
+		}
+		if err != nil {
 			s.logger.Error("Failed to get database stats", "error", err)
 			writeJSONError(w, http.StatusInternalServerError, err.Error())
 
 			return
-		case dbStats = <-statsChan:
-			// Success
 		}
 
 		uptime := time.Since(metrics.ConnectedSince).Truncate(time.Second).String()
@@ -231,6 +236,11 @@ func (s *Server) handleStatusJSON() http.HandlerFunc {
 		var whoisStats *WHOISStatsInfo
 		if s.asnFetcher != nil {
 			whoisStats = s.getWHOISStats(ctx)
+			if deadlinePassed(ctx) {
+				writeJSONError(w, http.StatusGatewayTimeout, requestTimeoutMessage)
+
+				return
+			}
 		}
 
 		stats := Stats{
@@ -386,7 +396,7 @@ func (s *Server) handleStats() http.HandlerFunc {
 		// Check if context is already cancelled
 		select {
 		case <-ctx.Done():
-			http.Error(w, "Request timeout", http.StatusRequestTimeout)
+			writeJSONError(w, http.StatusGatewayTimeout, requestTimeoutMessage)
 
 			return
 		default:
@@ -412,18 +422,24 @@ func (s *Server) handleStats() http.HandlerFunc {
 		}()
 
 		var dbStats database.Stats
+		var err error
 		select {
 		case <-ctx.Done():
+			err = ctx.Err()
+		case err = <-errChan:
+		case dbStats = <-statsChan:
+		}
+		if deadlinePassed(ctx) {
 			s.logger.Error("Database stats timeout")
-			// Don't write response here - timeout middleware already handles it
+			writeJSONError(w, http.StatusGatewayTimeout, requestTimeoutMessage)
+
 			return
-		case err := <-errChan:
+		}
+		if err != nil {
 			s.logger.Error("Failed to get database stats", "error", err)
 			writeJSONError(w, http.StatusInternalServerError, err.Error())
 
 			return
-		case dbStats = <-statsChan:
-			// Success
 		}
 
 		uptime := time.Since(metrics.ConnectedSince).Truncate(time.Second).String()
@@ -462,6 +478,11 @@ func (s *Server) handleStats() http.HandlerFunc {
 		var whoisStats *WHOISStatsInfo
 		if s.asnFetcher != nil {
 			whoisStats = s.getWHOISStats(ctx)
+			if deadlinePassed(ctx) {
+				writeJSONError(w, http.StatusGatewayTimeout, requestTimeoutMessage)
+
+				return
+			}
 		}
 
 		// Calculate connection duration
@@ -613,6 +634,11 @@ func (s *Server) handleIPInfo() http.HandlerFunc {
 			} else {
 				// It's a hostname - resolve it
 				resolved, err := net.DefaultResolver.LookupHost(ctx, t)
+				if deadlinePassed(ctx) {
+					writeJSONError(w, http.StatusGatewayTimeout, requestTimeoutMessage)
+
+					return
+				}
 				if err != nil {
 					response.Errors = append(response.Errors, t+": "+err.Error())
 
@@ -634,6 +660,11 @@ func (s *Server) handleIPInfo() http.HandlerFunc {
 		// Look up each IP
 		for _, ip := range ipsToLookup {
 			ipInfo, err := s.db.GetIPInfoContext(ctx, ip)
+			if deadlinePassed(ctx) {
+				writeJSONError(w, http.StatusGatewayTimeout, requestTimeoutMessage)
+
+				return
+			}
 			if err != nil {
 				response.Errors = append(response.Errors, ip+": "+err.Error())
 
@@ -642,6 +673,11 @@ func (s *Server) handleIPInfo() http.HandlerFunc {
 
 			// Do PTR lookup for this IP
 			ptrs, err := net.DefaultResolver.LookupAddr(ctx, ip)
+			if deadlinePassed(ctx) {
+				writeJSONError(w, http.StatusGatewayTimeout, requestTimeoutMessage)
+
+				return
+			}
 			if err == nil && len(ptrs) > 0 {
 				// Remove trailing dots from PTR records
 				for i, ptr := range ptrs {
@@ -702,6 +738,11 @@ func (s *Server) handleASDetailJSON() http.HandlerFunc {
 		}
 
 		asInfo, prefixes, err := s.db.GetASDetailsContext(r.Context(), asn)
+		if deadlinePassed(r.Context()) {
+			writeJSONError(w, http.StatusGatewayTimeout, requestTimeoutMessage)
+
+			return
+		}
 		if err != nil {
 			if errors.Is(err, database.ErrNoRoute) {
 				writeJSONError(w, http.StatusNotFound, err.Error())
@@ -760,6 +801,11 @@ func (s *Server) handlePrefixDetailJSON() http.HandlerFunc {
 		prefix := parsed.Masked().String()
 
 		routes, _, err := s.db.GetPrefixDetailsContext(r.Context(), prefix)
+		if deadlinePassed(r.Context()) {
+			writeJSONError(w, http.StatusGatewayTimeout, requestTimeoutMessage)
+
+			return
+		}
 		if err != nil {
 			if errors.Is(err, database.ErrNoRoute) {
 				writeJSONError(w, http.StatusNotFound, err.Error())
@@ -802,6 +848,11 @@ func (s *Server) handleASDetail() http.HandlerFunc {
 		}
 
 		asInfo, prefixes, err := s.db.GetASDetailsContext(r.Context(), asn)
+		if deadlinePassed(r.Context()) {
+			http.Error(w, requestTimeoutMessage, http.StatusGatewayTimeout)
+
+			return
+		}
 		if err != nil {
 			if errors.Is(err, database.ErrNoRoute) {
 				http.Error(w, "AS not found", http.StatusNotFound)
@@ -815,6 +866,11 @@ func (s *Server) handleASDetail() http.HandlerFunc {
 
 		// Get peers
 		peers, err := s.db.GetASPeersContext(r.Context(), asn)
+		if deadlinePassed(r.Context()) {
+			http.Error(w, requestTimeoutMessage, http.StatusGatewayTimeout)
+
+			return
+		}
 		if err != nil {
 			s.logger.Error("Failed to get AS peers", "error", err)
 			// Continue without peers rather than failing the whole request
@@ -931,6 +987,11 @@ func (s *Server) handlePrefixDetail() http.HandlerFunc {
 		prefix := parsed.Masked().String()
 
 		routes, originASNs, err := s.db.GetPrefixDetailsContext(r.Context(), prefix)
+		if deadlinePassed(r.Context()) {
+			http.Error(w, requestTimeoutMessage, http.StatusGatewayTimeout)
+
+			return
+		}
 		if err != nil {
 			if errors.Is(err, database.ErrNoRoute) {
 				http.Error(w, "Prefix not found", http.StatusNotFound)
@@ -1079,6 +1140,11 @@ func (s *Server) handlePrefixLength() http.HandlerFunc {
 		// Get random sample of prefixes
 		const maxPrefixes = 500
 		prefixes, err := s.db.GetRandomPrefixesByLengthContext(r.Context(), maskLength, ipVersion, maxPrefixes)
+		if deadlinePassed(r.Context()) {
+			http.Error(w, requestTimeoutMessage, http.StatusGatewayTimeout)
+
+			return
+		}
 		if err != nil {
 			s.logger.Error("Failed to get prefixes by length", "error", err)
 			http.Error(w, "Internal server error", http.StatusInternalServerError)
@@ -1170,6 +1236,11 @@ func (s *Server) handlePrefixLength6() http.HandlerFunc {
 		// Get random sample of prefixes
 		const maxPrefixes = 500
 		prefixes, err := s.db.GetRandomPrefixesByLengthContext(r.Context(), maskLength, ipVersion, maxPrefixes)
+		if deadlinePassed(r.Context()) {
+			http.Error(w, requestTimeoutMessage, http.StatusGatewayTimeout)
+
+			return
+		}
 		if err != nil {
 			s.logger.Error("Failed to get prefixes by length", "error", err)
 			http.Error(w, "Internal server error", http.StatusInternalServerError)

internal/server/handlers_test.go

@@ -270,6 +270,257 @@ func TestPrefixDetailShowsOriginNames(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 then answers as the real store does when the time
+// limit cuts it off: the IP read with the context's error, the prefix read
+// that there is no route, and the others with what they had read, here
+// nothing, and no error. Every other Store method is unused and would panic
+// if called.
+type slowDB struct {
+	database.Store
+}
+
+func (slowDB) GetStatsContext(ctx context.Context) (database.Stats, error) {
+	<-ctx.Done()
+
+	return database.Stats{}, nil
+}
+
+func (slowDB) GetIPInfoContext(ctx context.Context, _ string) (*database.IPInfo, error) {
+	<-ctx.Done()
+
+	return nil, ctx.Err()
+}
+
+func (slowDB) GetASDetailsContext(ctx context.Context, asn int) (*database.ASN, []database.LiveRoute, error) {
+	<-ctx.Done()
+
+	return &database.ASN{ASN: asn}, nil, nil
+}
+
+func (slowDB) GetPrefixDetailsContext(ctx context.Context, _ string) ([]database.LiveRoute, map[int]database.ASN, error) {
+	<-ctx.Done()
+
+	return nil, nil, database.ErrNoRoute
+}
+
+func (slowDB) GetRandomPrefixesByLengthContext(ctx context.Context, _, _, _ int) ([]database.LiveRoute, error) {
+	<-ctx.Done()
+
+	return nil, nil
+}
+
+// TestRoutesAnswer504WhenDatabaseReadIsCutOff sends a request through the
+// server's router for each route that reads the database, with a store whose
+// reads are cut off when the request's time limit passes. Each route answers
+// 504 with the time limit message. Without its check after the read, a route
+// answers with what the cut-off read returned: 500, 404, 200, or, on a page
+// that stops before rendering once the request has ended, 504 with an empty
+// body. 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 {
+		checkRouteAnswers504(t, s, tt.path, tt.wantJSONBody)
+	}
+}
+
+// slowSecondReadDB embeds database.Store (left nil) and overrides the reads
+// of the routes that read the database twice. The first read answers at once;
+// the second, the AS page's peers read or the stats routes' WHOIS read, waits
+// until its context ends and then answers with the context's error, as the
+// real store does.
+type slowSecondReadDB struct {
+	database.Store
+}
+
+func (slowSecondReadDB) GetASDetailsContext(_ context.Context, asn int) (*database.ASN, []database.LiveRoute, error) {
+	return &database.ASN{ASN: asn}, nil, nil
+}
+
+func (slowSecondReadDB) GetASPeersContext(ctx context.Context, _ int) ([]database.ASPeer, error) {
+	<-ctx.Done()
+
+	return nil, ctx.Err()
+}
+
+func (slowSecondReadDB) GetStatsContext(_ context.Context) (database.Stats, error) {
+	return database.Stats{}, nil
+}
+
+func (slowSecondReadDB) GetWHOISStats(ctx context.Context, _ time.Duration) (*database.WHOISStats, error) {
+	<-ctx.Done()
+
+	return nil, ctx.Err()
+}
+
+// idleASNFetcher is a WHOIS fetcher that does nothing. The stats handlers
+// read the WHOIS stats only when the server has a fetcher.
+type idleASNFetcher struct{}
+
+func (idleASNFetcher) QueueImmediate(int) {}
+
+func (idleASNFetcher) GetStats() ASNFetcherStats {
+	return ASNFetcherStats{}
+}
+
+// TestRoutesAnswer504WhenSecondDatabaseReadIsCutOff sends a request through
+// the server's router for each route that reads the database twice, with a
+// store that answers the first read and cuts off the second when the
+// request's time limit passes. Each route answers 504 with the time limit
+// message. Without the check after the second read, the AS page answers 504
+// with an empty body and the stats routes answer 200 without the WHOIS stats.
+func TestRoutesAnswer504WhenSecondDatabaseReadIsCutOff(t *testing.T) {
+	s := New(slowSecondReadDB{}, streamer.New(logger.New(), metrics.New()), logger.New(), &config.Config{})
+	s.SetASNFetcher(idleASNFetcher{})
+
+	tests := []struct {
+		path         string
+		wantJSONBody bool
+	}{
+		{path: "/as/64500"},
+		{path: "/status.json", wantJSONBody: true},
+		{path: "/api/v1/stats", wantJSONBody: true},
+	}
+
+	for _, tt := range tests {
+		checkRouteAnswers504(t, s, tt.path, tt.wantJSONBody)
+	}
+}
+
+// checkRouteAnswers504 sends a request for path through the server's router
+// with a short deadline and checks that it answers 504 with the time limit
+// message: the JSON error body on JSON routes, the plain message on pages.
+func checkRouteAnswers504(t *testing.T, s *Server, path string, wantJSONBody bool) {
+	t.Helper()
+
+	ctx, cancel := context.WithTimeout(context.Background(), shortRequestDeadline)
+	defer cancel()
+	rec := httptest.NewRecorder()
+	s.router.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, path, nil).WithContext(ctx))
+
+	if rec.Code != http.StatusGatewayTimeout {
+		t.Errorf("%s: status %d, want %d; body %s", path, rec.Code, http.StatusGatewayTimeout, rec.Body)
+	}
+	if wantJSONBody {
+		checkJSONErrorCode(t, path, rec, http.StatusGatewayTimeout)
+	} else if body := strings.TrimSpace(rec.Body.String()); body != requestTimeoutMessage {
+		t.Errorf("%s: body %q, want %q", path, body, requestTimeoutMessage)
+	}
+}
+
+// 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,6 +4,7 @@ import (
 	"bytes"
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"log/slog"
 	"net/http"
@@ -109,142 +110,28 @@ func JSONResponseMiddleware(next http.Handler) http.Handler {
 	})
 }
 
-// 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
-}
+// requestTimeoutMessage is the error message sent with 504 Gateway Timeout.
+const requestTimeoutMessage = "Request timeout"
 
-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).
+// deadlinePassed reports whether ctx ended because its deadline passed.
 //
-// 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)
-			}
-		})
-	}
+// A handler calls it right after each database read or host name lookup,
+// before it looks at what came back, error or not, and answers 504 if it
+// reports true. A read cut off by the time limit can return an error, or part
+// of its rows or none of them without one, so only this check tells a cut-off
+// read from a complete one.
+func deadlinePassed(ctx context.Context) bool {
+	return errors.Is(ctx.Err(), context.DeadlineExceeded)
 }
 
 // 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. 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.
+// 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.
 //
 // Non-JSON responses are passed through unchanged.
 func JSONValidationMiddleware(next http.Handler) http.Handler {
@@ -293,11 +180,14 @@ 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.StatusRequestTimeout {
-			errorMsg = "Request timeout"
+		if statusCode == http.StatusGatewayTimeout {
+			errorMsg = requestTimeoutMessage
 		}
 
 		response := map[string]interface{}{

internal/server/routes.go

@@ -16,9 +16,12 @@ func (s *Server) setupRoutes() {
 	r.Use(middleware.RealIP)
 	r.Use(RequestLoggerMiddleware(s.logger.Logger)) // Structured request logging
 	r.Use(middleware.Recoverer)
-	const requestTimeout = 30 * time.Second // Increased from 8s for slow queries
-	r.Use(TimeoutMiddleware(requestTimeout, s.logger.Logger))
 	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))
 
 	// Routes
 	r.Get("/", s.handleIndex())