From a2f412e21ff026da7f7f974c9f2916c5acd2b9e3 Mon Sep 17 00:00:00 2001
From: George Robinson <george.robinson@grafana.com>
Date: Wed, 9 Nov 2022 17:15:09 +0000
Subject: [PATCH] Alerting: Small improvements to staleResultsHandler (#58007)
 (#58513)

(cherry picked from commit 1290951b65b75ac6a054b401d3e91388cddfe33d)
---
 pkg/services/ngalert/state/manager.go    | 39 +++++++++++++-----------
 pkg/services/ngalert/state/state.go      |  7 +++++
 pkg/services/ngalert/state/state_test.go |  7 +++++
 3 files changed, 36 insertions(+), 17 deletions(-)

diff --git a/pkg/services/ngalert/state/manager.go b/pkg/services/ngalert/state/manager.go
index dbec25f7fb1..9ef60f3396b 100644
--- a/pkg/services/ngalert/state/manager.go
+++ b/pkg/services/ngalert/state/manager.go
@@ -180,7 +180,7 @@ func (st *Manager) ProcessEvalResults(ctx context.Context, evaluatedAt time.Time
 		states = append(states, s)
 		processedResults[s.CacheID] = s
 	}
-	resolvedStates := st.staleResultsHandler(ctx, evaluatedAt, alertRule, processedResults)
+	resolvedStates := st.staleResultsHandler(ctx, alertRule, processedResults, evaluatedAt)
 	if len(states) > 0 {
 		logger.Debug("saving new states to the database", "count", len(states))
 		for _, state := range states {
@@ -374,16 +374,24 @@ func (st *Manager) annotateState(ctx context.Context, alertRule *ngModels.AlertR
 	}
 }
 
-func (st *Manager) staleResultsHandler(ctx context.Context, evaluatedAt time.Time, alertRule *ngModels.AlertRule, states map[string]*State) []*State {
-	// If we are removing two or more stale series it makes sense to share the resolved image as the alert rule is the same.
-	// TODO: We will need to change this when we support images without screenshots as each series will have a different image
-	var resolvedImage *ngModels.Image
+func (st *Manager) staleResultsHandler(ctx context.Context, r *ngModels.AlertRule, states map[string]*State, evaluatedAt time.Time) []*State {
+	var (
+		// resolvedImage contains the image for all stale states that are resolved. The resolved image is shared between
+		// all resolved states as the alert rule is the same. TODO: We will need to change this when we support images
+		// without screenshots as each state will have a different image
+		resolvedImage *ngModels.Image
 
-	var resolvedStates []*State
-	allStates := st.GetStatesForRuleUID(alertRule.OrgID, alertRule.UID)
-	for _, s := range allStates {
+		// resolvedStates contains the stale states that were resolved
+		resolvedStates []*State
+
+		// knownStates contains the current set of states in the state cache
+		knownStates []*State
+	)
+
+	knownStates = st.GetStatesForRuleUID(r.OrgID, r.UID)
+	for _, s := range knownStates {
 		_, ok := states[s.CacheID]
-		if !ok && isItStale(evaluatedAt, s.LastEvaluationTime, alertRule.IntervalSeconds) {
+		if !ok && isItStale(evaluatedAt, s.LastEvaluationTime, r.IntervalSeconds) {
 			st.log.Debug("removing stale state entry", "orgID", s.OrgID, "alertRuleUID", s.AlertRuleUID, "cacheID", s.CacheID)
 			st.cache.deleteEntry(s.OrgID, s.AlertRuleUID, s.CacheID)
 			ilbs := ngModels.InstanceLabels(s.Labels)
@@ -398,22 +406,19 @@ func (st *Manager) staleResultsHandler(ctx context.Context, evaluatedAt time.Tim
 
 			if s.State == eval.Alerting {
 				previousState := InstanceStateAndReason{State: s.State, Reason: s.StateReason}
-				s.State = eval.Normal
-				s.StateReason = ngModels.StateReasonMissingSeries
-				s.EndsAt = evaluatedAt
-				s.Resolved = true
-				st.annotateState(ctx, alertRule, s.Labels, evaluatedAt,
+				s.Resolve(ngModels.StateReasonMissingSeries, evaluatedAt)
+				st.annotateState(ctx, r, s.Labels, evaluatedAt,
 					InstanceStateAndReason{State: eval.Normal, Reason: s.StateReason},
 					previousState,
 				)
 
 				// If there is no resolved image for this rule then take one
 				if resolvedImage == nil {
-					image, err := takeImage(ctx, st.imageService, alertRule)
+					image, err := takeImage(ctx, st.imageService, r)
 					if err != nil {
 						st.log.Warn("Failed to take an image",
-							"dashboard", alertRule.DashboardUID,
-							"panel", alertRule.PanelID,
+							"dashboard", r.DashboardUID,
+							"panel", r.PanelID,
 							"error", err)
 					} else if image != nil {
 						resolvedImage = image
diff --git a/pkg/services/ngalert/state/state.go b/pkg/services/ngalert/state/state.go
index 4aae855b89f..9447024d84c 100644
--- a/pkg/services/ngalert/state/state.go
+++ b/pkg/services/ngalert/state/state.go
@@ -73,6 +73,13 @@ func (a *State) GetRuleKey() models.AlertRuleKey {
 	}
 }
 
+func (a *State) Resolve(reason string, endsAt time.Time) {
+	a.State = eval.Normal
+	a.StateReason = reason
+	a.EndsAt = endsAt
+	a.Resolved = true
+}
+
 type Evaluation struct {
 	EvaluationTime  time.Time
 	EvaluationState eval.State
diff --git a/pkg/services/ngalert/state/state_test.go b/pkg/services/ngalert/state/state_test.go
index f1bd9552776..b95f1bc2520 100644
--- a/pkg/services/ngalert/state/state_test.go
+++ b/pkg/services/ngalert/state/state_test.go
@@ -298,6 +298,13 @@ func TestGetLastEvaluationValuesForCondition(t *testing.T) {
 	})
 }
 
+func TestResolve(t *testing.T) {
+	s := State{State: eval.Alerting, EndsAt: time.Now().Add(time.Minute)}
+	expected := State{State: eval.Normal, StateReason: "This is a reason", EndsAt: time.Now(), Resolved: true}
+	s.Resolve("This is a reason", expected.EndsAt)
+	assert.Equal(t, expected, s)
+}
+
 func TestShouldTakeImage(t *testing.T) {
 	tests := []struct {
 		name          string