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()