Fix stability and resource hygiene regressions

pkg/ruler/ruler_test.go

@@ -1613,13 +1613,9 @@ func TestGetRulesFromBackup(t *testing.T) {
 		require.Equal(t, a.Group.Interval, b.Group.Interval)
 		require.Equal(t, a.Group.User, b.Group.User)
 		require.Equal(t, a.Group.Limit, b.Group.Limit)
-		require.Equal(t, a.EvaluationTimestamp, b.EvaluationTimestamp)
-		require.Equal(t, a.EvaluationDuration, b.EvaluationDuration)
 		require.Equal(t, len(a.ActiveRules), len(b.ActiveRules))
 		for i, aRule := range a.ActiveRules {
 			bRule := b.ActiveRules[i]
-			require.Equal(t, aRule.EvaluationTimestamp, bRule.EvaluationTimestamp)
-			require.Equal(t, aRule.EvaluationDuration, bRule.EvaluationDuration)
 			require.Equal(t, aRule.Health, bRule.Health)
 			require.Equal(t, aRule.LastError, bRule.LastError)
 			require.Equal(t, aRule.Rule.Expr, bRule.Rule.Expr)
@@ -1841,13 +1837,9 @@ func getRulesHATest(replicationFactor int) func(t *testing.T) {
 			require.Equal(t, a.Group.Interval, b.Group.Interval)
 			require.Equal(t, a.Group.User, b.Group.User)
 			require.Equal(t, a.Group.Limit, b.Group.Limit)
-			require.Equal(t, a.EvaluationTimestamp, b.EvaluationTimestamp)
-			require.Equal(t, a.EvaluationDuration, b.EvaluationDuration)
 			require.Equal(t, len(a.ActiveRules), len(b.ActiveRules))
 			for i, aRule := range a.ActiveRules {
 				bRule := b.ActiveRules[i]
-				require.Equal(t, aRule.EvaluationTimestamp, bRule.EvaluationTimestamp)
-				require.Equal(t, aRule.EvaluationDuration, bRule.EvaluationDuration)
 				require.Equal(t, aRule.Health, bRule.Health)
 				require.Equal(t, aRule.LastError, bRule.LastError)
 				require.Equal(t, aRule.Rule.Expr, bRule.Rule.Expr)

pkg/scheduler/fragment_table/fragment_table.go

@@ -18,6 +18,8 @@ type FragmentTable struct {
 	mappings   map[distributed_execution.FragmentKey]*fragmentEntry
 	mu         sync.RWMutex
 	expiration time.Duration
+	stopCh     chan struct{}
+	stopOnce   sync.Once
 }
 
 // NewFragmentTable creates a new FragmentTable with the specified expiration duration.
@@ -27,11 +29,19 @@ func NewFragmentTable(expiration time.Duration) *FragmentTable {
 	ft := &FragmentTable{
 		mappings:   make(map[distributed_execution.FragmentKey]*fragmentEntry),
 		expiration: expiration,
+		stopCh:     make(chan struct{}),
 	}
 	go ft.periodicCleanup()
 	return ft
 }
 
+// Close stops the background cleanup goroutine.
+func (f *FragmentTable) Close() {
+	f.stopOnce.Do(func() {
+		close(f.stopCh)
+	})
+}
+
 // AddAddressByID associates a querier address with a specific fragment of a query.
 // The association will automatically expire after the configured duration.
 func (f *FragmentTable) AddAddressByID(queryID uint64, fragmentID uint64, addr string) {
@@ -73,7 +83,12 @@ func (f *FragmentTable) cleanupExpired() {
 func (f *FragmentTable) periodicCleanup() {
 	ticker := time.NewTicker(f.expiration / 2)
 	defer ticker.Stop()
-	for range ticker.C {
-		f.cleanupExpired()
+	for {
+		select {
+		case <-ticker.C:
+			f.cleanupExpired()
+		case <-f.stopCh:
+			return
+		}
 	}
 }

pkg/scheduler/fragment_table/fragment_table_test.go

@@ -27,6 +27,7 @@ func TestNewFragmentTable(t *testing.T) {
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			ft := NewFragmentTable(tt.expiration)
+			t.Cleanup(ft.Close)
 			require.NotNil(t, ft)
 			require.NotNil(t, ft.mappings)
 			assert.Equal(t, tt.expiration, ft.expiration)
@@ -36,6 +37,7 @@ func TestNewFragmentTable(t *testing.T) {
 
 func TestFragmentTable_AddAndGetAddress(t *testing.T) {
 	ft := NewFragmentTable(time.Hour)
+	t.Cleanup(ft.Close)
 
 	tests := []struct {
 		name      string
@@ -84,6 +86,7 @@ func TestFragmentTable_AddAndGetAddress(t *testing.T) {
 func TestFragmentTable_Expiration(t *testing.T) {
 	expiration := 100 * time.Millisecond
 	ft := NewFragmentTable(expiration)
+	t.Cleanup(ft.Close)
 
 	t.Run("entries expire after timeout", func(t *testing.T) {
 		ft.AddAddressByID(1, 1, "addr1")
@@ -106,6 +109,7 @@ func TestFragmentTable_Expiration(t *testing.T) {
 
 func TestFragmentTable_ConcurrentAccess(t *testing.T) {
 	ft := NewFragmentTable(time.Hour)
+	t.Cleanup(ft.Close)
 
 	const (
 		numGoroutines = 10
@@ -140,6 +144,7 @@ func TestFragmentTable_ConcurrentAccess(t *testing.T) {
 func TestFragmentTable_PeriodicCleanup(t *testing.T) {
 	expiration := 100 * time.Millisecond
 	ft := NewFragmentTable(expiration)
+	t.Cleanup(ft.Close)
 
 	ft.AddAddressByID(1, 1, "addr1")
 	ft.AddAddressByID(1, 2, "addr2")
@@ -163,3 +168,18 @@ func TestFragmentTable_PeriodicCleanup(t *testing.T) {
 	_, ok = ft.GetAddrByID(1, 2)
 	require.False(t, ok)
 }
+
+func TestFragmentTable_CloseStopsPeriodicCleanup(t *testing.T) {
+	expiration := 100 * time.Millisecond
+	ft := NewFragmentTable(expiration)
+	ft.Close()
+	ft.Close()
+
+	ft.AddAddressByID(1, 1, "addr1")
+
+	time.Sleep(expiration * 2)
+
+	addr, ok := ft.GetAddrByID(1, 1)
+	require.True(t, ok)
+	require.Equal(t, "addr1", addr)
+}

pkg/scheduler/scheduler.go

@@ -698,6 +698,7 @@ func (s *Scheduler) running(ctx context.Context) error {
 
 // Close the Scheduler.
 func (s *Scheduler) stopping(_ error) error {
+	s.fragmentTable.Close()
 	// This will also stop the requests queue, which stop accepting new requests and errors out any tracked requests.
 	return services.StopManagerAndAwaitStopped(context.Background(), s.subservices)
 }

pkg/util/http_test.go

@@ -10,6 +10,7 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"strconv"
+	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -223,13 +224,29 @@ func TestIsRequestBodyTooLargeRegression(t *testing.T) {
 }
 
 func TestParseProtoReader_GzipDecompressionBomb(t *testing.T) {
-	// Create a gzip payload where decompressed size far exceeds maxSize.
-	const maxSize = 4096                // 4 KB limit on decompressed output
-	uncompressed := make([]byte, 1<<20) // 1 MB of zeros
+	const maxSize = 4096
+	req := &cortexpb.PreallocWriteRequest{
+		WriteRequest: cortexpb.WriteRequest{
+			Timeseries: []cortexpb.PreallocTimeseries{
+				{
+					TimeSeries: &cortexpb.TimeSeries{
+						Labels: []cortexpb.LabelAdapter{
+							{Name: "__name__", Value: "large_series"},
+							{Name: "payload", Value: strings.Repeat("a", 1<<20)},
+						},
+						Samples: []cortexpb.Sample{{Value: 1, TimestampMs: 1}},
+					},
+				},
+			},
+		},
+	}
+	uncompressed, err := req.Marshal()
+	require.NoError(t, err)
+	require.Greater(t, len(uncompressed), maxSize)
 
 	var compressed bytes.Buffer
 	gzw := gzip.NewWriter(&compressed)
-	_, err := gzw.Write(uncompressed)
+	_, err = gzw.Write(uncompressed)
 	require.NoError(t, err)
 	require.NoError(t, gzw.Close())
 
@@ -239,8 +256,7 @@ func TestParseProtoReader_GzipDecompressionBomb(t *testing.T) {
 
 	var fromWire cortexpb.PreallocWriteRequest
 	err = util.ParseProtoReader(context.Background(), io.NopCloser(&compressed), 0, maxSize, &fromWire, util.Gzip)
-	// The decompressed output should be limited to maxSize+1 bytes, causing a
-	// proto unmarshal error (not an OOM). The key assertion is that we don't
-	// allocate 1 MB of memory.
-	assert.NotNil(t, err)
+	// The decompressed output should be limited to maxSize+1 bytes. Without
+	// the gzip LimitReader, the complete valid protobuf would unmarshal.
+	require.Error(t, err)
 }

pkg/util/queryeviction/evictor_test.go

@@ -207,7 +207,9 @@ func TestPrometheusMetrics_IncrementedCorrectly(t *testing.T) {
 		waitEvicted(t, evicted)
 	}
 
-	assert.Equal(t, float64(3), promtest.ToFloat64(evictor.evictionsTotal.WithLabelValues(string(resource.CPU))))
+	require.Eventually(t, func() bool {
+		return promtest.ToFloat64(evictor.evictionsTotal.WithLabelValues(string(resource.CPU))) == 3
+	}, time.Second, 10*time.Millisecond)
 }
 
 func TestNewQueryEvictor_ReturnsNilWhenDisabled(t *testing.T) {