Add CrashOnPanic middleware + wrap distributor push routes

CHANGELOG.md

@@ -48,6 +48,7 @@
 * [ENHANCEMENT] Querier/Ingester: Detach ingester series from gRPC buffers to reduce heap. #7670
 * [ENHANCEMENT] Ingester: Add `cortex_ingester_tsdb_head_max_timestamp` metric that re-exports the TSDB head max timestamp (`prometheus_tsdb_head_max_time`) per user, to help investigate ingestion issues like out-of-bounds (too old sample) errors. #7694
 * [ENHANCEMENT] Ingester: Include the TSDB head max time in the `out of bounds` and `too old sample` error messages, so that users can see how far behind the accepted time range a rejected sample is. #7695
+* [ENHANCEMENT] Distributor: Add CrashOnPanic middleware + wrap distributor push routes. #7709
 * [ENHANCEMENT] Compactor: Reduce object storage GET calls when updating the bucket index by skipping re-reading parquet converter markers for blocks that already have a valid-version parquet entry in the previous index. #7669
 * [BUGFIX] Querier: Fix queryWithRetry and labelsWithRetry returning (nil, nil) on cancelled context by propagating ctx.Err(). #7370
 * [BUGFIX] Metrics Helper: Fix non-deterministic bucket order in merged histograms by sorting buckets after map iteration, matching Prometheus client library behavior. #7380

pkg/api/api.go

@@ -43,6 +43,7 @@ import (
 	"github.com/cortexproject/cortex/pkg/storegateway"
 	"github.com/cortexproject/cortex/pkg/storegateway/storegatewaypb"
 	"github.com/cortexproject/cortex/pkg/util/flagext"
+	utilmiddleware "github.com/cortexproject/cortex/pkg/util/middleware"
 	"github.com/cortexproject/cortex/pkg/util/push"
 	"github.com/cortexproject/cortex/pkg/util/validation"
 )
@@ -293,7 +294,7 @@ func (a *API) RegisterDistributor(d *distributor.Distributor, pushConfig distrib
 		Help:      "Total number of push requests by type.",
 	}, []string{"type"})
 
-	a.RegisterRoute("/api/v1/push", push.Handler(pushConfig.RemoteWriteV2Enabled, pushConfig.AcceptUnknownRemoteWriteContentType, pushConfig.MaxRecvMsgSize, overrides, a.sourceIPs, a.cfg.wrapDistributorPush(d), requestTotal), true, "POST")
+	a.RegisterRoute("/api/v1/push", utilmiddleware.CrashOnPanic(a.logger, push.Handler(pushConfig.RemoteWriteV2Enabled, pushConfig.AcceptUnknownRemoteWriteContentType, pushConfig.MaxRecvMsgSize, overrides, a.sourceIPs, a.cfg.wrapDistributorPush(d), requestTotal)), true, "POST")
 	a.RegisterRoute("/api/v1/otlp/v1/metrics", push.OTLPHandler(pushConfig.OTLPMaxRecvMsgSize, overrides, pushConfig.OTLPConfig, a.sourceIPs, a.cfg.wrapDistributorPush(d), requestTotal), true, "POST")
 
 	a.indexPage.AddLink(SectionAdminEndpoints, "/distributor/ring", "Distributor Ring Status")
@@ -305,7 +306,7 @@ func (a *API) RegisterDistributor(d *distributor.Distributor, pushConfig distrib
 	a.RegisterRoute("/distributor/ha_tracker", d.HATracker, false, "GET")
 
 	// Legacy Routes
-	a.RegisterRoute(path.Join(a.cfg.LegacyHTTPPrefix, "/push"), push.Handler(pushConfig.RemoteWriteV2Enabled, pushConfig.AcceptUnknownRemoteWriteContentType, pushConfig.MaxRecvMsgSize, overrides, a.sourceIPs, a.cfg.wrapDistributorPush(d), requestTotal), true, "POST")
+	a.RegisterRoute(path.Join(a.cfg.LegacyHTTPPrefix, "/push"), utilmiddleware.CrashOnPanic(a.logger, push.Handler(pushConfig.RemoteWriteV2Enabled, pushConfig.AcceptUnknownRemoteWriteContentType, pushConfig.MaxRecvMsgSize, overrides, a.sourceIPs, a.cfg.wrapDistributorPush(d), requestTotal)), true, "POST")
 	a.RegisterRoute("/all_user_stats", http.HandlerFunc(d.AllUserStatsHandler), false, "GET")
 	a.RegisterRoute("/ha-tracker", d.HATracker, false, "GET")
 }

