Ensure that Query is closed for rules evaluations (#18733)

Signed-off-by: Andrew Hall <andrew.hall@grafana.com>

rules/manager.go

@@ -53,6 +53,8 @@ func EngineQueryFunc(engine promql.QueryEngine, q storage.Queryable) QueryFunc {
 		if err != nil {
 			return nil, err
 		}
+		defer q.Close()
+
 		res := q.Exec(ctx)
 		if res.Err != nil {
 			return nil, res.Err

rules/manager_test.go

@@ -43,11 +43,13 @@ import (
 	"github.com/prometheus/prometheus/model/timestamp"
 	"github.com/prometheus/prometheus/model/value"
 	"github.com/prometheus/prometheus/promql"
+	"github.com/prometheus/prometheus/promql/parser"
 	"github.com/prometheus/prometheus/promql/promqltest"
 	"github.com/prometheus/prometheus/storage"
 	"github.com/prometheus/prometheus/tsdb"
 	"github.com/prometheus/prometheus/tsdb/chunkenc"
 	"github.com/prometheus/prometheus/tsdb/tsdbutil"
+	"github.com/prometheus/prometheus/util/stats"
 	"github.com/prometheus/prometheus/util/teststorage"
 	prom_testutil "github.com/prometheus/prometheus/util/testutil"
 )
@@ -2773,3 +2775,35 @@ func BenchmarkRuleDependencyController_AnalyseRules(b *testing.B) {
 		}
 	}
 }
+
+// TestEngineQueryFunc_ClosesQuery verifies EngineQueryFunc Close()s the
+// promql.Query it creates after every evaluation.
+func TestEngineQueryFunc_ClosesQuery(t *testing.T) {
+	eng := &closeCountingEngine{}
+	qf := EngineQueryFunc(eng, nil)
+
+	_, err := qf(context.Background(), "vector(0)", time.Unix(0, 0))
+	require.NoError(t, err)
+	require.Equal(t, 1, eng.q.closeCalls, "Close must be called exactly once")
+}
+
+type closeCountingEngine struct{ q closeCountingQuery }
+
+func (e *closeCountingEngine) NewInstantQuery(context.Context, storage.Queryable, promql.QueryOpts, string, time.Time) (promql.Query, error) {
+	return &e.q, nil
+}
+
+func (e *closeCountingEngine) NewRangeQuery(context.Context, storage.Queryable, promql.QueryOpts, string, time.Time, time.Time, time.Duration) (promql.Query, error) {
+	return &e.q, nil
+}
+
+type closeCountingQuery struct{ closeCalls int }
+
+func (*closeCountingQuery) Exec(context.Context) *promql.Result {
+	return &promql.Result{Value: promql.Vector{}}
+}
+func (q *closeCountingQuery) Close()                    { q.closeCalls++ }
+func (*closeCountingQuery) Statement() parser.Statement { return nil }
+func (*closeCountingQuery) Stats() *stats.Statistics    { return nil }
+func (*closeCountingQuery) Cancel()                     {}
+func (*closeCountingQuery) String() string              { return "" }