From 33cd6c26d3bc7155499068a94618afe72f7d004e Mon Sep 17 00:00:00 2001
From: Aliaksandr Valialkin <valyala@gmail.com>
Date: Sun, 14 Mar 2021 22:24:03 +0200
Subject: [PATCH] lib/promscrape: return back the logic for flushing big
 buffers to storage from the commit 3fd8653b40113de40b6dc265724104fcf2f776d9

This should reduce memory usage when vmagent scrapes targets with big number of metrics and `-promscrape.streamParse` isn't enabled
---
 lib/promscrape/scrapework.go | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

diff --git a/lib/promscrape/scrapework.go b/lib/promscrape/scrapework.go
index d19ae8b349..e6b537b680 100644
--- a/lib/promscrape/scrapework.go
+++ b/lib/promscrape/scrapework.go
@@ -294,10 +294,22 @@ func (sw *scrapeWork) scrapeInternal(scrapeTimestamp, realTimestamp int64) error
 	srcRows := wc.rows.Rows
 	samplesScraped := len(srcRows)
 	scrapedSamples.Update(float64(samplesScraped))
+	samplesPostRelabeling := 0
 	for i := range srcRows {
 		sw.addRowToTimeseries(wc, &srcRows[i], scrapeTimestamp, true)
+		if len(wc.labels) > 40000 {
+			// Limit the maximum size of wc.writeRequest.
+			// This should reduce memory usage when scraping targets with millions of metrics and/or labels.
+			// For example, when scraping /federate handler from Prometheus - see https://prometheus.io/docs/prometheus/latest/federation/
+			samplesPostRelabeling += len(wc.writeRequest.Timeseries)
+			sw.updateSeriesAdded(wc)
+			startTime := time.Now()
+			sw.PushData(&wc.writeRequest)
+			pushDataDuration.UpdateDuration(startTime)
+			wc.resetNoRows()
+		}
 	}
-	samplesPostRelabeling := len(wc.writeRequest.Timeseries)
+	samplesPostRelabeling += len(wc.writeRequest.Timeseries)
 	if sw.Config.SampleLimit > 0 && samplesPostRelabeling > sw.Config.SampleLimit {
 		wc.resetNoRows()
 		up = 0