Add support for evicting multiple queries per cycle

Signed-off-by: Essam Eldaly <eeldaly@amazon.com>

Author: eeldaly

pkg/configs/query_protection.go

@@ -36,11 +36,12 @@ type Threshold struct {
 
 // EvictionConfig configures the resource-based query evictor.
 type EvictionConfig struct {
-	Threshold      Threshold     `yaml:"threshold"`
-	CheckInterval  time.Duration `yaml:"check_interval"`
-	CooldownPeriod int           `yaml:"cooldown_period"`
-	EvictionMetric string        `yaml:"eviction_metric"`
-	MinQueryAge    time.Duration `yaml:"min_query_age"`
+	Threshold            Threshold     `yaml:"threshold"`
+	CheckInterval        time.Duration `yaml:"check_interval"`
+	CooldownPeriod       int           `yaml:"cooldown_period"`
+	EvictionMetric       string        `yaml:"eviction_metric"`
+	MinQueryAge          time.Duration `yaml:"min_query_age"`
+	MaxEvictionsPerCycle int           `yaml:"max_evictions_per_cycle"`
 }
 
 // Enabled returns true when at least one eviction threshold is greater than 0.
@@ -60,6 +61,7 @@ func (cfg *QueryProtection) RegisterFlagsWithPrefix(f *flag.FlagSet, prefix stri
 	f.IntVar(&cfg.Eviction.CooldownPeriod, prefix+"query-protection.eviction.cooldown-period", 3, "EXPERIMENTAL: Number of check intervals to wait after an eviction before evicting again.")
 	f.StringVar(&cfg.Eviction.EvictionMetric, prefix+"query-protection.eviction.eviction-metric", "fetched_samples", "EXPERIMENTAL: The query metric used to determine the heaviest query for eviction. Supported values: fetched_samples, fetched_series, fetched_chunks, fetched_chunk_bytes.")
 	f.DurationVar(&cfg.Eviction.MinQueryAge, prefix+"query-protection.eviction.min-query-age", 10*time.Second, "EXPERIMENTAL: Minimum time a query must be running before it becomes eligible for eviction. Queries younger than this are ignored.")
+	f.IntVar(&cfg.Eviction.MaxEvictionsPerCycle, prefix+"query-protection.eviction.max-evictions-per-cycle", 1, "EXPERIMENTAL: Maximum number of queries to evict in a single check cycle when resource thresholds are breached.")
 }
 
 func (cfg *QueryProtection) Validate(monitoredResources flagext.StringSliceCSV) error {
@@ -104,6 +106,10 @@ func (cfg *QueryProtection) Validate(monitoredResources flagext.StringSliceCSV)
 			return fmt.Errorf("unrecognized eviction_metric %q; supported values: fetched_samples, fetched_series, fetched_chunks, fetched_chunk_bytes", cfg.Eviction.EvictionMetric)
 		}
 
+		if cfg.Eviction.MaxEvictionsPerCycle < 1 {
+			return errors.New("eviction max_evictions_per_cycle must be >= 1")
+		}
+
 		if evThreshold.CPUUtilization > 0 && !strings.Contains(monitoredResources.String(), string(resource.CPU)) {
 			return errors.New("monitored_resources config must include \"cpu\" when eviction cpu threshold is set")
 		}

pkg/configs/query_protection_test.go

@@ -99,10 +99,11 @@ func Test_EvictionConfig_Validation(t *testing.T) {
 	validBase := func() QueryProtection {
 		return QueryProtection{
 			Eviction: EvictionConfig{
-				Threshold:      Threshold{CPUUtilization: 0.8, HeapUtilization: 0.85},
-				CheckInterval:  1 * time.Second,
-				CooldownPeriod: 3,
-				EvictionMetric: "fetched_samples",
+				Threshold:            Threshold{CPUUtilization: 0.8, HeapUtilization: 0.85},
+				CheckInterval:        1 * time.Second,
+				CooldownPeriod:       3,
+				EvictionMetric:       "fetched_samples",
+				MaxEvictionsPerCycle: 1,
 			},
 		}
 	}
@@ -112,17 +113,18 @@ func Test_EvictionConfig_Validation(t *testing.T) {
 		monitoredResources []string
 		err                string
 	}{
-		"valid config passes":          {func(qp *QueryProtection) {}, []string{"cpu", "heap"}, ""},
-		"cpu > 1 fails":                {func(qp *QueryProtection) { qp.Eviction.Threshold.CPUUtilization = 1.5 }, []string{"cpu", "heap"}, "eviction cpu_utilization must be between 0 and 1"},
-		"cpu < 0 fails":                {func(qp *QueryProtection) { qp.Eviction.Threshold.CPUUtilization = -0.1 }, []string{"cpu", "heap"}, "eviction cpu_utilization must be between 0 and 1"},
-		"heap > 1 fails":               {func(qp *QueryProtection) { qp.Eviction.Threshold.HeapUtilization = 2.0 }, []string{"cpu", "heap"}, "eviction heap_utilization must be between 0 and 1"},
-		"heap < 0 fails":               {func(qp *QueryProtection) { qp.Eviction.Threshold.HeapUtilization = -0.5 }, []string{"cpu", "heap"}, "eviction heap_utilization must be between 0 and 1"},
-		"check_interval 0 fails":       {func(qp *QueryProtection) { qp.Eviction.CheckInterval = 0 }, []string{"cpu", "heap"}, "eviction check_interval must be greater than 0 when eviction is enabled"},
-		"cooldown < 0 fails":           {func(qp *QueryProtection) { qp.Eviction.CooldownPeriod = -1 }, []string{"cpu", "heap"}, "eviction cooldown_period must be >= 0"},
-		"unknown metric fails":         {func(qp *QueryProtection) { qp.Eviction.EvictionMetric = "unknown" }, []string{"cpu", "heap"}, `unrecognized eviction_metric "unknown"; supported values: fetched_samples, fetched_series, fetched_chunks, fetched_chunk_bytes`},
-		"cpu without monitored fails":  {func(qp *QueryProtection) {}, []string{"heap"}, `monitored_resources config must include "cpu" when eviction cpu threshold is set`},
-		"heap without monitored fails": {func(qp *QueryProtection) {}, []string{"cpu"}, `monitored_resources config must include "heap" when eviction heap threshold is set`},
-		"cooldown 0 is valid":          {func(qp *QueryProtection) { qp.Eviction.CooldownPeriod = 0 }, []string{"cpu", "heap"}, ""},
+		"valid config passes":               {func(qp *QueryProtection) {}, []string{"cpu", "heap"}, ""},
+		"cpu > 1 fails":                     {func(qp *QueryProtection) { qp.Eviction.Threshold.CPUUtilization = 1.5 }, []string{"cpu", "heap"}, "eviction cpu_utilization must be between 0 and 1"},
+		"cpu < 0 fails":                     {func(qp *QueryProtection) { qp.Eviction.Threshold.CPUUtilization = -0.1 }, []string{"cpu", "heap"}, "eviction cpu_utilization must be between 0 and 1"},
+		"heap > 1 fails":                    {func(qp *QueryProtection) { qp.Eviction.Threshold.HeapUtilization = 2.0 }, []string{"cpu", "heap"}, "eviction heap_utilization must be between 0 and 1"},
+		"heap < 0 fails":                    {func(qp *QueryProtection) { qp.Eviction.Threshold.HeapUtilization = -0.5 }, []string{"cpu", "heap"}, "eviction heap_utilization must be between 0 and 1"},
+		"check_interval 0 fails":            {func(qp *QueryProtection) { qp.Eviction.CheckInterval = 0 }, []string{"cpu", "heap"}, "eviction check_interval must be greater than 0 when eviction is enabled"},
+		"cooldown < 0 fails":                {func(qp *QueryProtection) { qp.Eviction.CooldownPeriod = -1 }, []string{"cpu", "heap"}, "eviction cooldown_period must be >= 0"},
+		"unknown metric fails":              {func(qp *QueryProtection) { qp.Eviction.EvictionMetric = "unknown" }, []string{"cpu", "heap"}, `unrecognized eviction_metric "unknown"; supported values: fetched_samples, fetched_series, fetched_chunks, fetched_chunk_bytes`},
+		"cpu without monitored fails":       {func(qp *QueryProtection) {}, []string{"heap"}, `monitored_resources config must include "cpu" when eviction cpu threshold is set`},
+		"heap without monitored fails":      {func(qp *QueryProtection) {}, []string{"cpu"}, `monitored_resources config must include "heap" when eviction heap threshold is set`},
+		"cooldown 0 is valid":               {func(qp *QueryProtection) { qp.Eviction.CooldownPeriod = 0 }, []string{"cpu", "heap"}, ""},
+		"max_evictions_per_cycle < 1 fails": {func(qp *QueryProtection) { qp.Eviction.MaxEvictionsPerCycle = 0 }, []string{"cpu", "heap"}, "eviction max_evictions_per_cycle must be >= 1"},
 		"disabled skips interval check": {func(qp *QueryProtection) {
 			qp.Eviction.Threshold = Threshold{}
 			qp.Eviction.CheckInterval = 0
@@ -154,5 +156,6 @@ func Test_RegisterFlagsWithPrefix_EvictionDefaults(t *testing.T) {
 	assert.Equal(t, 1*time.Second, cfg.Eviction.CheckInterval)
 	assert.Equal(t, 3, cfg.Eviction.CooldownPeriod)
 	assert.Equal(t, "fetched_samples", cfg.Eviction.EvictionMetric)
+	assert.Equal(t, 1, cfg.Eviction.MaxEvictionsPerCycle)
 	assert.False(t, cfg.Eviction.Enabled())
 }

pkg/querier/querier_eviction_test.go

@@ -85,9 +85,10 @@ func TestQuerier_EvictionIntegration(t *testing.T) {
 						CPUUtilization:  tc.evictionCPU,
 						HeapUtilization: tc.evictionHeap,
 					},
-					CheckInterval:  1 * time.Second,
-					CooldownPeriod: 3,
-					EvictionMetric: "fetched_samples",
+					CheckInterval:        1 * time.Second,
+					CooldownPeriod:       3,
+					EvictionMetric:       "fetched_samples",
+					MaxEvictionsPerCycle: 1,
 				},
 			}
 

pkg/util/queryeviction/engine_wrapper_test.go

@@ -118,9 +118,9 @@ func TestEngineWrapper_EvictedQueryReturnsErrQueryEvicted(t *testing.T) {
 	mq := &mockQuery{
 		execFunc: func(ctx context.Context) *promql.Result {
 			// Simulate eviction: find the registered query and cancel it.
-			heaviest := registry.FindHeaviest(0)
-			require.NotNil(t, heaviest, "query should be registered during Exec")
-			heaviest.Cancel() // This cancels the child context, simulating evictor behavior.
+			victims := registry.FindHeaviest(1, 0)
+			require.Len(t, victims, 1, "query should be registered during Exec")
+			victims[0].Cancel() // This cancels the child context, simulating evictor behavior.
 
 			// The inner query would see a cancelled context and return an error.
 			return &promql.Result{Err: context.Canceled}

pkg/util/queryeviction/evictor.go

@@ -82,31 +82,31 @@ func (e *QueryEvictor) running(ctx context.Context) error {
 				continue // no breach
 			}
 
-			// Find the heaviest running query.
-			heaviest := e.registry.FindHeaviest(e.cfg.MinQueryAge)
-			if heaviest == nil {
+			// Find the heaviest running queries (up to MaxEvictionsPerCycle).
+			victims := e.registry.FindHeaviest(e.cfg.MaxEvictionsPerCycle, e.cfg.MinQueryAge)
+			if len(victims) == 0 {
 				continue // no running queries to evict
 			}
 
-			// Evict the heaviest query.
-			metricValue := e.registry.metric(heaviest.Stats)
-			heaviest.Cancel()
-
-			// Log the eviction.
-			level.Warn(e.logger).Log(
-				"msg", "evicting heaviest query due to resource pressure",
-				"resource", breachedResource,
-				"utilization", utilization,
-				"threshold", threshold,
-				"request_id", heaviest.RequestID,
-				"query", heaviest.QueryExpr,
-				"user", heaviest.UserID,
-				"metric", e.cfg.EvictionMetric,
-				"metric_value", metricValue,
-			)
-
-			// Increment metrics.
-			e.evictionsTotal.WithLabelValues(string(breachedResource)).Inc()
+			// Evict each victim.
+			for _, victim := range victims {
+				metricValue := e.registry.metric(victim.Stats)
+				victim.Cancel()
+
+				level.Warn(e.logger).Log(
+					"msg", "evicting query due to resource pressure",
+					"resource", breachedResource,
+					"utilization", utilization,
+					"threshold", threshold,
+					"request_id", victim.RequestID,
+					"query", victim.QueryExpr,
+					"user", victim.UserID,
+					"metric", e.cfg.EvictionMetric,
+					"metric_value", metricValue,
+				)
+
+				e.evictionsTotal.WithLabelValues(string(breachedResource)).Inc()
+			}
 
 			// Enter cooldown.
 			cooldownRemaining = e.cfg.CooldownPeriod

pkg/util/queryeviction/evictor_test.go

@@ -39,9 +39,10 @@ func testEvictorConfig(cpu, heap float64, cooldown int) configs.EvictionConfig {
 			CPUUtilization:  cpu,
 			HeapUtilization: heap,
 		},
-		CheckInterval:  10 * time.Millisecond,
-		CooldownPeriod: cooldown,
-		EvictionMetric: "fetched_samples",
+		CheckInterval:        10 * time.Millisecond,
+		CooldownPeriod:       cooldown,
+		EvictionMetric:       "fetched_samples",
+		MaxEvictionsPerCycle: 1,
 	}
 }
 
@@ -178,8 +179,9 @@ func TestCheckThresholds(t *testing.T) {
 					CPUUtilization:  tc.cpuThresh,
 					HeapUtilization: tc.heapThresh,
 				},
-				CheckInterval:  time.Second,
-				EvictionMetric: "fetched_samples",
+				CheckInterval:        time.Second,
+				EvictionMetric:       "fetched_samples",
+				MaxEvictionsPerCycle: 1,
 			}
 
 			var evictor *QueryEvictor
@@ -212,7 +214,7 @@ func TestPrometheusMetrics_IncrementedCorrectly(t *testing.T) {
 }
 
 func TestNewQueryEvictor_ReturnsNilWhenDisabled(t *testing.T) {
-	cfg := configs.EvictionConfig{CheckInterval: time.Second, EvictionMetric: "fetched_samples"}
+	cfg := configs.EvictionConfig{CheckInterval: time.Second, EvictionMetric: "fetched_samples", MaxEvictionsPerCycle: 1}
 	evictor, err := NewQueryEvictor(newMockMonitor(0, 0), NewQueryRegistry(testMetricFunc), cfg, log.NewNopLogger(), prometheus.NewPedanticRegistry(), "test")
 	assert.NoError(t, err)
 	assert.Nil(t, evictor)
@@ -236,3 +238,29 @@ func TestEviction_HeaviestQueryIsEvicted(t *testing.T) {
 		t.Fatal("expected heaviest query to be evicted")
 	}
 }
+
+func TestEviction_MultipleQueriesPerCycle(t *testing.T) {
+	reg := NewQueryRegistry(testMetricFunc)
+	_, evictedSmall := registerTestQuery(reg, 100, "small", "user1")
+	_, evictedLarge := registerTestQuery(reg, 10000, "large", "user2")
+	_, evictedMedium := registerTestQuery(reg, 500, "medium", "user3")
+
+	cfg := configs.EvictionConfig{
+		Threshold: configs.Threshold{
+			CPUUtilization: 0.9,
+		},
+		CheckInterval:        10 * time.Millisecond,
+		CooldownPeriod:       300,
+		EvictionMetric:       "fetched_samples",
+		MaxEvictionsPerCycle: 2,
+	}
+
+	startEvictor(t, newMockMonitor(0.95, 0.0), reg, cfg)
+
+	// The two heaviest (large=10000, medium=500) should be evicted in the same cycle.
+	waitEvicted(t, evictedLarge)
+	waitEvicted(t, evictedMedium)
+
+	// The smallest should not be evicted because max per cycle is 2 and cooldown prevents another cycle.
+	assertNotEvicted(t, evictedSmall, 50*time.Millisecond)
+}

pkg/util/queryeviction/registry.go

@@ -1,6 +1,7 @@
 package queryeviction
 
 import (
+	"container/heap"
 	"context"
 	"fmt"
 	"sync"
@@ -80,29 +81,61 @@ func (r *QueryRegistry) Deregister(id uint64) {
 	delete(r.queries, id)
 }
 
-// FindHeaviest returns the entry with the highest metric value
+// FindHeaviest returns up to n entries with the highest metric values
 // among queries that have been running for at least minAge,
-// or nil if no eligible queries exist.
-func (r *QueryRegistry) FindHeaviest(minAge time.Duration) *QueryEntry {
+// sorted heaviest first. Returns nil if no eligible queries exist.
+func (r *QueryRegistry) FindHeaviest(n int, minAge time.Duration) []*QueryEntry {
 	r.mu.RLock()
 	defer r.mu.RUnlock()
 
-	var heaviest *QueryEntry
-	var maxWeight uint64
 	now := time.Now()
 
+	// Use a min-heap of size n to find the top-n queries
+	h := &weightedHeap{}
+
 	for _, entry := range r.queries {
 		if now.Sub(entry.RegisteredAt) < minAge {
 			continue
 		}
-		weight := r.metric(entry.Stats)
-		if heaviest == nil || weight > maxWeight {
-			heaviest = entry
-			maxWeight = weight
+		w := weighted{entry: entry, weight: r.metric(entry.Stats)}
+
+		if h.Len() < n {
+			heap.Push(h, w)
+		} else if w.weight > (*h)[0].weight {
+			(*h)[0] = w
+			heap.Fix(h, 0)
 		}
 	}
 
-	return heaviest
+	if h.Len() == 0 {
+		return nil
+	}
+
+	result := make([]*QueryEntry, h.Len())
+	for i := len(result) - 1; i >= 0; i-- {
+		result[i] = heap.Pop(h).(weighted).entry
+	}
+	return result
+}
+
+// weightedHeap is a min-heap of weighted entries (smallest weight at root).
+type weightedHeap []weighted
+
+type weighted struct {
+	entry  *QueryEntry
+	weight uint64
+}
+
+func (h weightedHeap) Len() int            { return len(h) }
+func (h weightedHeap) Less(i, j int) bool  { return h[i].weight < h[j].weight }
+func (h weightedHeap) Swap(i, j int)       { h[i], h[j] = h[j], h[i] }
+func (h *weightedHeap) Push(x interface{}) { *h = append(*h, x.(weighted)) }
+func (h *weightedHeap) Pop() interface{} {
+	old := *h
+	n := len(old)
+	x := old[n-1]
+	*h = old[:n-1]
+	return x
 }
 
 // Len returns the number of currently registered queries.

pkg/util/queryeviction/registry_test.go

@@ -54,14 +54,14 @@ func TestDeregister_RemovesEntry(t *testing.T) {
 	require.Equal(t, 1, reg.Len())
 
 	// FindHeaviest should return the registered entry.
-	heaviest := reg.FindHeaviest(0)
-	require.NotNil(t, heaviest)
-	assert.Equal(t, id, heaviest.QueryID)
+	results := reg.FindHeaviest(1, 0)
+	require.Len(t, results, 1)
+	assert.Equal(t, id, results[0].QueryID)
 
 	// Deregister and verify it's gone.
 	reg.Deregister(id)
 	assert.Equal(t, 0, reg.Len())
-	assert.Nil(t, reg.FindHeaviest(0), "FindHeaviest should return nil after deregistering the only entry")
+	assert.Nil(t, reg.FindHeaviest(1, 0), "FindHeaviest should return nil after deregistering the only entry")
 }
 
 func TestDeregister_UnknownID_IsNoOp(t *testing.T) {
@@ -76,9 +76,9 @@ func TestDeregister_UnknownID_IsNoOp(t *testing.T) {
 
 	// Original entry should still be present.
 	assert.Equal(t, 1, reg.Len())
-	heaviest := reg.FindHeaviest(0)
-	require.NotNil(t, heaviest)
-	assert.Equal(t, id, heaviest.QueryID)
+	results := reg.FindHeaviest(1, 0)
+	require.Len(t, results, 1)
+	assert.Equal(t, id, results[0].QueryID)
 }
 
 func TestFindHeaviest_ReturnsHighestMetricValue(t *testing.T) {
@@ -91,16 +91,16 @@ func TestFindHeaviest_ReturnsHighestMetricValue(t *testing.T) {
 	heaviestID := reg.Register(cancel, newTestStats(1000), "large-query", "user3", "")
 	reg.Register(cancel, newTestStats(200), "another-query", "user4", "")
 
-	heaviest := reg.FindHeaviest(0)
-	require.NotNil(t, heaviest)
-	assert.Equal(t, heaviestID, heaviest.QueryID)
-	assert.Equal(t, "large-query", heaviest.QueryExpr)
-	assert.Equal(t, uint64(1000), heaviest.Stats.LoadFetchedSamples())
+	results := reg.FindHeaviest(1, 0)
+	require.Len(t, results, 1)
+	assert.Equal(t, heaviestID, results[0].QueryID)
+	assert.Equal(t, "large-query", results[0].QueryExpr)
+	assert.Equal(t, uint64(1000), results[0].Stats.LoadFetchedSamples())
 }
 
 func TestFindHeaviest_EmptyRegistry(t *testing.T) {
 	reg := NewQueryRegistry(testMetricFunc)
-	assert.Nil(t, reg.FindHeaviest(0), "FindHeaviest should return nil for empty registry")
+	assert.Nil(t, reg.FindHeaviest(1, 0), "FindHeaviest should return nil for empty registry")
 }
 
 func TestLen_ReflectsCurrentCount(t *testing.T) {
@@ -147,7 +147,7 @@ func TestConcurrent_RegisterDeregisterFindHeaviest(t *testing.T) {
 				id := reg.Register(cancel, stats, "concurrent-query", "user", "")
 
 				// Interleave FindHeaviest and Len calls.
-				_ = reg.FindHeaviest(0)
+				_ = reg.FindHeaviest(1, 0)
 				_ = reg.Len()
 
 				reg.Deregister(id)