Add nativeHistogram IngestionRate limit (#6794)

* Add native histograms ingestion rate limit

Signed-off-by: Paurush Garg <paurushg@amazon.com>

* Adding Tests and Updating Docs

Signed-off-by: Paurush Garg <paurushg@amazon.com>

* Resolving failed testcase

Signed-off-by: Paurush Garg <paurushg@amazon.com>

* Resolving comments

Signed-off-by: Paurush Garg <paurushg@amazon.com>

* Resolving comments

Signed-off-by: Paurush Garg <paurushg@amazon.com>

* Updating doc

Signed-off-by: Paurush Garg <paurushg@amazon.com>

* Changing NativeHistograms default ingestion limits

Signed-off-by: Paurush Garg <paurushg@amazon.com>

* Discard only NH Samples for NH Rate Limiter

Signed-off-by: Paurush Garg <paurushg@amazon.com>

* Adding more test coverage

Signed-off-by: Paurush Garg <paurushg@amazon.com>

---------

Signed-off-by: Paurush Garg <paurushg@amazon.com>

Author: PaurushGarg

CHANGELOG.md

@@ -40,6 +40,7 @@
 * [ENHANCEMENT] Distributor: Add min/max schema validation for NativeHistograms. #6766
 * [ENHANCEMENT] Ingester: Handle runtime errors in query path #6769
 * [ENHANCEMENT] Compactor: Support metadata caching bucket for Cleaner. Can be enabled via `-compactor.cleaner-caching-bucket-enabled` flag. #6778
+* [ENHANCEMENT] Distributor: Add ingestion rate limit for Native Histograms. #6794
 * [ENHANCEMENT] Compactor, Store Gateway: Introduce user scanner strategy and user index. #6780
 * [ENHANCEMENT] Querier: Support chunks cache for parquet queryable. #6805
 * [ENHANCEMENT] Parquet Storage: Add some metrics for parquet blocks and converter. #6809 #6821

docs/configuration/config-file-reference.md

@@ -3476,6 +3476,11 @@ The `limits_config` configures default and per-tenant limits imposed by Cortex s
 # CLI flag: -distributor.ingestion-rate-limit
 [ingestion_rate: <float> | default = 25000]
 
+# Per-user native histogram ingestion rate limit in samples per second. Disabled
+# by default
+# CLI flag: -distributor.native-histogram-ingestion-rate-limit
+[native_histogram_ingestion_rate: <float> | default = 1.7976931348623157e+308]
+
 # Whether the ingestion rate limit should be applied individually to each
 # distributor instance (local), or evenly shared across the cluster (global).
 # CLI flag: -distributor.ingestion-rate-limit-strategy
@@ -3485,6 +3490,10 @@ The `limits_config` configures default and per-tenant limits imposed by Cortex s
 # CLI flag: -distributor.ingestion-burst-size
 [ingestion_burst_size: <int> | default = 50000]
 
+# Per-user allowed native histogram ingestion burst size (in number of samples)
+# CLI flag: -distributor.native-histogram-ingestion-burst-size
+[native_histogram_ingestion_burst_size: <int> | default = 0]
+
 # Flag to enable, for all users, handling of samples with external labels
 # identifying replicas in an HA Prometheus setup.
 # CLI flag: -distributor.ha-tracker.enable-for-all-users

pkg/distributor/distributor.go

@@ -95,7 +95,8 @@ type Distributor struct {
 	HATracker *ha.HATracker
 
 	// Per-user rate limiter.
-	ingestionRateLimiter *limiter.RateLimiter
+	ingestionRateLimiter                *limiter.RateLimiter
+	nativeHistogramIngestionRateLimiter *limiter.RateLimiter
 
 	// Manager for subservices (HA Tracker, distributor ring and client pool)
 	subservices        *services.Manager
@@ -267,11 +268,13 @@ func New(cfg Config, clientConfig ingester_client.Config, limits *validation.Ove
 	// it's an internal dependency and can't join the distributors ring, we skip rate
 	// limiting.
 	var ingestionRateStrategy limiter.RateLimiterStrategy
+	var nativeHistogramIngestionRateStrategy limiter.RateLimiterStrategy
 	var distributorsLifeCycler *ring.Lifecycler
 	var distributorsRing *ring.Ring
 
 	if !canJoinDistributorsRing {
 		ingestionRateStrategy = newInfiniteIngestionRateStrategy()
+		nativeHistogramIngestionRateStrategy = newInfiniteIngestionRateStrategy()
 	} else if limits.IngestionRateStrategy() == validation.GlobalIngestionRateStrategy {
 		distributorsLifeCycler, err = ring.NewLifecycler(cfg.DistributorRing.ToLifecyclerConfig(), nil, "distributor", ringKey, true, true, log, prometheus.WrapRegistererWithPrefix("cortex_", reg))
 		if err != nil {
@@ -285,21 +288,24 @@ func New(cfg Config, clientConfig ingester_client.Config, limits *validation.Ove
 		subservices = append(subservices, distributorsLifeCycler, distributorsRing)
 
 		ingestionRateStrategy = newGlobalIngestionRateStrategy(limits, distributorsLifeCycler)
+		nativeHistogramIngestionRateStrategy = newGlobalNativeHistogramIngestionRateStrategy(limits, distributorsLifeCycler)
 	} else {
 		ingestionRateStrategy = newLocalIngestionRateStrategy(limits)
+		nativeHistogramIngestionRateStrategy = newLocalNativeHistogramIngestionRateStrategy(limits)
 	}
 
 	d := &Distributor{
-		cfg:                    cfg,
-		log:                    log,
-		ingestersRing:          ingestersRing,
-		ingesterPool:           NewPool(cfg.PoolConfig, ingestersRing, cfg.IngesterClientFactory, log),
-		distributorsLifeCycler: distributorsLifeCycler,
-		distributorsRing:       distributorsRing,
-		limits:                 limits,
-		ingestionRateLimiter:   limiter.NewRateLimiter(ingestionRateStrategy, 10*time.Second),
-		HATracker:              haTracker,
-		ingestionRate:          util_math.NewEWMARate(0.2, instanceIngestionRateTickInterval),
+		cfg:                                 cfg,
+		log:                                 log,
+		ingestersRing:                       ingestersRing,
+		ingesterPool:                        NewPool(cfg.PoolConfig, ingestersRing, cfg.IngesterClientFactory, log),
+		distributorsLifeCycler:              distributorsLifeCycler,
+		distributorsRing:                    distributorsRing,
+		limits:                              limits,
+		ingestionRateLimiter:                limiter.NewRateLimiter(ingestionRateStrategy, 10*time.Second),
+		nativeHistogramIngestionRateLimiter: limiter.NewRateLimiter(nativeHistogramIngestionRateStrategy, 10*time.Second),
+		HATracker:                           haTracker,
+		ingestionRate:                       util_math.NewEWMARate(0.2, instanceIngestionRateTickInterval),
 
 		queryDuration: instrument.NewHistogramCollector(promauto.With(reg).NewHistogramVec(prometheus.HistogramOpts{
 			Namespace: "cortex",
@@ -754,7 +760,7 @@ func (d *Distributor) Push(ctx context.Context, req *cortexpb.WriteRequest) (*co
 	}
 
 	// A WriteRequest can only contain series or metadata but not both. This might change in the future.
-	seriesKeys, validatedTimeseries, validatedFloatSamples, validatedHistogramSamples, validatedExemplars, firstPartialErr, err := d.prepareSeriesKeys(ctx, req, userID, limits, removeReplica)
+	seriesKeys, nhSeriesKeys, validatedTimeseries, nhValidatedTimeseries, validatedFloatSamples, validatedHistogramSamples, validatedExemplars, firstPartialErr, err := d.prepareSeriesKeys(ctx, req, userID, limits, removeReplica)
 	if err != nil {
 		return nil, err
 	}
@@ -765,6 +771,15 @@ func (d *Distributor) Push(ctx context.Context, req *cortexpb.WriteRequest) (*co
 	d.receivedExemplars.WithLabelValues(userID).Add(float64(validatedExemplars))
 	d.receivedMetadata.WithLabelValues(userID).Add(float64(len(validatedMetadata)))
 
+	if !d.nativeHistogramIngestionRateLimiter.AllowN(now, userID, validatedHistogramSamples) {
+		level.Warn(d.log).Log("msg", "native histogram ingestion rate limit (%v) exceeded while adding %d native histogram samples", d.nativeHistogramIngestionRateLimiter.Limit(now, userID), validatedHistogramSamples)
+		d.validateMetrics.DiscardedSamples.WithLabelValues(validation.NativeHistogramRateLimited, userID).Add(float64(validatedHistogramSamples))
+		validatedHistogramSamples = 0
+	} else {
+		seriesKeys = append(seriesKeys, nhSeriesKeys...)
+		validatedTimeseries = append(validatedTimeseries, nhValidatedTimeseries...)
+	}
+
 	if len(seriesKeys) == 0 && len(metadataKeys) == 0 {
 		// Ensure the request slice is reused if there's no series or metadata passing the validation.
 		cortexpb.ReuseSlice(req.Timeseries)
@@ -936,14 +951,16 @@ type samplesLabelSetEntry struct {
 	labels           labels.Labels
 }
 
-func (d *Distributor) prepareSeriesKeys(ctx context.Context, req *cortexpb.WriteRequest, userID string, limits *validation.Limits, removeReplica bool) ([]uint32, []cortexpb.PreallocTimeseries, int, int, int, error, error) {
+func (d *Distributor) prepareSeriesKeys(ctx context.Context, req *cortexpb.WriteRequest, userID string, limits *validation.Limits, removeReplica bool) ([]uint32, []uint32, []cortexpb.PreallocTimeseries, []cortexpb.PreallocTimeseries, int, int, int, error, error) {
 	pSpan, _ := opentracing.StartSpanFromContext(ctx, "prepareSeriesKeys")
 	defer pSpan.Finish()
 
 	// For each timeseries or samples, we compute a hash to distribute across ingesters;
-	// check each sample/metadata and discard if outside limits.
+	// check each sample/metadata and discard  if outside limits.
 	validatedTimeseries := make([]cortexpb.PreallocTimeseries, 0, len(req.Timeseries))
+	nhValidatedTimeseries := make([]cortexpb.PreallocTimeseries, 0, len(req.Timeseries))
 	seriesKeys := make([]uint32, 0, len(req.Timeseries))
+	nhSeriesKeys := make([]uint32, 0, len(req.Timeseries))
 	validatedFloatSamples := 0
 	validatedHistogramSamples := 0
 	validatedExemplars := 0
@@ -1051,7 +1068,7 @@ func (d *Distributor) prepareSeriesKeys(ctx context.Context, req *cortexpb.Write
 		// label and dropped labels (if any)
 		key, err := d.tokenForLabels(userID, ts.Labels)
 		if err != nil {
-			return nil, nil, 0, 0, 0, nil, err
+			return nil, nil, nil, nil, 0, 0, 0, nil, err
 		}
 		validatedSeries, validationErr := d.validateSeries(ts, userID, skipLabelNameValidation, limits)
 
@@ -1086,8 +1103,13 @@ func (d *Distributor) prepareSeriesKeys(ctx context.Context, req *cortexpb.Write
 			}
 		}
 
-		seriesKeys = append(seriesKeys, key)
-		validatedTimeseries = append(validatedTimeseries, validatedSeries)
+		if len(ts.Histograms) > 0 {
+			nhSeriesKeys = append(nhSeriesKeys, key)
+			nhValidatedTimeseries = append(nhValidatedTimeseries, validatedSeries)
+		} else {
+			seriesKeys = append(seriesKeys, key)
+			validatedTimeseries = append(validatedTimeseries, validatedSeries)
+		}
 		validatedFloatSamples += len(ts.Samples)
 		validatedHistogramSamples += len(ts.Histograms)
 		validatedExemplars += len(ts.Exemplars)
@@ -1103,7 +1125,7 @@ func (d *Distributor) prepareSeriesKeys(ctx context.Context, req *cortexpb.Write
 		}
 	}
 
-	return seriesKeys, validatedTimeseries, validatedFloatSamples, validatedHistogramSamples, validatedExemplars, firstPartialErr, nil
+	return seriesKeys, nhSeriesKeys, validatedTimeseries, nhValidatedTimeseries, validatedFloatSamples, validatedHistogramSamples, validatedExemplars, firstPartialErr, nil
 }
 
 func sortLabelsIfNeeded(labels []cortexpb.LabelAdapter) {

pkg/distributor/distributor_test.go

@@ -681,6 +681,147 @@ func TestDistributor_PushIngestionRateLimiter(t *testing.T) {
 	}
 }
 
+func TestDistributor_PushIngestionRateLimiter_Histograms(t *testing.T) {
+	t.Parallel()
+	type testPush struct {
+		samples                              int
+		nhSamples                            int
+		metadata                             int
+		expectedError                        error
+		expectedNHDiscardedSampleMetricValue int
+	}
+
+	ctx := user.InjectOrgID(context.Background(), "user")
+	tests := map[string]struct {
+		distributors                         int
+		ingestionRateStrategy                string
+		ingestionRate                        float64
+		ingestionBurstSize                   int
+		nativeHistogramIngestionRateStrategy string
+		nativeHistogramIngestionRate         float64
+		nativeHistogramIngestionBurstSize    int
+		pushes                               []testPush
+	}{
+		"local strategy: native histograms limit should be set to each distributor": {
+			distributors:                      2,
+			ingestionRateStrategy:             validation.LocalIngestionRateStrategy,
+			ingestionRate:                     20,
+			ingestionBurstSize:                20,
+			nativeHistogramIngestionRate:      10,
+			nativeHistogramIngestionBurstSize: 10,
+			pushes: []testPush{
+				{nhSamples: 4, expectedError: nil},
+				{nhSamples: 6, expectedError: nil},
+				{nhSamples: 6, expectedError: nil, expectedNHDiscardedSampleMetricValue: 6},
+				{nhSamples: 4, metadata: 1, expectedError: nil, expectedNHDiscardedSampleMetricValue: 10},
+				{nhSamples: 1, expectedError: nil, expectedNHDiscardedSampleMetricValue: 11},
+				{nhSamples: 1, expectedError: nil, expectedNHDiscardedSampleMetricValue: 12},
+			},
+		},
+		"global strategy: native histograms limit should be evenly shared across distributors": {
+			distributors:                      2,
+			ingestionRateStrategy:             validation.GlobalIngestionRateStrategy,
+			ingestionRate:                     20,
+			ingestionBurstSize:                10,
+			nativeHistogramIngestionRate:      10,
+			nativeHistogramIngestionBurstSize: 5,
+			pushes: []testPush{
+				{nhSamples: 2, expectedError: nil},
+				{nhSamples: 1, expectedError: nil},
+				{nhSamples: 3, metadata: 1, expectedError: nil, expectedNHDiscardedSampleMetricValue: 3},
+				{nhSamples: 2, expectedError: nil, expectedNHDiscardedSampleMetricValue: 3},
+				{nhSamples: 1, expectedError: nil, expectedNHDiscardedSampleMetricValue: 4},
+				{nhSamples: 1, expectedError: nil, expectedNHDiscardedSampleMetricValue: 5},
+			},
+		},
+		"global strategy: native histograms burst should set to each distributor": {
+			distributors:                      2,
+			ingestionRateStrategy:             validation.GlobalIngestionRateStrategy,
+			ingestionRate:                     20,
+			ingestionBurstSize:                40,
+			nativeHistogramIngestionRate:      10,
+			nativeHistogramIngestionBurstSize: 20,
+			pushes: []testPush{
+				{nhSamples: 10, expectedError: nil},
+				{nhSamples: 5, expectedError: nil},
+				{nhSamples: 6, metadata: 1, expectedError: nil, expectedNHDiscardedSampleMetricValue: 6},
+				{nhSamples: 5, expectedError: nil, expectedNHDiscardedSampleMetricValue: 6},
+				{nhSamples: 1, expectedError: nil, expectedNHDiscardedSampleMetricValue: 7},
+				{nhSamples: 1, expectedError: nil, expectedNHDiscardedSampleMetricValue: 8},
+			},
+		},
+		"local strategy: If NH samples hit NH rate limit, other samples should succeed when under rate limit": {
+			distributors:                      2,
+			ingestionRateStrategy:             validation.LocalIngestionRateStrategy,
+			ingestionRate:                     20,
+			ingestionBurstSize:                20,
+			nativeHistogramIngestionRate:      5,
+			nativeHistogramIngestionBurstSize: 5,
+			pushes: []testPush{
+				{samples: 5, nhSamples: 4, expectedError: nil},
+				{samples: 6, nhSamples: 2, expectedError: nil, expectedNHDiscardedSampleMetricValue: 2},
+				{samples: 4, metadata: 1, nhSamples: 1, expectedError: httpgrpc.Errorf(http.StatusTooManyRequests, "ingestion rate limit (20) exceeded while adding 5 samples and 1 metadata"), expectedNHDiscardedSampleMetricValue: 2},
+				{metadata: 1, expectedError: nil, expectedNHDiscardedSampleMetricValue: 2},
+			},
+		},
+		"global strategy: If NH samples hit NH rate limit, other samples should succeed when under rate limit": {
+			distributors:                      2,
+			ingestionRateStrategy:             validation.GlobalIngestionRateStrategy,
+			ingestionRate:                     20,
+			ingestionBurstSize:                10,
+			nativeHistogramIngestionRate:      10,
+			nativeHistogramIngestionBurstSize: 5,
+			pushes: []testPush{
+				{samples: 3, nhSamples: 2, expectedError: nil},
+				{samples: 3, nhSamples: 4, expectedError: nil, expectedNHDiscardedSampleMetricValue: 4},
+				{samples: 1, metadata: 1, nhSamples: 1, expectedError: httpgrpc.Errorf(http.StatusTooManyRequests, "ingestion rate limit (10) exceeded while adding 2 samples and 1 metadata"), expectedNHDiscardedSampleMetricValue: 4},
+				{metadata: 1, expectedError: nil, expectedNHDiscardedSampleMetricValue: 4},
+			},
+		},
+	}
+
+	for testName, testData := range tests {
+		testData := testData
+
+		t.Run(testName, func(t *testing.T) {
+			t.Parallel()
+			limits := &validation.Limits{}
+			flagext.DefaultValues(limits)
+			limits.IngestionRateStrategy = testData.ingestionRateStrategy
+			limits.IngestionRate = testData.ingestionRate
+			limits.IngestionBurstSize = testData.ingestionBurstSize
+			limits.NativeHistogramIngestionRate = testData.nativeHistogramIngestionRate
+			limits.NativeHistogramIngestionBurstSize = testData.nativeHistogramIngestionBurstSize
+
+			// Start all expected distributors
+			distributors, _, _, _ := prepare(t, prepConfig{
+				numIngesters:     3,
+				happyIngesters:   3,
+				numDistributors:  testData.distributors,
+				shardByAllLabels: true,
+				limits:           limits,
+			})
+
+			// Push samples in multiple requests to the first distributor
+			for _, push := range testData.pushes {
+				var request = makeWriteRequest(0, push.samples, push.metadata, push.nhSamples)
+
+				response, err := distributors[0].Push(ctx, request)
+
+				if push.expectedError == nil {
+					assert.Equal(t, emptyResponse, response)
+					assert.Nil(t, err)
+				} else {
+					assert.Nil(t, response)
+					assert.Equal(t, push.expectedError, err)
+				}
+				assert.Equal(t, float64(push.expectedNHDiscardedSampleMetricValue), testutil.ToFloat64(distributors[0].validateMetrics.DiscardedSamples.WithLabelValues(validation.NativeHistogramRateLimited, "user")))
+			}
+		})
+	}
+
+}
+
 func TestPush_QuorumError(t *testing.T) {
 	t.Parallel()
 

pkg/distributor/ingestion_rate_strategy.go

@@ -45,7 +45,7 @@ func newGlobalIngestionRateStrategy(limits *validation.Overrides, ring ReadLifec
 func (s *globalStrategy) Limit(tenantID string) float64 {
 	numDistributors := s.ring.HealthyInstancesCount()
 
-	if numDistributors == 0 {
+	if numDistributors == 0 || s.limits.IngestionRate(tenantID) == float64(rate.Inf) {
 		return s.limits.IngestionRate(tenantID)
 	}
 
@@ -72,3 +72,49 @@ func (s *infiniteStrategy) Burst(tenantID string) int {
 	// Burst is ignored when limit = rate.Inf
 	return 0
 }
+
+type localStrategyNativeHistogram struct {
+	limits *validation.Overrides
+}
+
+func newLocalNativeHistogramIngestionRateStrategy(limits *validation.Overrides) limiter.RateLimiterStrategy {
+	return &localStrategyNativeHistogram{
+		limits: limits,
+	}
+}
+
+func (s *localStrategyNativeHistogram) Limit(tenantID string) float64 {
+	return s.limits.NativeHistogramIngestionRate(tenantID)
+}
+
+func (s *localStrategyNativeHistogram) Burst(tenantID string) int {
+	return s.limits.NativeHistogramIngestionBurstSize(tenantID)
+}
+
+type globalStrategyNativeHistogram struct {
+	limits *validation.Overrides
+	ring   ReadLifecycler
+}
+
+func newGlobalNativeHistogramIngestionRateStrategy(limits *validation.Overrides, ring ReadLifecycler) limiter.RateLimiterStrategy {
+	return &globalStrategyNativeHistogram{
+		limits: limits,
+		ring:   ring,
+	}
+}
+
+func (s *globalStrategyNativeHistogram) Limit(tenantID string) float64 {
+	numDistributors := s.ring.HealthyInstancesCount()
+
+	if numDistributors == 0 || s.limits.NativeHistogramIngestionRate(tenantID) == float64(rate.Inf) {
+		return s.limits.NativeHistogramIngestionRate(tenantID)
+	}
+
+	return s.limits.NativeHistogramIngestionRate(tenantID) / float64(numDistributors)
+}
+
+func (s *globalStrategyNativeHistogram) Burst(tenantID string) int {
+	// The meaning of burst doesn't change for the global strategy, in order
+	// to keep it easier to understand for users / operators.
+	return s.limits.NativeHistogramIngestionBurstSize(tenantID)
+}