pkg/util/middleware/crash_on_panic.go

@@ -0,0 +1,44 @@
+package middleware
+
+import (
+	"net/http"
+	"os"
+	"runtime"
+
+	"github.com/go-kit/log"
+	"github.com/go-kit/log/level"
+)
+
+// CrashOnPanic wraps an http.Handler and terminates the process if the handler
+// panics. This is intended for stateless components (e.g., distributors) where
+// a panic may leave the process in a corrupted state. Rather than continuing to
+// serve bad traffic, the process exits and allows the orchestrator (e.g.,
+// Kubernetes) to restart it with a clean state.
+func CrashOnPanic(logger log.Logger, next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		defer crashOnPanic(logger)
+		next.ServeHTTP(w, r)
+	})
+}
+
+func crashOnPanic(logger log.Logger) {
+	p := recover()
+	if p == nil {
+		return
+	}
+
+	// Capture the stack trace before exiting.
+	const stackSize = 8192
+	stack := make([]byte, stackSize)
+	n := runtime.Stack(stack, false)
+
+	level.Error(logger).Log(
+		"msg", "panic detected, crashing process",
+		"panic", p,
+		"stack", string(stack[:n]),
+	)
+
+	// Flush stderr so the stack trace is visible in logs before exit.
+	os.Stderr.Sync() //nolint:errcheck
+	os.Exit(1)
+}

pkg/util/middleware/crash_on_panic_test.go

@@ -0,0 +1,59 @@
+package middleware
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"os"
+	"os/exec"
+	"testing"
+
+	"github.com/go-kit/log"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestCrashOnPanic_NoPanic(t *testing.T) {
+	logger := log.NewNopLogger()
+
+	handler := CrashOnPanic(logger, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+		_, _ = w.Write([]byte("ok"))
+	}))
+
+	req := httptest.NewRequest(http.MethodPost, "/push", nil)
+	rec := httptest.NewRecorder()
+	handler.ServeHTTP(rec, req)
+
+	assert.Equal(t, http.StatusOK, rec.Code)
+	assert.Equal(t, "ok", rec.Body.String())
+}
+
+func TestCrashOnPanic_WithPanic(t *testing.T) {
+	// os.Exit cannot be tested directly. Use the subprocess test pattern:
+	// run this test binary with a special env var, and verify the exit code.
+	if os.Getenv("TEST_CRASH_ON_PANIC") == "1" {
+		logger := log.NewNopLogger()
+
+		handler := CrashOnPanic(logger, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			panic("runtime error: growslice: len out of range")
+		}))
+
+		req := httptest.NewRequest(http.MethodPost, "/push", nil)
+		rec := httptest.NewRecorder()
+		handler.ServeHTTP(rec, req)
+		// Should not reach here.
+		return
+	}
+
+	// Re-run this test in a subprocess with the env var set.
+	cmd := exec.Command(os.Args[0], "-test.run=TestCrashOnPanic_WithPanic")
+	cmd.Env = append(os.Environ(), "TEST_CRASH_ON_PANIC=1")
+	err := cmd.Run()
+
+	// The subprocess should exit with a non-zero exit code.
+	require.Error(t, err)
+
+	exitErr, ok := err.(*exec.ExitError)
+	require.True(t, ok, "expected *exec.ExitError, got %T", err)
+	assert.False(t, exitErr.Success(), "process should have exited with non-zero status")
+}