From 7e32daa63ab9e87d56732fc4c0e1e08ba2e048b9 Mon Sep 17 00:00:00 2001
From: Anzor <anzoor1998@gmail.com>
Date: Wed, 7 Aug 2024 10:36:14 +0300
Subject: [PATCH] app/vmagent: read __sample_limit__ from labels (#6665)
 (#6666)

By introducing this feature, users will have the ability to customize
the sampleLimit parameter on a per-target basis, providing more
flexibility and control over the job execution behavior.

(cherry picked from commit 994796367bd8469100f8a8ec51d67be3711baa6b)
---
 docs/CHANGELOG.md             |  1 +
 lib/promscrape/config.go      | 12 +++++++++++-
 lib/promscrape/config_test.go |  4 +++-
 3 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md
index c0af992d6..67d3bda31 100644
--- a/docs/CHANGELOG.md
+++ b/docs/CHANGELOG.md
@@ -32,6 +32,7 @@ See also [LTS releases](https://docs.victoriametrics.com/lts-releases/).
 
 * FEATURE: add `/influx/health` health-check handler for Influx endpoints. This is needed as some clients use the health endpoint to determine if the server is healthy and ready for data ingestion. See [this issue](https://github.com/VictoriaMetrics/VictoriaMetrics/issues/6653) for the details.
 * FEATURE: [vmctl](https://docs.victoriametrics.com/vmctl/): add `--vm-backoff-retries`, `--vm-backoff-factor`, `--vm-backoff-min-duration` and `--vm-native-backoff-retries`, `--vm-native-backoff-factor`, `--vm-native-backoff-min-duration` command-line flags. These flags allow to change backoff policy config for import requests to VictoriaMetrics. See [this issue](https://github.com/VictoriaMetrics/VictoriaMetrics/issues/6622).
+* FEATURE: [vmagent](https://docs.victoriametrics.com/vmagent/): allow overriding the `sample_limit` option at [scrape_configs](https://docs.victoriametrics.com/sd_configs/#scrape_configs) when a label `__sample_limit__` is specified for target. See [this issue](https://github.com/VictoriaMetrics/VictoriaMetrics/issues/6665). Thanks to @zoglam for the [pull request](https://github.com/VictoriaMetrics/VictoriaMetrics/pull/6666).
 
 * BUGFIX: [vmalert](https://docs.victoriametrics.com/vmalert): respect HTTP headers defined in [notifier configuration file](https://docs.victoriametrics.com/vmalert/#notifier-configuration-file) for each request to notifiers. Previously, this param was ignored by mistake.
 
diff --git a/lib/promscrape/config.go b/lib/promscrape/config.go
index 7c94f6ae5..0b4af38d9 100644
--- a/lib/promscrape/config.go
+++ b/lib/promscrape/config.go
@@ -1187,6 +1187,16 @@ func (swc *scrapeWorkConfig) getScrapeWork(target string, extraLabels, metaLabel
 		}
 		seriesLimit = n
 	}
+	// Read sample_limit option from __sample_limit__ label.
+	// See https://docs.victoriametrics.com/vmagent/#automatically-generated-metrics
+	sampleLimit := swc.sampleLimit
+	if s := labels.Get("__sample_limit__"); len(s) > 0 {
+		n, err := strconv.Atoi(s)
+		if err != nil {
+			return nil, fmt.Errorf("cannot parse __sample_limit__=%q: %w", s, err)
+		}
+		sampleLimit = n
+	}
 	// Read stream_parse option from __stream_parse__ label.
 	// See https://docs.victoriametrics.com/vmagent/#stream-parsing-mode
 	streamParse := swc.streamParse
@@ -1232,7 +1242,7 @@ func (swc *scrapeWorkConfig) getScrapeWork(target string, extraLabels, metaLabel
 		AuthConfig:           swc.authConfig,
 		RelabelConfigs:       swc.relabelConfigs,
 		MetricRelabelConfigs: swc.metricRelabelConfigs,
-		SampleLimit:          swc.sampleLimit,
+		SampleLimit:          sampleLimit,
 		DisableCompression:   swc.disableCompression,
 		DisableKeepAlive:     swc.disableKeepAlive,
 		StreamParse:          streamParse,
diff --git a/lib/promscrape/config_test.go b/lib/promscrape/config_test.go
index 369053e18..29930d380 100644
--- a/lib/promscrape/config_test.go
+++ b/lib/promscrape/config_test.go
@@ -1150,6 +1150,8 @@ scrape_configs:
         replacement: 127.0.0.1:9116  # The SNMP exporter's real hostname:port.
       - target_label: __series_limit__
         replacement: 1234
+      - target_label: __sample_limit__
+        replacement: 5678
       - target_label: __stream_parse__
         replacement: true
 `, []*ScrapeWork{
@@ -1162,7 +1164,7 @@ scrape_configs:
 				"instance": "192.168.1.2",
 				"job":      "snmp",
 			}),
-			SampleLimit:         100,
+			SampleLimit:         5678,
 			DisableKeepAlive:    true,
 			DisableCompression:  true,
 			StreamParse:         true,