From a7f00101f5603b8b5cb14c0cb0d03d48c07b5f9f Mon Sep 17 00:00:00 2001
From: Roman Khavronenko <hagen1778@gmail.com>
Date: Mon, 10 May 2021 09:11:45 +0100
Subject: [PATCH] vmalert: add support for `round_digits` param in datasource
 package (#1278)

Starting from v1.56.0 VM supports `round_digits` which allows to limit
the number of digits after the decimal point in response value. The feature
can be used to reduce entropy of produced by recording rules values
and significantly improve the compression. See more details in link below.

https://github.com/VictoriaMetrics/VictoriaMetrics/issues/525
---
 app/vmalert/datasource/init.go    | 23 +++++++++++++++++++++--
 app/vmalert/datasource/vm.go      |  4 ++++
 app/vmalert/datasource/vm_test.go | 10 ++++++++++
 3 files changed, 35 insertions(+), 2 deletions(-)

diff --git a/app/vmalert/datasource/init.go b/app/vmalert/datasource/init.go
index 62603f749b..093dedabb0 100644
--- a/app/vmalert/datasource/init.go
+++ b/app/vmalert/datasource/init.go
@@ -4,6 +4,7 @@ import (
 	"flag"
 	"fmt"
 	"net/http"
+	"strings"
 
 	"github.com/VictoriaMetrics/VictoriaMetrics/app/vmalert/utils"
 )
@@ -26,6 +27,8 @@ var (
 		"For example, if datasource.queryStep=15s then param \"step\" with value \"15s\" will be added to every query."+
 		"If queryStep isn't specified, rule's evaluationInterval will be used instead.")
 	maxIdleConnections = flag.Int("datasource.maxIdleConnections", 100, `Defines the number of idle (keep-alive connections) to each configured datasource. Consider setting this value equal to the value: groups_total * group.concurrency. Too low a value may result in a high number of sockets in TIME_WAIT state.`)
+	roundDigits        = flag.Int("datasource.roundDigits", 0, `Adds \"round_digits\" GET param to datasource requests. 
+		In VM \"round_digits\" limits the number of digits after the decimal point in response values.`)
 )
 
 // Init creates a Querier from provided flag values.
@@ -33,11 +36,27 @@ func Init() (QuerierBuilder, error) {
 	if *addr == "" {
 		return nil, fmt.Errorf("datasource.url is empty")
 	}
+
 	tr, err := utils.Transport(*addr, *tlsCertFile, *tlsKeyFile, *tlsCAFile, *tlsServerName, *tlsInsecureSkipVerify)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create transport: %w", err)
 	}
 	tr.MaxIdleConns = *maxIdleConnections
-	c := &http.Client{Transport: tr}
-	return NewVMStorage(*addr, *basicAuthUsername, *basicAuthPassword, *lookBack, *queryStep, *appendTypePrefix, c), nil
+
+	var rd string
+	if *roundDigits > 0 {
+		rd = fmt.Sprintf("%d", *roundDigits)
+	}
+
+	return &VMStorage{
+		c:                &http.Client{Transport: tr},
+		basicAuthUser:    *basicAuthUsername,
+		basicAuthPass:    *basicAuthPassword,
+		datasourceURL:    strings.TrimSuffix(*addr, "/"),
+		appendTypePrefix: *appendTypePrefix,
+		lookBack:         *lookBack,
+		queryStep:        *queryStep,
+		roundDigits:      rd,
+		dataSourceType:   NewPrometheusType(),
+	}, nil
 }
diff --git a/app/vmalert/datasource/vm.go b/app/vmalert/datasource/vm.go
index 72ab551758..8b59b389c3 100644
--- a/app/vmalert/datasource/vm.go
+++ b/app/vmalert/datasource/vm.go
@@ -81,6 +81,7 @@ type VMStorage struct {
 	appendTypePrefix bool
 	lookBack         time.Duration
 	queryStep        time.Duration
+	roundDigits      string
 
 	dataSourceType     Type
 	evaluationInterval time.Duration
@@ -218,6 +219,9 @@ func (s *VMStorage) setPrometheusReqParams(r *http.Request, query string, timest
 		// override step with user-specified value
 		q.Set("step", s.queryStep.String())
 	}
+	if s.roundDigits != "" {
+		q.Set("round_digits", s.roundDigits)
+	}
 	r.URL.RawQuery = q.Encode()
 }
 
diff --git a/app/vmalert/datasource/vm_test.go b/app/vmalert/datasource/vm_test.go
index db030cd5ac..9a004bcc56 100644
--- a/app/vmalert/datasource/vm_test.go
+++ b/app/vmalert/datasource/vm_test.go
@@ -243,6 +243,16 @@ func TestPrepareReq(t *testing.T) {
 				checkEqualString(t, exp, r.URL.RawQuery)
 			},
 		},
+		{
+			"round digits",
+			&VMStorage{
+				roundDigits: "10",
+			},
+			func(t *testing.T, r *http.Request) {
+				exp := fmt.Sprintf("query=%s&round_digits=10&time=%d", query, timestamp.Unix())
+				checkEqualString(t, exp, r.URL.RawQuery)
+			},
+		},
 	}
 
 	for _, tc := range testCases {