Merge pull request #328 from grafana/dimitar/rules-manager-always-restore

Add option to always restore the state of rules when loading

rules/fixtures/rules_alerts.yaml

@@ -0,0 +1,5 @@
+groups:
+  - name: test
+    rules:
+      - alert: test
+        expr: sum by (job)(rate(http_requests_total[5m]))

rules/fixtures/rules_alerts2.yaml

@@ -0,0 +1,5 @@
+groups:
+  - name: test
+    rules:
+      - alert: test_2
+        expr: sum by (job)(rate(http_requests_total[5m]))

rules/manager.go

@@ -962,6 +962,10 @@ type ManagerOptions struct {
 	GroupLoader                GroupLoader
 	DefaultEvaluationDelay     func() time.Duration
 
+	// AlwaysRestoreAlertState forces all new or changed groups in calls to Update to restore.
+	// Useful when you know you will be adding alerting rules after the manager has already started.
+	AlwaysRestoreAlertState bool
+
 	Metrics *Metrics
 }
 
@@ -1116,7 +1120,7 @@ func (m *Manager) LoadGroups(
 ) (map[string]*Group, []error) {
 	groups := make(map[string]*Group)
 
-	shouldRestore := !m.restored
+	shouldRestore := !m.restored || m.opts.AlwaysRestoreAlertState
 
 	for _, fn := range filenames {
 		rgs, errs := m.opts.GroupLoader.Load(fn)
@@ -1146,7 +1150,7 @@ func (m *Manager) LoadGroups(
 						labels.FromMap(r.Annotations),
 						externalLabels,
 						externalURL,
-						m.restored,
+						!shouldRestore,
 						log.With(m.logger, "alert", r.Alert),
 					))
 					continue

rules/manager_test.go

@@ -809,6 +809,69 @@ func TestUpdate(t *testing.T) {
 	reloadAndValidate(rgs, t, tmpFile, ruleManager, expected, ogs)
 }
 
+func TestUpdate_AlwaysRestore(t *testing.T) {
+	st := teststorage.New(t)
+	defer st.Close()
+
+	ruleManager := NewManager(&ManagerOptions{
+		Appendable:              st,
+		Queryable:               st,
+		Context:                 context.Background(),
+		Logger:                  log.NewNopLogger(),
+		AlwaysRestoreAlertState: true,
+	})
+	ruleManager.start()
+	defer ruleManager.Stop()
+
+	err := ruleManager.Update(10*time.Second, []string{"fixtures/rules_alerts.yaml"}, nil, "", nil)
+	require.NoError(t, err)
+
+	for _, g := range ruleManager.groups {
+		require.True(t, g.shouldRestore)
+		g.shouldRestore = false // set to false to check if Update will set it to true again
+	}
+
+	// Use different file, so groups haven't changed, therefore, we expect state restoration
+	err = ruleManager.Update(10*time.Second, []string{"fixtures/rules_alerts2.yaml"}, nil, "", nil)
+	for _, g := range ruleManager.groups {
+		require.True(t, g.shouldRestore)
+	}
+
+	require.NoError(t, err)
+}
+
+func TestUpdate_AlwaysRestoreDoesntAffectUnchangedGroups(t *testing.T) {
+	files := []string{"fixtures/rules_alerts.yaml"}
+	st := teststorage.New(t)
+	defer st.Close()
+
+	ruleManager := NewManager(&ManagerOptions{
+		Appendable:              st,
+		Queryable:               st,
+		Context:                 context.Background(),
+		Logger:                  log.NewNopLogger(),
+		AlwaysRestoreAlertState: true,
+	})
+	ruleManager.start()
+	defer ruleManager.Stop()
+
+	err := ruleManager.Update(10*time.Second, files, nil, "", nil)
+	require.NoError(t, err)
+
+	for _, g := range ruleManager.groups {
+		require.True(t, g.shouldRestore)
+		g.shouldRestore = false // set to false to check if Update will set it to true again
+	}
+
+	// Use the same file, so groups haven't changed, therefore, we don't expect state restoration
+	err = ruleManager.Update(10*time.Second, files, nil, "", nil)
+	for _, g := range ruleManager.groups {
+		require.False(t, g.shouldRestore)
+	}
+
+	require.NoError(t, err)
+}
+
 func TestUpdateSetsSourceTenants(t *testing.T) {
 	st := teststorage.New(t)
 	defer st.Close()