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