From 83225c86b7203f7bfd914b9c3ef1632c36b5ab73 Mon Sep 17 00:00:00 2001 From: Rohithmatham12 Date: Wed, 22 Jul 2026 11:06:49 -0400 Subject: [PATCH] Fix stability and resource hygiene regressions Signed-off-by: Rohithmatham12 --- pkg/ruler/ruler_test.go | 8 ----- .../fragment_table/fragment_table.go | 19 +++++++++-- .../fragment_table/fragment_table_test.go | 20 ++++++++++++ pkg/scheduler/scheduler.go | 1 + pkg/util/http_test.go | 32 ++++++++++++++----- pkg/util/queryeviction/evictor_test.go | 4 ++- 6 files changed, 65 insertions(+), 19 deletions(-) diff --git a/pkg/ruler/ruler_test.go b/pkg/ruler/ruler_test.go index cf42a970677..70f22d457d8 100644 --- a/pkg/ruler/ruler_test.go +++ b/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) diff --git a/pkg/scheduler/fragment_table/fragment_table.go b/pkg/scheduler/fragment_table/fragment_table.go index 3cf161070b8..2b6bb4ae2fc 100644 --- a/pkg/scheduler/fragment_table/fragment_table.go +++ b/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 + } } } diff --git a/pkg/scheduler/fragment_table/fragment_table_test.go b/pkg/scheduler/fragment_table/fragment_table_test.go index 576e2726455..6b307905fa1 100644 --- a/pkg/scheduler/fragment_table/fragment_table_test.go +++ b/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) +} diff --git a/pkg/scheduler/scheduler.go b/pkg/scheduler/scheduler.go index d84cdb5f8df..c46ac1bcd1b 100644 --- a/pkg/scheduler/scheduler.go +++ b/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) } diff --git a/pkg/util/http_test.go b/pkg/util/http_test.go index d2a5e128d4f..42cbd391fdd 100644 --- a/pkg/util/http_test.go +++ b/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) } diff --git a/pkg/util/queryeviction/evictor_test.go b/pkg/util/queryeviction/evictor_test.go index a6bbe924919..3dba2bdd499 100644 --- a/pkg/util/queryeviction/evictor_test.go +++ b/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) {