From f4e8687c88cf11ffef514acd9b3a4cc5b9348483 Mon Sep 17 00:00:00 2001
From: Aliaksandr Valialkin <valyala@gmail.com>
Date: Thu, 8 Oct 2020 20:12:57 +0300
Subject: [PATCH] app/vmalert: accept days, weeks and years in `for: ` part of
 config like Prometheus does

Updates https://github.com/VictoriaMetrics/VictoriaMetrics/issues/817
---
 app/vmalert/alerting.go                     |  2 +-
 app/vmalert/config/config.go                | 33 ++++++++++++++++++++-
 app/vmalert/config/config_test.go           |  2 +-
 app/vmalert/config/testdata/kube-good.rules |  4 +--
 app/vmalert/group_test.go                   |  4 +--
 5 files changed, 38 insertions(+), 7 deletions(-)

diff --git a/app/vmalert/alerting.go b/app/vmalert/alerting.go
index 08960f031..7e11ad850 100644
--- a/app/vmalert/alerting.go
+++ b/app/vmalert/alerting.go
@@ -53,7 +53,7 @@ func newAlertingRule(group *Group, cfg config.Rule) *AlertingRule {
 		RuleID:      cfg.ID,
 		Name:        cfg.Alert,
 		Expr:        cfg.Expr,
-		For:         cfg.For,
+		For:         cfg.For.Duration(),
 		Labels:      cfg.Labels,
 		Annotations: cfg.Annotations,
 		GroupID:     group.ID(),
diff --git a/app/vmalert/config/config.go b/app/vmalert/config/config.go
index cd40af7af..11f8f7eb8 100644
--- a/app/vmalert/config/config.go
+++ b/app/vmalert/config/config.go
@@ -94,7 +94,7 @@ type Rule struct {
 	Record      string            `yaml:"record,omitempty"`
 	Alert       string            `yaml:"alert,omitempty"`
 	Expr        string            `yaml:"expr"`
-	For         time.Duration     `yaml:"for,omitempty"`
+	For         PromDuration      `yaml:"for,omitempty"`
 	Labels      map[string]string `yaml:"labels,omitempty"`
 	Annotations map[string]string `yaml:"annotations,omitempty"`
 
@@ -102,6 +102,37 @@ type Rule struct {
 	XXX map[string]interface{} `yaml:",inline"`
 }
 
+// PromDuration is Prometheus duration.
+type PromDuration struct {
+	milliseconds int64
+}
+
+// NewPromDuration returns PromDuration for given d.
+func NewPromDuration(d time.Duration) PromDuration {
+	return PromDuration{
+		milliseconds: d.Milliseconds(),
+	}
+}
+
+// UnmarshalYAML implements yaml.Unmarshaler interface.
+func (pd *PromDuration) UnmarshalYAML(unmarshal func(interface{}) error) error {
+	var s string
+	if err := unmarshal(&s); err != nil {
+		return err
+	}
+	ms, err := metricsql.DurationValue(s, 0)
+	if err != nil {
+		return err
+	}
+	pd.milliseconds = ms
+	return nil
+}
+
+// Duration returns duration for pd.
+func (pd *PromDuration) Duration() time.Duration {
+	return time.Duration(pd.milliseconds) * time.Millisecond
+}
+
 // UnmarshalYAML implements the yaml.Unmarshaler interface.
 func (r *Rule) UnmarshalYAML(unmarshal func(interface{}) error) error {
 	type rule Rule
diff --git a/app/vmalert/config/config_test.go b/app/vmalert/config/config_test.go
index d01665980..a631e6ddc 100644
--- a/app/vmalert/config/config_test.go
+++ b/app/vmalert/config/config_test.go
@@ -270,7 +270,7 @@ func TestHashRule(t *testing.T) {
 			true,
 		},
 		{
-			Rule{Alert: "alert", Expr: "up == 1", For: time.Minute},
+			Rule{Alert: "alert", Expr: "up == 1", For: NewPromDuration(time.Minute)},
 			Rule{Alert: "alert", Expr: "up == 1"},
 			true,
 		},
diff --git a/app/vmalert/config/testdata/kube-good.rules b/app/vmalert/config/testdata/kube-good.rules
index a6e3e3e90..25ab513fc 100644
--- a/app/vmalert/config/testdata/kube-good.rules
+++ b/app/vmalert/config/testdata/kube-good.rules
@@ -665,7 +665,7 @@
           /
         sum(rate(kube_state_metrics_list_total{job="kube-state-metrics"}[5m])))
         > 0.01
-      for: 15m
+      for: 1d
       labels:
         severity: critical
     - alert: KubeStateMetricsWatchErrors
@@ -1724,4 +1724,4 @@
         rate(prometheus_operator_node_address_lookup_errors_total{job="prometheus-operator",namespace="monitoring"}[5m]) > 0.1
       for: 10m
       labels:
-        severity: warning
\ No newline at end of file
+        severity: warning
diff --git a/app/vmalert/group_test.go b/app/vmalert/group_test.go
index c18f8d487..4a3246a1c 100644
--- a/app/vmalert/group_test.go
+++ b/app/vmalert/group_test.go
@@ -32,7 +32,7 @@ func TestUpdateWith(t *testing.T) {
 			[]config.Rule{{
 				Alert: "foo",
 				Expr:  "up > 0",
-				For:   time.Second,
+				For:   config.NewPromDuration(time.Second),
 				Labels: map[string]string{
 					"bar": "baz",
 				},
@@ -44,7 +44,7 @@ func TestUpdateWith(t *testing.T) {
 			[]config.Rule{{
 				Alert: "foo",
 				Expr:  "up > 10",
-				For:   time.Second,
+				For:   config.NewPromDuration(time.Second),
 				Labels: map[string]string{
 					"baz": "bar",
 				},