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

internal/server/handlers.go

@@ -201,12 +201,16 @@ func (s *Server) handleStatusJSON() http.HandlerFunc {
 		select {
 		case <-ctx.Done():
 			s.logger.Error("Database stats timeout in status.json")
-			writeJSONError(w, http.StatusRequestTimeout, "Database timeout")
+			writeJSONError(w, http.StatusGatewayTimeout, requestTimeoutMessage)
 
 			return
 		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:
@@ -386,7 +390,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:
@@ -415,11 +419,16 @@ func (s *Server) handleStats() http.HandlerFunc {
 		select {
 		case <-ctx.Done():
 			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:
 			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:
@@ -622,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")
 
@@ -664,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])
@@ -702,6 +725,13 @@ 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)
+
+			return
+		}
 		if err != nil {
 			if errors.Is(err, database.ErrNoRoute) {
 				writeJSONError(w, http.StatusNotFound, err.Error())
@@ -760,6 +790,13 @@ 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)
+
+			return
+		}
 		if err != nil {
 			if errors.Is(err, database.ErrNoRoute) {
 				writeJSONError(w, http.StatusNotFound, err.Error())
@@ -802,6 +839,13 @@ 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)
+
+			return
+		}
 		if err != nil {
 			if errors.Is(err, database.ErrNoRoute) {
 				http.Error(w, "AS not found", http.StatusNotFound)
@@ -931,6 +975,13 @@ 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)
+
+			return
+		}
 		if err != nil {
 			if errors.Is(err, database.ErrNoRoute) {
 				http.Error(w, "Prefix not found", http.StatusNotFound)
@@ -1080,8 +1131,12 @@ func (s *Server) handlePrefixLength() http.HandlerFunc {
 		const maxPrefixes = 500
 		prefixes, err := s.db.GetRandomPrefixesByLengthContext(r.Context(), maskLength, ipVersion, maxPrefixes)
 		if err != nil {
-			s.logger.Error("Failed to get prefixes by length", "error", err)
+			if deadlinePassed(r.Context()) {
-			http.Error(w, "Internal server error", http.StatusInternalServerError)
+				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
 		}
@@ -1171,8 +1226,12 @@ func (s *Server) handlePrefixLength6() http.HandlerFunc {
 		const maxPrefixes = 500
 		prefixes, err := s.db.GetRandomPrefixesByLengthContext(r.Context(), maskLength, ipVersion, maxPrefixes)
 		if err != nil {
-			s.logger.Error("Failed to get prefixes by length", "error", err)
+			if deadlinePassed(r.Context()) {
-			http.Error(w, "Internal server error", http.StatusInternalServerError)
+				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

@@ -270,6 +270,183 @@ 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 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
+}
+
+func (slowDB) GetStatsContext(ctx context.Context) (database.Stats, error) {
+	<-ctx.Done()
+
+	return database.Stats{}, ctx.Err()
+}
+
+func (slowDB) GetIPInfoContext(ctx context.Context, _ string) (*database.IPInfo, error) {
+	<-ctx.Done()
+
+	return nil, ctx.Err()
+}
+
+func (slowDB) GetASDetailsContext(ctx context.Context, 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, 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, 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{})
+
+	tests := []struct {
+		path         string
+		wantJSONBody bool
+	}{
+		{path: "/as/64500"},
+		{path: "/prefix/192.0.2.0/24"},
+		{path: "/prefixlength/24"},
+		{path: "/prefixlength6/48"},
+		{path: "/status.json", wantJSONBody: true},
+		{path: "/ip/192.0.2.1", wantJSONBody: true},
+		{path: "/api/v1/stats", wantJSONBody: true},
+		{path: "/api/v1/ip/192.0.2.1", wantJSONBody: true},
+		{path: "/api/v1/as/64500", wantJSONBody: true},
+		{path: "/api/v1/prefix/192.0.2.0/24", wantJSONBody: true},
+	}
+
+	for _, tt := range tests {
+		ctx, cancel := context.WithTimeout(context.Background(), shortRequestDeadline)
+		rec := httptest.NewRecorder()
+		s.router.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, tt.path, nil).WithContext(ctx))
+		cancel()
+
+		if rec.Code != http.StatusGatewayTimeout {
+			t.Errorf("%s: status %d, want %d; body %s", tt.path, rec.Code, http.StatusGatewayTimeout, rec.Body)
+		}
+		if tt.wantJSONBody {
+			checkJSONErrorCode(t, tt.path, rec, http.StatusGatewayTimeout)
+		}
+	}
+}
+
+// checkJSONErrorCode checks that rec holds the JSON error body with the given
+// code.
+func checkJSONErrorCode(t *testing.T, name string, rec *httptest.ResponseRecorder, want int) {
+	t.Helper()
+
+	var body struct {
+		Status string `json:"status"`
+		Error  struct {
+			Code int `json:"code"`
+		} `json:"error"`
+	}
+	if err := json.Unmarshal(rec.Body.Bytes(), &body); err != nil {
+		t.Errorf("%s: decoding the answer %q: %v", name, rec.Body, err)
+
+		return
+	}
+	if body.Status != "error" || body.Error.Code != want {
+		t.Errorf("%s: answer %s, want status \"error\" and error code %d", name, rec.Body, want)
+	}
+}
+
 // settledGoroutineCount lets transient goroutines finish, then reports the
 // current count.
 func settledGoroutineCount() int {

internal/server/middleware.go

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