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

The hand-written time limit middleware answered 408 and wrote that
answer from its own goroutine while the handler could still be writing.
It is replaced by chi's middleware.Timeout, placed after
JSONResponseMiddleware so its 504 is not lost behind a 200.

After a database read or other lookup, a handler asks deadlinePassed
whether the request's time limit passed and, if so, answers 504 instead
of its 500, 404 or partial result: the AS and prefix reads, cut off
while reading rows, return part of them, or none, without an error.
JSONValidationMiddleware answers 504 with the JSON error body when the
limit has passed and the handler wrote nothing. status.json and
/api/v1/stats answer 504 when their 4-second database read runs out.

Model: opus-5-5

README.md

@@ -233,15 +233,6 @@ 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,12 +19,6 @@ 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
@@ -38,16 +32,12 @@ type ASHandler struct {
 	db     database.Store
 	logger *logger.Logger
 
-	// Batching. mu guards batch only: the flush goroutine takes the batch
-	// under mu and writes it after releasing mu, so HandleMessage waits for
-	// a database write only while batch is at asnBatchMaxSize.
-	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
+	// Batching
+	mu        sync.Mutex
+	batch     []asnOp
+	lastFlush time.Time
+	stopCh    chan struct{}
+	wg        sync.WaitGroup
 }
 
 type asnOp struct {
@@ -62,14 +52,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),
-		batchFullCh: make(chan struct{}, 1),
-		lastFlush:   time.Now(),
-		stopCh:      make(chan struct{}),
+		db:        db,
+		logger:    logger,
+		batch:     make([]asnOp, 0, asnBatchSize),
+		lastFlush: time.Now(),
+		stopCh:    make(chan struct{}),
 	}
-	h.batchTaken.L = &h.mu
 
 	// Start the flush timer goroutine
 	h.wg.Add(1)
@@ -97,8 +85,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. When the batch reaches the configured size
-// threshold, the flush goroutine is woken to write it.
+// with their associated timestamps. The batch is automatically flushed when
+// it reaches the configured size threshold.
 func (h *ASHandler) HandleMessage(msg *ristypes.RISMessage) {
 	// Use the pre-parsed timestamp
 	timestamp := msg.ParsedTimestamp
@@ -112,13 +100,6 @@ 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{
@@ -137,19 +118,13 @@ func (h *ASHandler) HandleMessage(msg *ristypes.RISMessage) {
 		}
 	}
 
-	// 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.
+	// Check if we need to flush
 	if len(h.batch) >= asnBatchSize {
-		select {
-		case h.batchFullCh <- struct{}{}:
-		default: // already woken
-		}
+		h.flushBatchLocked()
 	}
 }
 
