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

A request past the server's 30-second limit answered 408, which tells the
client it was too slow sending its request, and /api/v1/stats answered 500
when its 4-second read ran out. The limit is now chi's middleware.Timeout,
which answers 504 and never writes from a second goroutine while the
handler is writing.

Right after each database read or host name lookup, whatever came back, a
handler asks deadlinePassed whether its time limit has passed and, if so,
answers 504, with the JSON error body on JSON routes: a cut-off read can
return part of its rows without an error. The health check is unchanged.

Model: opus-5-5

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
+// 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.
-	tw.mu.Lock()
-	defer tw.mu.Unlock()
-
-	if tw.written {
-		return 0, nil // Discard writes after timeout
-	}
-
-	return tw.ResponseWriter.Write(b)
-}
-
-func (tw *timeoutWriter) WriteHeader(statusCode int) {
-	tw.mu.Lock()
-	defer tw.mu.Unlock()
-
-	if tw.written {
-		return // Discard writes after timeout
-	}
-
-	tw.ResponseWriter.WriteHeader(statusCode)
-}
-
-func (tw *timeoutWriter) Header() http.Header {
-	tw.mu.Lock()
-	defer tw.mu.Unlock()
-
-	if tw.written {
-		// Return a copy to prevent modifications after timeout
-		if tw.header == nil {
-			tw.header = make(http.Header)
-		}
-
-		return tw.header
-	}
-
-	return tw.ResponseWriter.Header()
-}
-
-func (tw *timeoutWriter) markWritten() {
-	tw.mu.Lock()
-	defer tw.mu.Unlock()
-	tw.written = true
-}
-
-// TimeoutMiddleware creates an HTTP middleware that enforces a request timeout.
-// If the handler does not complete within the specified duration, the middleware
-// returns a JSON error response with HTTP status 408 (Request Timeout).
 //
-// The timeout parameter specifies the maximum duration allowed for request processing.
+// A handler calls it right after each database read or host name lookup,
-// The returned middleware handles panics from the wrapped handler by re-panicking
+// before it looks at what came back, error or not, and answers 504 if it
-// after cleanup, and prevents concurrent writes to the response after timeout occurs.
+// reports true. A read cut off by the time limit can return an error, or part
-func TimeoutMiddleware(timeout time.Duration, logger *slog.Logger) func(http.Handler) http.Handler {
+// of its rows or none of them without one, so only this check tells a cut-off
-	return func(next http.Handler) http.Handler {
+// read from a complete one.
-		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+func deadlinePassed(ctx context.Context) bool {
-			startTime := time.Now()
+	return errors.Is(ctx.Err(), context.DeadlineExceeded)
-
-			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 +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 {
+		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())