Merge pull request #17390 from prymitive/sample_limit

Include histograms in sample_limit logic

scrape/scrape_test.go

@@ -2195,7 +2195,6 @@ func TestScrapeLoop_HistogramBucketLimit(t *testing.T) {
 		}
 		return l
 	}
-	sl.sampleLimit = app.limit
 
 	metric := dto.Metric{}
 	err := sl.metrics.targetScrapeNativeHistogramBucketLimit.Write(&metric)

scrape/target.go

@@ -347,6 +347,22 @@ func (app *limitAppender) Append(ref storage.SeriesRef, lset labels.Labels, t in
 	return ref, nil
 }
 
+func (app *limitAppender) AppendHistogram(ref storage.SeriesRef, lset labels.Labels, t int64, h *histogram.Histogram, fh *histogram.FloatHistogram) (storage.SeriesRef, error) {
+	// Bypass sample_limit checks only if we have a staleness marker for a known series (ref value is non-zero).
+	// This ensures that if a series is already in TSDB then we always write the marker.
+	if ref == 0 || (h != nil && !value.IsStaleNaN(h.Sum)) || (fh != nil && !value.IsStaleNaN(fh.Sum)) {
+		app.i++
+		if app.i > app.limit {
+			return 0, errSampleLimit
+		}
+	}
+	ref, err := app.Appender.AppendHistogram(ref, lset, t, h, fh)
+	if err != nil {
+		return 0, err
+	}
+	return ref, nil
+}
+
 type timeLimitAppender struct {
 	storage.Appender
 

scrape/target_test.go

@@ -14,6 +14,7 @@
 package scrape
 
 import (
+	"context"
 	"crypto/tls"
 	"crypto/x509"
 	"fmt"
@@ -34,6 +35,8 @@ import (
 	"github.com/prometheus/prometheus/discovery/targetgroup"
 	"github.com/prometheus/prometheus/model/histogram"
 	"github.com/prometheus/prometheus/model/labels"
+	"github.com/prometheus/prometheus/model/timestamp"
+	"github.com/prometheus/prometheus/storage"
 )
 
 const (
@@ -719,3 +722,40 @@ func TestMaxSchemaAppender(t *testing.T) {
 		}
 	}
 }
+
+// Test sample_limit when a scrape containst Native Histograms.
+func TestAppendWithSampleLimitAndNativeHistogram(t *testing.T) {
+	const sampleLimit = 2
+	resApp := &collectResultAppender{}
+	sl := newBasicScrapeLoop(t, context.Background(), nil, func(_ context.Context) storage.Appender {
+		return resApp
+	}, 0)
+	sl.sampleLimit = sampleLimit
+
+	now := time.Now()
+	app := appender(sl.appender(context.Background()), sl.sampleLimit, sl.bucketLimit, sl.maxSchema)
+
+	// sample_limit is set to 2, so first two scrapes should work
+	_, err := app.Append(0, labels.FromStrings(model.MetricNameLabel, "foo"), timestamp.FromTime(now), 1)
+	require.NoError(t, err)
+
+	// Second sample, should be ok.
+	_, err = app.AppendHistogram(
+		0,
+		labels.FromStrings(model.MetricNameLabel, "my_histogram1"),
+		timestamp.FromTime(now),
+		&histogram.Histogram{},
+		nil,
+	)
+	require.NoError(t, err)
+
+	// This is third sample with sample_limit=2, it should trigger errSampleLimit.
+	_, err = app.AppendHistogram(
+		0,
+		labels.FromStrings(model.MetricNameLabel, "my_histogram2"),
+		timestamp.FromTime(now),
+		&histogram.Histogram{},
+		nil,
+	)
+	require.ErrorIs(t, err, errSampleLimit)
+}