From 5b235b902b7efc0ad24f29a3d5d2c551cb1b6d6f Mon Sep 17 00:00:00 2001
From: Nikolay <nik@victoriametrics.com>
Date: Wed, 4 Nov 2020 21:29:18 +0300
Subject: [PATCH] Adds ready probe (#874)

* adds leading forward slash check for scrapeURL path
https://github.com/VictoriaMetrics/VictoriaMetrics/issues/835

* adds ready probe for scrape config initialization,
it should prevent metrics loss during vmagent rolling update,
/ready api will return 425 http code, if some scrape config still waits for initialization.

* updates docs

* Update app/vmagent/README.md

* renames var

* Update app/vmagent/README.md

Co-authored-by: Aliaksandr Valialkin <valyala@gmail.com>
---
 app/vmagent/README.md     |  2 ++
 app/vmagent/main.go       | 11 +++++++++++
 docs/vmagent.md           |  2 ++
 lib/promscrape/scraper.go | 13 ++++++++++++-
 4 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/app/vmagent/README.md b/app/vmagent/README.md
index 6d132ad5a..099947c5b 100644
--- a/app/vmagent/README.md
+++ b/app/vmagent/README.md
@@ -219,6 +219,8 @@ It accepts optional `show_original_labels=1` query arg, which shows the original
 This information may be useful for debugging target relabeling.
 * `http://vmagent-host:8429/api/v1/targets`. This handler returns data compatible with [the corresponding page from Prometheus API](https://prometheus.io/docs/prometheus/latest/querying/api/#targets).
 
+* `http://vmagent-host:8429/ready` - this handler returns http 200 status code when `vmagent` finishes initialization for all service_discovery configs.
+  It may be useful, when you have many entries at `-promscrape.config` and want to perform `vmagent` rolling update without scrape loss.
 
 ### Troubleshooting
 
diff --git a/app/vmagent/main.go b/app/vmagent/main.go
index d3fa1d6b9..4d2b7acb5 100644
--- a/app/vmagent/main.go
+++ b/app/vmagent/main.go
@@ -7,6 +7,7 @@ import (
 	"os"
 	"strconv"
 	"strings"
+	"sync/atomic"
 	"time"
 
 	"github.com/VictoriaMetrics/VictoriaMetrics/app/vmagent/csvimport"
@@ -222,6 +223,16 @@ func requestHandler(w http.ResponseWriter, r *http.Request) bool {
 		procutil.SelfSIGHUP()
 		w.WriteHeader(http.StatusOK)
 		return true
+	case "/ready":
+		if rdy := atomic.LoadInt32(&promscrape.PendingScrapeConfigs); rdy > 0 {
+			errMsg := fmt.Sprintf("waiting for scrapes to init, left: %d", rdy)
+			http.Error(w, errMsg, http.StatusTooEarly)
+		} else {
+			w.Header().Set("Content-Type", "text/plain")
+			w.WriteHeader(http.StatusOK)
+			w.Write([]byte("OK"))
+		}
+		return true
 	}
 	return false
 }
diff --git a/docs/vmagent.md b/docs/vmagent.md
index 6d132ad5a..87559d888 100644
--- a/docs/vmagent.md
+++ b/docs/vmagent.md
@@ -219,6 +219,8 @@ It accepts optional `show_original_labels=1` query arg, which shows the original
 This information may be useful for debugging target relabeling.
 * `http://vmagent-host:8429/api/v1/targets`. This handler returns data compatible with [the corresponding page from Prometheus API](https://prometheus.io/docs/prometheus/latest/querying/api/#targets).
 
+* `http://vmagent-host:8429/ready` - this handler returns http 200 status code, when `vmagent` finished initialization for all service_discovery configs. 
+  It may be useful, when you have a lof of entries at promscrape.config and want to perform `vmagent` rolling update without metrics loss.
 
 ### Troubleshooting
 
diff --git a/lib/promscrape/scraper.go b/lib/promscrape/scraper.go
index 341e66462..52f99a97e 100644
--- a/lib/promscrape/scraper.go
+++ b/lib/promscrape/scraper.go
@@ -5,6 +5,7 @@ import (
 	"flag"
 	"fmt"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/VictoriaMetrics/VictoriaMetrics/lib/logger"
@@ -75,6 +76,9 @@ func Stop() {
 var (
 	globalStopCh chan struct{}
 	scraperWG    sync.WaitGroup
+	// PendingScrapeConfigs - zero value means, that
+	// all scrapeConfigs are inited and ready for work.
+	PendingScrapeConfigs int32
 )
 
 func runScraper(configFile string, pushData func(wr *prompbmarshal.WriteRequest), globalStopCh <-chan struct{}) {
@@ -166,6 +170,7 @@ func newScrapeConfigs(pushData func(wr *prompbmarshal.WriteRequest)) *scrapeConf
 }
 
 func (scs *scrapeConfigs) add(name string, checkInterval time.Duration, getScrapeWork func(cfg *Config, swsPrev []ScrapeWork) []ScrapeWork) {
+	atomic.AddInt32(&PendingScrapeConfigs, 1)
 	scfg := &scrapeConfig{
 		name:          name,
 		pushData:      scs.pushData,
@@ -216,10 +221,15 @@ func (scfg *scrapeConfig) run() {
 
 	cfg := <-scfg.cfgCh
 	var swsPrev []ScrapeWork
-	for {
+	updateScrapeWork := func(cfg *Config) {
 		sws := scfg.getScrapeWork(cfg, swsPrev)
 		sg.update(sws)
 		swsPrev = sws
+	}
+	updateScrapeWork(cfg)
+	atomic.AddInt32(&PendingScrapeConfigs, -1)
+
+	for {
 
 		select {
 		case <-scfg.stopCh:
@@ -227,6 +237,7 @@ func (scfg *scrapeConfig) run() {
 		case cfg = <-scfg.cfgCh:
 		case <-tickerCh:
 		}
+		updateScrapeWork(cfg)
 	}
 }