-// 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.
+// flushLoop runs in a goroutine and periodically flushes batches
 func (h *ASHandler) flushLoop() {
 	defer h.wg.Done()
 	ticker := time.NewTicker(asnBatchTimeout)
@@ -158,47 +133,31 @@ func (h *ASHandler) flushLoop() {
 	for {
 		select {
 		case <-ticker.C:
+			h.mu.Lock()
 			if time.Since(h.lastFlush) >= asnBatchTimeout {
-				h.flushBatch()
+				h.flushBatchLocked()
 			}
-		case <-h.batchFullCh:
-			h.flushBatch()
+			h.mu.Unlock()
 		case <-h.stopCh:
 			// Final flush
-			h.flushBatch()
+			h.mu.Lock()
+			h.flushBatchLocked()
+			h.mu.Unlock()
 
 			return
 		}
 	}
 }
 
-// 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()
-
+// flushBatchLocked flushes the ASN batch to the database (must be called with mutex held)
+func (h *ASHandler) flushBatchLocked() {
 	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 batch {
+	for _, op := range h.batch {
 		if existing, ok := asnMap[op.number]; !ok || op.timestamp.After(existing) {
 			asnMap[op.number] = op.timestamp
 		}
@@ -209,6 +168,8 @@ func (h *ASHandler) flushBatch() {
 		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

@@ -1,374 +0,0 @@
-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,12 +23,6 @@ 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
 )
@@ -41,16 +35,12 @@ type PeerHandler struct {
 	db     database.Store
 	logger *logger.Logger
 
-	// Batching. mu guards peerBatch only: the flush goroutine takes the batch
-	// under mu and writes it after releasing mu, so HandleMessage waits for
-	// a database write only while peerBatch is at peerBatchMaxSize.
-	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
+	// Batching
+	mu        sync.Mutex
+	peerBatch []peerUpdate
+	lastFlush time.Time
+	stopCh    chan struct{}
+	wg        sync.WaitGroup
 }
 
 type peerUpdate struct {
@@ -66,14 +56,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),
-		batchFullCh: make(chan struct{}, 1),
-		lastFlush:   time.Now(),
-		stopCh:      make(chan struct{}),
+		db:        db,
+		logger:    logger,
+		peerBatch: make([]peerUpdate, 0, peerBatchSize),
+		lastFlush: time.Now(),
+		stopCh:    make(chan struct{}),
 	}
-	h.batchTaken.L = &h.mu
 
 	// Start the flush timer goroutine
 	h.wg.Add(1)
@@ -99,8 +87,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. The flush goroutine writes the batch to the
-// database when it reaches peerBatchSize or the batch timeout expires.
+// the update to an internal batch. When the batch reaches peerBatchSize
+// or the batch timeout expires, the batch is flushed to the database.
 func (h *PeerHandler) HandleMessage(msg *ristypes.RISMessage) {
 	// Parse peer ASN from string
 	peerASN := 0
@@ -113,13 +101,6 @@ 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,
@@ -128,19 +109,13 @@ func (h *PeerHandler) HandleMessage(msg *ristypes.RISMessage) {
 		timestamp:   msg.ParsedTimestamp,
 	})
 
-	// 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.
+	// Check if we need to flush
 	if len(h.peerBatch) >= peerBatchSize {
-		select {
-		case h.batchFullCh <- struct{}{}:
-		default: // already woken
-		}
+		h.flushBatchLocked()
 	}
 }
 
-// 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.
+// flushLoop runs in a goroutine and periodically flushes batches
 func (h *PeerHandler) flushLoop() {
 	defer h.wg.Done()
 	ticker := time.NewTicker(peerBatchTimeout)
@@ -149,47 +124,31 @@ func (h *PeerHandler) flushLoop() {
 	for {
 		select {
 		case <-ticker.C:
+			h.mu.Lock()
 			if time.Since(h.lastFlush) >= peerBatchTimeout {
-				h.flushBatch()
+				h.flushBatchLocked()
 			}
-		case <-h.batchFullCh:
-			h.flushBatch()
+			h.mu.Unlock()
 		case <-h.stopCh:
 			// Final flush
-			h.flushBatch()
+			h.mu.Lock()
+			h.flushBatchLocked()
+			h.mu.Unlock()
 
 			return
 		}
 	}
 }
 
-// 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()
-
+// flushBatchLocked flushes the peer batch to the database (must be called with mutex held)
+func (h *PeerHandler) flushBatchLocked() {
 	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 batch {
+	for _, update := range h.peerBatch {
 		if existing, ok := peerMap[update.peerIP]; !ok || update.timestamp.After(existing.timestamp) {
 			peerMap[update.peerIP] = update
 		}
@@ -211,6 +170,8 @@ func (h *PeerHandler) flushBatch() {
 		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,12 +24,6 @@ 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
@@ -46,16 +40,12 @@ type PrefixHandler struct {
 	logger  *logger.Logger
 	metrics *metrics.Tracker
 
-	// 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
+	// Batching
+	mu        sync.Mutex
+	batch     []prefixUpdate
+	lastFlush time.Time
+	stopCh    chan struct{}
+	wg        sync.WaitGroup
 }
 
 type prefixUpdate struct {
@@ -70,14 +60,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),
-		batchFullCh: make(chan struct{}, 1),
-		lastFlush:   time.Now(),
-		stopCh:      make(chan struct{}),
+		db:        db,
+		logger:    logger,
+		batch:     make([]prefixUpdate, 0, prefixBatchSize),
+		lastFlush: time.Now(),
+		stopCh:    make(chan struct{}),
 	}
-	h.batchTaken.L = &h.mu
 
 	// Start the flush timer goroutine
 	h.wg.Add(1)
@@ -117,13 +105,6 @@ 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).
@@ -159,19 +140,13 @@ func (h *PrefixHandler) HandleMessage(msg *ristypes.RISMessage) {
 		}
 	}
 
-	// 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.
+	// Check if we need to flush
 	if len(h.batch) >= prefixBatchSize {
-		select {
-		case h.batchFullCh <- struct{}{}:
-		default: // already woken
-		}
+		h.flushBatchLocked()
 	}
 }
 
-// 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.
+// flushLoop runs in a goroutine and periodically flushes batches
 func (h *PrefixHandler) flushLoop() {
 	defer h.wg.Done()
 	ticker := time.NewTicker(prefixBatchTimeout)
@@ -180,46 +155,30 @@ func (h *PrefixHandler) flushLoop() {
 	for {
 		select {
 		case <-ticker.C:
+			h.mu.Lock()
 			if time.Since(h.lastFlush) >= prefixBatchTimeout {
-				h.flushBatch()
+				h.flushBatchLocked()
 			}
-		case <-h.batchFullCh:
-			h.flushBatch()
+			h.mu.Unlock()
 		case <-h.stopCh:
 			// Final flush
-			h.flushBatch()
+			h.mu.Lock()
+			h.flushBatchLocked()
+			h.mu.Unlock()
 
 			return
 		}
 	}
 }
 
-// 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()
-
+// flushBatchLocked flushes the prefix batch to the database (must be called with mutex held)
+func (h *PrefixHandler) flushBatchLocked() {
 	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(batch)
+	batchSize := len(h.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
@@ -227,7 +186,7 @@ func (h *PrefixHandler) flushBatch() {
 	// updates within one second carry the same one.
 	type prefixAndPeer struct{ prefix, peer string }
 	lastUpdates := make(map[prefixAndPeer]prefixUpdate)
-	for _, update := range batch {
+	for _, update := range h.batch {
 		lastUpdates[prefixAndPeer{update.prefix, update.peer}] = update
 	}
 
@@ -301,6 +260,8 @@ func (h *PrefixHandler) flushBatch() {
 		"duration_ms", elapsed.Milliseconds(),
 	)
 
+	// Clear batch
+	h.batch = h.batch[:0]
 	h.lastFlush = time.Now()
 }
 

internal/routewatch/prefixhandler_test.go

@@ -31,6 +31,11 @@ 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{
@@ -42,7 +47,7 @@ func TestPrefixHandlerStoresPrefixesTheIPLookupFinds(t *testing.T) {
 			Prefixes: []string{"2001:db8:1::/48", "192.0.2.0/24"},
 		}},
 	})
-	h.flushBatch()
+	flush()
 
 	for ip, want := range map[string]string{
 		"2001:db8:1::1": "2001:db8:1::/48",
@@ -62,7 +67,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"},
 	})
-	h.flushBatch()
+	flush()
 
 	for _, ip := range []string{"2001:db8:1::1", "192.0.2.1"} {
 		if info, err := db.GetIPInfo(ip); !errors.Is(err, database.ErrNoRoute) {
@@ -160,7 +165,9 @@ func TestPrefixHandlerAppliesEachPeersLastUpdateInABatch(t *testing.T) {
 				for _, msg := range batch {
 					h.HandleMessage(msg)
 				}
-				h.flushBatch()
+				h.mu.Lock()
+				h.flushBatchLocked()
+				h.mu.Unlock()
 			}
 
 			ctx := context.Background()
@@ -258,7 +265,9 @@ func TestPrefixHandlerKeepsOneRoutePerPrefixAndPeer(t *testing.T) {
 				h := &PrefixHandler{db: db, logger: logger.New()}
 				for _, msg := range tt.messages {
 					h.HandleMessage(msg)
-					h.flushBatch()
+					h.mu.Lock()
+					h.flushBatchLocked()
+					h.mu.Unlock()
 				}
 
 				for _, p := range []string{prefix, otherPrefix} {

internal/server/handlers.go

@@ -198,24 +198,23 @@ 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.StatusGatewayTimeout, requestTimeoutMessage)
 
 			return
-		}
-		if err != nil {
+		case err := <-errChan:
 			s.logger.Error("Failed to get database stats", "error", err)
-			writeJSONError(w, http.StatusInternalServerError, err.Error())
+			if deadlinePassed(ctx) {
+				writeJSONError(w, http.StatusGatewayTimeout, requestTimeoutMessage)
+			} else {
+				writeJSONError(w, http.StatusInternalServerError, err.Error())
+			}
 
 			return
+		case dbStats = <-statsChan:
+			// Success
 		}
 
 		uptime := time.Since(metrics.ConnectedSince).Truncate(time.Second).String()
@@ -236,11 +235,6 @@ 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{
@@ -422,24 +416,23 @@ 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")
 			writeJSONError(w, http.StatusGatewayTimeout, requestTimeoutMessage)
 
 			return
-		}
-		if err != nil {
+		case err := <-errChan:
 			s.logger.Error("Failed to get database stats", "error", err)
-			writeJSONError(w, http.StatusInternalServerError, err.Error())
+			if deadlinePassed(ctx) {
+				writeJSONError(w, http.StatusGatewayTimeout, requestTimeoutMessage)
+			} else {
+				writeJSONError(w, http.StatusInternalServerError, err.Error())
+			}
 
 			return
+		case dbStats = <-statsChan:
+			// Success
 		}
 
 		uptime := time.Since(metrics.ConnectedSince).Truncate(time.Second).String()
@@ -478,11 +471,6 @@ 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
@@ -634,11 +622,6 @@ 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())
 
@@ -648,6 +631,13 @@ 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")
 
@@ -660,11 +650,6 @@ 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())
 
@@ -673,11 +658,6 @@ 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 {
@@ -700,6 +680,13 @@ 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])
@@ -738,6 +725,8 @@ func (s *Server) handleASDetailJSON() http.HandlerFunc {
 		}
 
 		asInfo, prefixes, err := s.db.GetASDetailsContext(r.Context(), asn)
+		// A read cut off by the time limit can return part of its rows, or
+		// none, without an error.
 		if deadlinePassed(r.Context()) {
 			writeJSONError(w, http.StatusGatewayTimeout, requestTimeoutMessage)
 
@@ -801,6 +790,8 @@ func (s *Server) handlePrefixDetailJSON() http.HandlerFunc {
 		prefix := parsed.Masked().String()
 
 		routes, _, err := s.db.GetPrefixDetailsContext(r.Context(), prefix)
+		// A read cut off by the time limit can return part of its rows, or
+		// none, without an error.
 		if deadlinePassed(r.Context()) {
 			writeJSONError(w, http.StatusGatewayTimeout, requestTimeoutMessage)
 
@@ -848,6 +839,8 @@ func (s *Server) handleASDetail() http.HandlerFunc {
 		}
 
 		asInfo, prefixes, err := s.db.GetASDetailsContext(r.Context(), asn)
+		// A read cut off by the time limit can return part of its rows, or
+		// none, without an error.
 		if deadlinePassed(r.Context()) {
 			http.Error(w, requestTimeoutMessage, http.StatusGatewayTimeout)
 
@@ -866,11 +859,6 @@ 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
@@ -987,6 +975,8 @@ func (s *Server) handlePrefixDetail() http.HandlerFunc {
 		prefix := parsed.Masked().String()
 
 		routes, originASNs, err := s.db.GetPrefixDetailsContext(r.Context(), prefix)
+		// A read cut off by the time limit can return part of its rows, or
+		// none, without an error.
 		if deadlinePassed(r.Context()) {
 			http.Error(w, requestTimeoutMessage, http.StatusGatewayTimeout)
 
@@ -1140,14 +1130,13 @@ 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)
+			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)
+			}
 
 			return
 		}
@@ -1236,14 +1225,13 @@ 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)
+			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)
+			}
 
 			return
 		}

internal/server/handlers_test.go

@@ -345,11 +345,11 @@ func TestStatsHandlersAnswer504WhenDatabaseReadRunsOut(t *testing.T) {
 
 // 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.
+// its context ends and answers as the real store can when the time limit cuts
+// its read off. The AS and prefix reads, cut off while reading their rows,
+// give no error: the AS read answers with the AS and the prefixes read so far,
+// here none, and the prefix read answers that there is no route. The other
+// reads answer with the context's error.
 type slowDB struct {
 	database.Store
 }
@@ -357,7 +357,7 @@ type slowDB struct {
 func (slowDB) GetStatsContext(ctx context.Context) (database.Stats, error) {
 	<-ctx.Done()
 
-	return database.Stats{}, nil
+	return database.Stats{}, ctx.Err()
 }
 
 func (slowDB) GetIPInfoContext(ctx context.Context, _ string) (*database.IPInfo, error) {
@@ -381,17 +381,17 @@ func (slowDB) GetPrefixDetailsContext(ctx context.Context, _ string) ([]database
 func (slowDB) GetRandomPrefixesByLengthContext(ctx context.Context, _, _, _ int) ([]database.LiveRoute, error) {
 	<-ctx.Done()
 
-	return nil, nil
+	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 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.
+// 504, the JSON routes with the JSON error body. Before the fix they answered
+// 408 or, when the handler wrote first, its own answer: 500, 404 on the IP
+// lookups and the prefix routes, or 200 on /api/v1/as/{asn}.
+// 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{})
 
@@ -412,91 +412,17 @@ func TestRoutesAnswer504WhenDatabaseReadIsCutOff(t *testing.T) {
 	}
 
 	for _, tt := range tests {
-		checkRouteAnswers504(t, s, tt.path, tt.wantJSONBody)
-	}
-}
+		ctx, cancel := context.WithTimeout(context.Background(), shortRequestDeadline)
+		rec := httptest.NewRecorder()
+		s.router.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, tt.path, nil).WithContext(ctx))
+		cancel()
 
-// 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)
+		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)
+		}
 	}
 }
 

internal/server/middleware.go

@@ -113,13 +113,10 @@ func JSONResponseMiddleware(next http.Handler) http.Handler {
 // requestTimeoutMessage is the error message sent with 504 Gateway Timeout.
 const requestTimeoutMessage = "Request timeout"
 
-// deadlinePassed reports whether ctx ended because its deadline passed.
-//
-// 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.
+// deadlinePassed reports whether ctx ended because its deadline passed. A
+// handler checks it after a database read or other lookup, so that a request
+// cut off by its time limit answers 504 instead of with the read's error or
+// its partial result.
 func deadlinePassed(ctx context.Context) bool {
 	return errors.Is(ctx.Err(), context.DeadlineExceeded)
 }