cmd: cancel in-flight scrapes and bound the HTTP server shutdown (#2728)

Signed-off-by: Jan-Otto Kröpke <mail@jkroepke.de>

cmd/windows_exporter/main.go

@@ -27,6 +27,7 @@ import (
 	"log/slog"
 	"maps"
 	"math"
+	"net"
 	"net/http"
 	"net/http/httptest"
 	"net/http/pprof"
@@ -53,8 +54,16 @@ import (
 	"golang.org/x/sys/windows"
 )
 
-// collectionCloseTimeout bounds how long the exporter waits for the collectors to close on shutdown.
-const collectionCloseTimeout = 10 * time.Second
+const (
+	// collectionCloseTimeout bounds how long the exporter waits for the collectors to close on shutdown.
+	collectionCloseTimeout = 10 * time.Second
+	// serverShutdownTimeout bounds how long the exporter waits for in-flight requests on shutdown,
+	// before it closes their connections.
+	serverShutdownTimeout = 5 * time.Second
+	// writeTimeoutMargin is the time a scrape gets on top of the largest scrape budget
+	// to encode and write its response.
+	writeTimeoutMargin = 30 * time.Second
+)
 
 func main() {
 	// os.Kill can't be caught. On Windows, Go delivers console close, logoff and shutdown events as SIGTERM.
@@ -277,13 +286,11 @@ func run(ctx context.Context, args []string) int {
 
 	logger.LogAttrs(ctx, slog.LevelInfo, fmt.Sprintf("started windows_exporter in %s", time.Since(startTime)))
 
-	server := &http.Server{
-		ReadHeaderTimeout: 5 * time.Second,
-		IdleTimeout:       60 * time.Second,
-		ReadTimeout:       5 * time.Second,
-		WriteTimeout:      5 * time.Minute,
-		Handler:           mux,
-	}
+	// Requests derive their context from serveCtx, so that canceling it cancels in-flight scrapes.
+	serveCtx, cancelServe := context.WithCancel(ctx)
+	defer cancelServe()
+
+	server := newServer(serveCtx, mux)
 
 	errCh := make(chan error, 1)
 
@@ -312,23 +319,61 @@ func run(ctx context.Context, args []string) int {
 		}
 	}
 
-	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
-	defer cancel()
+	// In-flight scrapes must be canceled before the server drains, so that they end with
+	// 503 Service Unavailable instead of holding the shutdown until their scrape budget ends.
+	// On a signal or service stop, serveCtx is canceled already as a child of ctx; this call
+	// also covers a server that stopped for another reason.
+	cancelServe()
 
-	//nolint:contextcheck // create a new context for server shutdown
-	if err = server.Shutdown(ctx); err != nil {
-		//nolint:contextcheck
+	if err = shutdownServer(ctx, server, serverShutdownTimeout); err != nil {
 		logger.LogAttrs(ctx, slog.LevelError, "Failed to shutdown windows_exporter",
 			slog.Any("err", err),
 		)
 	} else {
-		//nolint:contextcheck
 		logger.LogAttrs(ctx, slog.LevelInfo, "windows_exporter has shut down")
 	}
 
+	// The deferred closeCollection closes the collectors after the server has shut down.
 	return 0
 }
 
+// newServer returns the HTTP server of the exporter. Requests derive their context from ctx.
+func newServer(ctx context.Context, handler http.Handler) *http.Server {
+	return &http.Server{
+		BaseContext:       func(net.Listener) context.Context { return ctx },
+		ReadHeaderTimeout: 5 * time.Second,
+		IdleTimeout:       60 * time.Second,
+		ReadTimeout:       5 * time.Second,
+		// A scrape may use the whole scrape budget before it writes the response.
+		WriteTimeout: httphandler.MaxScrapeTimeout + writeTimeoutMargin,
+		Handler:      handler,
+	}
+}
+
+// shutdownServer stops the server. It waits up to timeout for in-flight requests to complete.
+// Afterwards, it closes the remaining connections, so that a handler that doesn't return
+// can't block the shutdown.
+func shutdownServer(ctx context.Context, server *http.Server, timeout time.Duration) error {
+	// ctx is usually canceled already, as it is what triggered the shutdown.
+	shutdownCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), timeout)
+	defer cancel()
+
+	err := server.Shutdown(shutdownCtx)
+	if err == nil {
+		return nil
+	}
+
+	if closeErr := server.Close(); closeErr != nil {
+		err = errors.Join(err, fmt.Errorf("failed to close server: %w", closeErr))
+	}
+
+	if errors.Is(err, context.DeadlineExceeded) {
+		return fmt.Errorf("in-flight requests didn't complete within %s, closed their connections: %w", timeout, err)
+	}
+
+	return fmt.Errorf("failed to shut down server, closed it: %w", err)
+}
+
 func logCurrentUser(ctx context.Context, logger *slog.Logger) {
 	u, err := user.Current()
 	if err != nil {

cmd/windows_exporter/main_test.go

@@ -212,6 +212,154 @@ func TestScrapeStdout(t *testing.T) {
 	require.Contains(t, output, `windows_exporter_collector_success{collector="cpu"} 1`)
 }
 
+// serveTest serves handler with the exporter's server settings on a free loopback port.
+func serveTest(t *testing.T, ctx context.Context, handler http.Handler) (*http.Server, string) {
+	t.Helper()
+
+	listener, err := (&net.ListenConfig{}).Listen(ctx, "tcp", "127.0.0.1:0")
+	require.NoError(t, err)
+
+	server := newServer(ctx, handler)
+
+	go func() {
+		_ = server.Serve(listener)
+	}()
+
+	t.Cleanup(func() {
+		_ = server.Close()
+	})
+
+	return server, "http://" + listener.Addr().String() + "/"
+}
+
+func TestNewServerWriteTimeout(t *testing.T) {
+	t.Parallel()
+
+	server := newServer(t.Context(), http.NotFoundHandler())
+
+	// A scrape that uses its full budget must still be able to write its response.
+	require.Greater(t, server.WriteTimeout, httphandler.MaxScrapeTimeout)
+}
+
+func TestNewServerCancelsInFlightRequests(t *testing.T) {
+	t.Parallel()
+
+	ctx, cancel := context.WithCancel(t.Context())
+	defer cancel()
+
+	started := make(chan struct{})
+
+	_, uri := serveTest(t, ctx, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		close(started)
+
+		select {
+		case <-r.Context().Done():
+			_, _ = io.WriteString(w, "canceled")
+		case <-time.After(30 * time.Second):
+			_, _ = io.WriteString(w, "not canceled")
+		}
+	}))
+
+	type result struct {
+		body string
+		err  error
+	}
+
+	resultCh := make(chan result, 1)
+
+	go func() {
+		req, err := http.NewRequestWithContext(t.Context(), http.MethodGet, uri, nil)
+		if err != nil {
+			resultCh <- result{err: err}
+
+			return
+		}
+
+		resp, err := http.DefaultClient.Do(req)
+		if err != nil {
+			resultCh <- result{err: err}
+
+			return
+		}
+
+		defer resp.Body.Close()
+
+		body, err := io.ReadAll(resp.Body)
+		resultCh <- result{body: string(body), err: err}
+	}()
+
+	<-started
+	cancel()
+
+	select {
+	case res := <-resultCh:
+		require.NoError(t, res.err)
+		require.Equal(t, "canceled", res.body)
+	case <-time.After(10 * time.Second):
+		t.Fatal("in-flight request wasn't canceled with the server context")
+	}
+}
+
+func TestShutdownServerClosesStuckRequests(t *testing.T) {
+	t.Parallel()
+
+	started := make(chan struct{})
+	release := make(chan struct{})
+
+	t.Cleanup(func() {
+		close(release)
+	})
+
+	// The handler ignores its context, like a scrape stuck in a native call.
+	server, uri := serveTest(t, t.Context(), http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		close(started)
+		<-release
+
+		_, _ = io.WriteString(w, "released")
+	}))
+
+	errCh := make(chan error, 1)
+
+	go func() {
+		req, err := http.NewRequestWithContext(t.Context(), http.MethodGet, uri, nil)
+		if err != nil {
+			errCh <- err
+
+			return
+		}
+
+		resp, err := http.DefaultClient.Do(req)
+		if err == nil {
+			_ = resp.Body.Close()
+		}
+
+		errCh <- err
+	}()
+
+	<-started
+
+	start := time.Now()
+	err := shutdownServer(t.Context(), server, 100*time.Millisecond)
+	require.ErrorIs(t, err, context.DeadlineExceeded)
+	require.Less(t, time.Since(start), 5*time.Second)
+
+	// Close drops the connection of the stuck request, so the client doesn't wait for it.
+	select {
+	case err := <-errCh:
+		require.Error(t, err)
+	case <-time.After(10 * time.Second):
+		t.Fatal("connection of the stuck request wasn't closed")
+	}
+}
+
+func TestShutdownServerIdle(t *testing.T) {
+	t.Parallel()
+
+	server, _ := serveTest(t, t.Context(), http.NotFoundHandler())
+
+	require.NoError(t, shutdownServer(t.Context(), server, time.Second))
+}
+
 func TestSetPriorityWindows(t *testing.T) {
 	t.Parallel()
 

internal/httphandler/httphandler.go

@@ -19,6 +19,7 @@ package httphandler
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"log/slog"
 	"math"
@@ -32,6 +33,7 @@ import (
 	"github.com/prometheus/client_golang/prometheus/collectors"
 	"github.com/prometheus/client_golang/prometheus/collectors/version"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
+	dto "github.com/prometheus/client_model/go"
 )
 
 // Interface guard.
@@ -42,9 +44,10 @@ const (
 	defaultScrapeTimeout = 10 * time.Second
 	// minScrapeTimeout is the smallest time given to the collectors.
 	minScrapeTimeout = 10 * time.Millisecond
-	// maxScrapeTimeout is the largest time given to the collectors.
-	// The exporter's HTTP server ends responses after 5 minutes anyway.
-	maxScrapeTimeout = 5 * time.Minute
+	// MaxScrapeTimeout is the largest time given to the collectors, including the time a scrape
+	// waits for an earlier one. The HTTP server's write timeout must exceed it, so that a scrape
+	// that uses its full budget can still write its response.
+	MaxScrapeTimeout = 5 * time.Minute
 )
 
 type MetricsHTTPHandler struct {
@@ -120,7 +123,7 @@ func (c *MetricsHTTPHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 // getScrapeTimeout returns the time the collectors get for this request.
 // It is the client's X-Prometheus-Scrape-Timeout-Seconds minus the timeout margin,
 // but the margin takes at most half of the client's timeout.
-// The result is between minScrapeTimeout and maxScrapeTimeout.
+// The result is between minScrapeTimeout and MaxScrapeTimeout.
 func (c *MetricsHTTPHandler) getScrapeTimeout(logger *slog.Logger, r *http.Request) time.Duration {
 	timeout := defaultScrapeTimeout
 
@@ -130,7 +133,7 @@ func (c *MetricsHTTPHandler) getScrapeTimeout(logger *slog.Logger, r *http.Reque
 		switch {
 		case math.IsInf(seconds, 1):
 			// Includes values too large for a float64, which ParseFloat reports as an error.
-			timeout = maxScrapeTimeout
+			timeout = MaxScrapeTimeout
 		case err != nil, math.IsNaN(seconds), seconds <= 0:
 			logger.WarnContext(r.Context(), fmt.Sprintf("Invalid X-Prometheus-Scrape-Timeout-Seconds: %q. Defaulting timeout to %s", v, defaultScrapeTimeout))
 		default:
@@ -140,18 +143,18 @@ func (c *MetricsHTTPHandler) getScrapeTimeout(logger *slog.Logger, r *http.Reque
 
 	timeout = max(timeout-c.timeoutMargin, timeout/2)
 
-	return min(max(timeout, minScrapeTimeout), maxScrapeTimeout)
+	return min(max(timeout, minScrapeTimeout), MaxScrapeTimeout)
 }
 
-// toDuration converts seconds to a duration between 0 and maxScrapeTimeout.
+// toDuration converts seconds to a duration between 0 and MaxScrapeTimeout.
 // NaN and negative values become 0.
 func toDuration(seconds float64) time.Duration {
 	if math.IsNaN(seconds) || seconds <= 0 {
 		return 0
 	}
 
-	if seconds >= maxScrapeTimeout.Seconds() {
-		return maxScrapeTimeout
+	if seconds >= MaxScrapeTimeout.Seconds() {
+		return MaxScrapeTimeout
 	}
 
 	return time.Duration(seconds * float64(time.Second))
@@ -171,35 +174,50 @@ func (c *MetricsHTTPHandler) handlerFactory(ctx context.Context, logger *slog.Lo
 	}
 
 	// Scrapes are serialized by the collection, so the handler doesn't limit requests in flight.
-	var regHandler http.Handler
-	if c.exporterMetricsRegistry != nil {
-		regHandler = promhttp.HandlerFor(
-			mergedGatherer{first: c.exporterMetricsRegistry, second: reg},
-			promhttp.HandlerOpts{
-				ErrorLog:          slog.NewLogLogger(logger.Handler(), slog.LevelError),
-				ErrorHandling:     promhttp.ContinueOnError,
-				Registry:          c.exporterMetricsRegistry,
-				EnableOpenMetrics: true,
-				ProcessStartTime:  c.metricCollectors.GetStartTime(),
-			},
-		)
-
-		// Note that we have to use h.exporterMetricsRegistry here to
-		// use the same promhttp metrics for all expositions.
-		regHandler = promhttp.InstrumentMetricHandler(
-			c.exporterMetricsRegistry, regHandler,
-		)
-	} else {
-		regHandler = promhttp.HandlerFor(
-			reg,
-			promhttp.HandlerOpts{
-				ErrorLog:          slog.NewLogLogger(logger.Handler(), slog.LevelError),
-				ErrorHandling:     promhttp.ContinueOnError,
-				EnableOpenMetrics: true,
-				ProcessStartTime:  c.metricCollectors.GetStartTime(),
-			},
-		)
+	opts := promhttp.HandlerOpts{
+		ErrorLog:          slog.NewLogLogger(logger.Handler(), slog.LevelError),
+		ErrorHandling:     promhttp.ContinueOnError,
+		EnableOpenMetrics: true,
+		ProcessStartTime:  c.metricCollectors.GetStartTime(),
 	}
 
-	return regHandler, nil
+	if c.exporterMetricsRegistry == nil {
+		return canceledScrapeHandler(logger, reg, opts), nil
+	}
+
+	// Note that we have to use h.exporterMetricsRegistry here to
+	// use the same promhttp metrics for all expositions.
+	opts.Registry = c.exporterMetricsRegistry
+
+	return promhttp.InstrumentMetricHandler(
+		c.exporterMetricsRegistry,
+		canceledScrapeHandler(logger, mergedGatherer{first: c.exporterMetricsRegistry, second: reg}, opts),
+	), nil
+}
+
+// canceledScrapeHandler serves the metrics of gatherer like [promhttp.HandlerFor], but fails the
+// response with 503 Service Unavailable if the request was canceled while it gathered, for example
+// because the exporter shuts down. The collectors return early then, so the response would report
+// collectors as timed out or failed and miss their series. A failed scrape records up=0 instead.
+// A scrape that ends at its deadline is still served, with the timed-out collectors marked.
+func canceledScrapeHandler(logger *slog.Logger, gatherer prometheus.Gatherer, opts promhttp.HandlerOpts) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		mfs, err := gatherer.Gather()
+
+		if ctx := r.Context(); errors.Is(ctx.Err(), context.Canceled) {
+			logger.DebugContext(ctx, "scrape canceled, metrics not served",
+				slog.Any("cause", context.Cause(ctx)),
+			)
+
+			http.Error(w, "Scrape canceled: "+context.Cause(ctx).Error(), http.StatusServiceUnavailable)
+
+			return
+		}
+
+		gathered := prometheus.GathererFunc(func() ([]*dto.MetricFamily, error) {
+			return mfs, err
+		})
+
+		promhttp.HandlerFor(gathered, opts).ServeHTTP(w, r)
+	})
 }

internal/httphandler/httphandler_test.go

@@ -211,6 +211,43 @@ func TestMetricsHTTPHandlerScrapeTimeout(t *testing.T) {
 	}
 }
 
+func TestMetricsHTTPHandlerCanceledScrape(t *testing.T) {
+	for _, tc := range []struct {
+		name    string
+		options *httphandler.Options
+	}{
+		{name: "exporter metrics", options: &httphandler.Options{}},
+		{name: "exporter metrics disabled", options: &httphandler.Options{DisableExporterMetrics: true}},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			// The collector's delay and the cancellation use synthetic time.
+			synctest.Test(t, func(t *testing.T) {
+				collection := collector.New(collector.Map{"test": &testCollector{name: "test", delay: time.Second}})
+				handler := httphandler.New(slog.New(slog.DiscardHandler), collection, tc.options)
+
+				// The request is canceled while the collector runs, like at exporter shutdown.
+				ctx, cancel := context.WithCancel(t.Context())
+				time.AfterFunc(100*time.Millisecond, cancel)
+
+				request := httptest.NewRequestWithContext(ctx, http.MethodGet, "/metrics", nil)
+				request.Header.Set("X-Prometheus-Scrape-Timeout-Seconds", "10")
+
+				response := httptest.NewRecorder()
+				handler.ServeHTTP(response, request)
+
+				// A failed scrape records up=0 instead of a partial result with false collector failures.
+				require.Equal(t, http.StatusServiceUnavailable, response.Code)
+				require.Contains(t, response.Body.String(), "Scrape canceled")
+				require.NotContains(t, response.Body.String(), "windows_exporter_collector_success")
+
+				// Finish the delayed collector and its drain goroutine before leaving the bubble.
+				time.Sleep(time.Second)
+				synctest.Wait()
+			})
+		})
+	}
+}
+
 func TestMetricsHTTPHandlerScrapeBudget(t *testing.T) {
 	const maxScrapeTimeout = 5 * time.Minute