From e26bcb8bbb9a116220fcd2e9619763e463594a32 Mon Sep 17 00:00:00 2001
From: Nikolay <nik@victoriametrics.com>
Date: Tue, 12 Apr 2022 12:24:11 +0300
Subject: [PATCH] lib/promscrape: allows to use k8s pod name as
 clusterMemberNum (#2436)

* lib/promscrape: allows to use k8s pod name as clusterMemberNum
it must improve user expirience and simplify clustering scrapers.
it must allow to use vmagent cluster with distroless images
https://github.com/VictoriaMetrics/VictoriaMetrics/issues/2359

* Apply suggestions from code review

Co-authored-by: Aliaksandr Valialkin <valyala@victoriametrics.com>
---
 lib/promscrape/config.go  | 26 +++++++++++++++++++++++---
 lib/promscrape/scraper.go |  3 +++
 2 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/lib/promscrape/config.go b/lib/promscrape/config.go
index d30b178e69..52c32858f7 100644
--- a/lib/promscrape/config.go
+++ b/lib/promscrape/config.go
@@ -49,12 +49,32 @@ var (
 	clusterMembersCount = flag.Int("promscrape.cluster.membersCount", 0, "The number of members in a cluster of scrapers. "+
 		"Each member must have an unique -promscrape.cluster.memberNum in the range 0 ... promscrape.cluster.membersCount-1 . "+
 		"Each member then scrapes roughly 1/N of all the targets. By default cluster scraping is disabled, i.e. a single scraper scrapes all the targets")
-	clusterMemberNum = flag.Int("promscrape.cluster.memberNum", 0, "The number of number in the cluster of scrapers. "+
-		"It must be an unique value in the range 0 ... promscrape.cluster.membersCount-1 across scrapers in the cluster")
+	clusterMemberNum = flag.String("promscrape.cluster.memberNum", "0", "The number of number in the cluster of scrapers. "+
+		"It must be an unique value in the range 0 ... promscrape.cluster.membersCount-1 across scrapers in the cluster. "+
+		"Can be specified as pod name of Kubernetes StatefulSet - pod-name-Num, where Num is a numeric part of pod name")
 	clusterReplicationFactor = flag.Int("promscrape.cluster.replicationFactor", 1, "The number of members in the cluster, which scrape the same targets. "+
 		"If the replication factor is greater than 2, then the deduplication must be enabled at remote storage side. See https://docs.victoriametrics.com/#deduplication")
 )
 
+var clusterMemberID int
+
+// must be called before any scraper
+func initClusterMemberID() error {
+	s := *clusterMemberNum
+	// special case for kubernetes deployment, where pod-name formatted at some-pod-name-1
+	// obtain memberNum from last segment
+	// https://github.com/VictoriaMetrics/VictoriaMetrics/issues/2359
+	if idx := strings.LastIndexByte(s, '-'); idx >= 0 {
+		s = s[idx+1:]
+	}
+	n, err := strconv.ParseInt(s, 10, 64)
+	if err != nil {
+		return fmt.Errorf("cannot parse -promscrape.cluster.memberNum=%q: %w", *clusterMemberNum, err)
+	}
+	clusterMemberID = int(n)
+	return nil
+}
+
 // Config represents essential parts from Prometheus config defined at https://prometheus.io/docs/prometheus/latest/configuration/configuration/
 type Config struct {
 	Global            GlobalConfig   `yaml:"global,omitempty"`
@@ -996,7 +1016,7 @@ func (swc *scrapeWorkConfig) getScrapeWork(target string, extraLabels, metaLabel
 	if *clusterMembersCount > 1 {
 		bb := scrapeWorkKeyBufPool.Get()
 		bb.B = appendScrapeWorkKey(bb.B[:0], labels)
-		needSkip := needSkipScrapeWork(bytesutil.ToUnsafeString(bb.B), *clusterMembersCount, *clusterReplicationFactor, *clusterMemberNum)
+		needSkip := needSkipScrapeWork(bytesutil.ToUnsafeString(bb.B), *clusterMembersCount, *clusterReplicationFactor, clusterMemberID)
 		scrapeWorkKeyBufPool.Put(bb)
 		if needSkip {
 			return nil, nil
diff --git a/lib/promscrape/scraper.go b/lib/promscrape/scraper.go
index c5a3e8f12d..c6415d416d 100644
--- a/lib/promscrape/scraper.go
+++ b/lib/promscrape/scraper.go
@@ -99,6 +99,9 @@ func runScraper(configFile string, pushData func(wr *prompbmarshal.WriteRequest)
 	// See https://github.com/VictoriaMetrics/VictoriaMetrics/issues/1240
 	sighupCh := procutil.NewSighupChan()
 
+	if err := initClusterMemberID(); err != nil {
+		logger.Fatalf("cannot init clusterMembership: %s", err)
+	}
 	logger.Infof("reading Prometheus configs from %q", configFile)
 	cfg, data, err := loadConfig(configFile)
 	if err != nil {