From 7785869ccc4d4ff4dc80ad9f33442f6c5b369c72 Mon Sep 17 00:00:00 2001
From: Aliaksandr Valialkin <valyala@gmail.com>
Date: Tue, 1 Sep 2020 10:55:21 +0300
Subject: [PATCH] lib/promscrape: reduce memory usage when scraping targets
 with millions of metrics

This should help when scraping /federate endpoints from Prometheus instances,
which scrape millions of metrics. See https://prometheus.io/docs/prometheus/latest/federation/
---
 lib/promscrape/scrapework.go | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/lib/promscrape/scrapework.go b/lib/promscrape/scrapework.go
index c918450bc..9746d871d 100644
--- a/lib/promscrape/scrapework.go
+++ b/lib/promscrape/scrapework.go
@@ -230,6 +230,15 @@ func (sw *scrapeWork) scrapeInternal(scrapeTimestamp, realTimestamp int64) error
 	scrapedSamples.Update(float64(samplesScraped))
 	for i := range srcRows {
 		sw.addRowToTimeseries(wc, &srcRows[i], scrapeTimestamp, true)
+		if len(wc.labels) > 10000 {
+			// 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/
+			startTime := time.Now()
+			sw.PushData(&wc.writeRequest)
+			pushDataDuration.UpdateDuration(startTime)
+			wc.reset()
+		}
 	}
 	if sw.Config.SampleLimit > 0 && len(wc.writeRequest.Timeseries) > sw.Config.SampleLimit {
 		prompbmarshal.ResetWriteRequest(&wc.writeRequest)
@@ -256,7 +265,7 @@ func (sw *scrapeWork) scrapeInternal(scrapeTimestamp, realTimestamp int64) error
 	return err
 }
 
-// leveldWriteRequestCtxPool allows reducing memory usage when writeRequesCtx
+// leveledWriteRequestCtxPool allows reducing memory usage when writeRequesCtx
 // structs contain mixed number of labels.
 //
 // Its logic has been copied from leveledbytebufferpool.