feat: Map streams to scopes and add per-scope ingestion block.

docs/sources/shared/configuration.md

@@ -3941,6 +3941,25 @@ otlp_config:
 # CLI flag: -limits.block-ingestion-status-code
 [block_ingestion_status_code: <int> | default = 260]
 
+# Block ingestion until the given time for the given policy. Pushes will be
+# assigned to a policy based on the stream matcher configuration. Experimental.
+[block_policy_ingestion_until: <map of string to Time>]
+
+# HTTP status code to return when ingestion is blocked for the given policy.
+# Experimental.
+[block_policy_ingestion_status_code: <map of string to int>]
+
+# Map of policies to stream selectors. Push streams that matches a policy
+# selector will be considered as belonging to that policy. If that policy is
+# blocked, the push will be rejected with the status code specified in
+# block_policy_ingestion_status_code. Experimental.
+[policy_stream_mapping: <map of string to string>]
+
+# Map of policies to enforced labels. Push streams that matches a policy
+# selector will be considered as belonging to that policy and as such, the
+# labels related to the policy will be enforced. Experimental.
+[policy_enforced_labels: <map of string to list of strings>]
+
 # List of labels that must be present in the stream. If any of the labels are
 # missing, the stream will be discarded. This flag configures it globally for
 # all tenants. Experimental.

pkg/distributor/distributor.go

@@ -453,8 +453,10 @@ func (d *Distributor) Push(ctx context.Context, req *logproto.PushRequest) (*log
 	// We use the heuristic of 1 sample per TS to size the array.
 	// We also work out the hash value at the same time.
 	streams := make([]KeyedStream, 0, len(req.Streams))
-	validatedLineSize := 0
-	validatedLineCount := 0
+	validatedLineCountTotal := 0
+	validatedLineSizeTotal := 0
+	validatedLineSizePerRetention := make(map[string]int)  // map of retention period to validated line size
+	validatedLineCountPerRetention := make(map[string]int) // map of retention period to validated line count
 
 	var validationErrors util.GroupedErrors
 
@@ -512,33 +514,35 @@ func (d *Distributor) Push(ctx context.Context, req *logproto.PushRequest) (*log
 			// Truncate first so subsequent steps have consistent line lengths
 			d.truncateLines(validationContext, &stream)
 
+			tenantRetention := d.validator.Limits.RetentionPeriod(tenantID)
+
+			initialEntriesSize := util.EntriesTotalSize(stream.Entries)
+
 			var lbs labels.Labels
 			lbs, stream.Labels, stream.Hash, err = d.parseStreamLabels(validationContext, stream.Labels, stream)
 			if err != nil {
 				d.writeFailuresManager.Log(tenantID, err)
 				validationErrors.Add(err)
-				validation.DiscardedSamples.WithLabelValues(validation.InvalidLabels, tenantID).Add(float64(len(stream.Entries)))
-				discardedBytes := util.EntriesTotalSize(stream.Entries)
-				validation.DiscardedBytes.WithLabelValues(validation.InvalidLabels, tenantID).Add(float64(discardedBytes))
+				validation.DiscardedSamples.WithLabelValues(validation.InvalidLabels, tenantID, tenantRetention.String()).Add(float64(len(stream.Entries)))
+				validation.DiscardedBytes.WithLabelValues(validation.InvalidLabels, tenantID, tenantRetention.String()).Add(float64(initialEntriesSize))
 				continue
 			}
 
-			if missing, lbsMissing := d.missingEnforcedLabels(lbs, tenantID); missing {
-				err := fmt.Errorf(validation.MissingEnforcedLabelsErrorMsg, strings.Join(lbsMissing, ","), tenantID)
+			if errorLbValue, err := d.missingEnforcedLabels(lbs, tenantID); err != nil {
 				d.writeFailuresManager.Log(tenantID, err)
 				validationErrors.Add(err)
-				validation.DiscardedSamples.WithLabelValues(validation.MissingEnforcedLabels, tenantID).Add(float64(len(stream.Entries)))
-				discardedBytes := util.EntriesTotalSize(stream.Entries)
-				validation.DiscardedBytes.WithLabelValues(validation.MissingEnforcedLabels, tenantID).Add(float64(discardedBytes))
+				validation.DiscardedSamples.WithLabelValues(errorLbValue, tenantID, tenantRetention.String()).Add(float64(len(stream.Entries)))
+				validation.DiscardedBytes.WithLabelValues(errorLbValue, tenantID, tenantRetention.String()).Add(float64(initialEntriesSize))
 				continue
 			}
 
 			n := 0
 			pushSize := 0
 			prevTs := stream.Entries[0].Timestamp
+			streamRetention := d.retentionPeriodForStream(lbs, tenantID)
 
 			for _, entry := range stream.Entries {
-				if err := d.validator.ValidateEntry(ctx, validationContext, lbs, entry); err != nil {
+				if err := d.validator.ValidateEntry(ctx, validationContext, lbs, entry, streamRetention); err != nil {
 					d.writeFailuresManager.Log(tenantID, err)
 					validationErrors.Add(err)
 					continue
@@ -593,16 +597,36 @@ func (d *Distributor) Push(ctx context.Context, req *logproto.PushRequest) (*log
 				}
 
 				n++
-				validatedLineSize += util.EntryTotalSize(&entry)
-				validatedLineCount++
-				pushSize += len(entry.Line)
+
+				entrySize := util.EntryTotalSize(&entry)
+				validatedLineSizePerRetention[streamRetention.String()] += entrySize
+				validatedLineCountPerRetention[streamRetention.String()]++
+				validatedLineSizeTotal += entrySize
+				validatedLineCountTotal++
+				pushSize += entrySize
 			}
+
 			stream.Entries = stream.Entries[:n]
 			if len(stream.Entries) == 0 {
 				// Empty stream after validating all the entries
 				continue
 			}
 
+			if policy := d.policyForStream(lbs, tenantID); policy != "" {
+				streamSize := util.EntriesTotalSize(stream.Entries)
+				if yes, until, retStatusCode := d.validator.ShouldBlockIngestionForPolicy(validationContext, policy, now); yes {
+					d.trackDiscardedDataFromPolicy(ctx, policy, tenantID, validation.BlockedPolicyIngestion, streamRetention, streamSize, lbs)
+
+					err := fmt.Errorf(validation.BlockedPolicyIngestionErrorMsg, tenantID, until.Format(time.RFC3339), retStatusCode, policy)
+					d.writeFailuresManager.Log(tenantID, err)
+
+					validationErrors.Add(err)
+					validation.DiscardedSamples.WithLabelValues(validation.BlockedPolicyIngestion, tenantID, streamRetention.String()).Add(float64(len(stream.Entries)))
+					validation.DiscardedBytes.WithLabelValues(validation.BlockedPolicyIngestion, tenantID, streamRetention.String()).Add(float64(streamSize))
+					continue
+				}
+			}
+
 			maybeShardStreams(stream, lbs, pushSize)
 		}
 	}()
@@ -617,8 +641,8 @@ func (d *Distributor) Push(ctx context.Context, req *logproto.PushRequest) (*log
 		return &logproto.PushResponse{}, validationErr
 	}
 
-	if block, until, retStatusCode := d.validator.ShouldBlockIngestion(validationContext, now); block {
-		d.trackDiscardedData(ctx, req, validationContext, tenantID, validatedLineCount, validatedLineSize, validation.BlockedIngestion)
+	if block, until, retStatusCode := d.validator.ShouldBlockIngestionForTenant(validationContext, now); block {
+		d.trackDiscardedData(ctx, req, validationContext, tenantID, validatedLineCountPerRetention, validatedLineSizePerRetention, validation.BlockedIngestion)
 
 		err = fmt.Errorf(validation.BlockedIngestionErrorMsg, tenantID, until.Format(time.RFC3339), retStatusCode)
 		d.writeFailuresManager.Log(tenantID, err)
@@ -632,10 +656,10 @@ func (d *Distributor) Push(ctx context.Context, req *logproto.PushRequest) (*log
 		return nil, httpgrpc.Errorf(retStatusCode, "%s", err.Error())
 	}
 
-	if !d.ingestionRateLimiter.AllowN(now, tenantID, validatedLineSize) {
-		d.trackDiscardedData(ctx, req, validationContext, tenantID, validatedLineCount, validatedLineSize, validation.RateLimited)
+	if !d.ingestionRateLimiter.AllowN(now, tenantID, validatedLineSizeTotal) {
+		d.trackDiscardedData(ctx, req, validationContext, tenantID, validatedLineCountPerRetention, validatedLineSizePerRetention, validation.RateLimited)
 
-		err = fmt.Errorf(validation.RateLimitedErrorMsg, tenantID, int(d.ingestionRateLimiter.Limit(now, tenantID)), validatedLineCount, validatedLineSize)
+		err = fmt.Errorf(validation.RateLimitedErrorMsg, tenantID, int(d.ingestionRateLimiter.Limit(now, tenantID)), validatedLineCountTotal, validatedLineSizeTotal)
 		d.writeFailuresManager.Log(tenantID, err)
 		// Return a 429 to indicate to the client they are being rate limited
 		return nil, httpgrpc.Errorf(http.StatusTooManyRequests, "%s", err.Error())
@@ -743,14 +767,39 @@ func (d *Distributor) Push(ctx context.Context, req *logproto.PushRequest) (*log
 	}
 }
 
-// missingEnforcedLabels returns true if the stream is missing any of the required labels.
+func (d *Distributor) retentionPeriodForStream(lbs labels.Labels, tenantID string) time.Duration {
+	retentions := d.validator.Limits.StreamRetention(tenantID)
+
+	for _, retention := range retentions {
+		if retention.Matches(lbs) {
+			return time.Duration(retention.Period)
+		}
+	}
+
+	// No retention for this specific stream, use the default retention period.
+	return d.validator.Limits.RetentionPeriod(tenantID)
+}
+
+func (d *Distributor) missingEnforcedLabels(lbs labels.Labels, tenantID string) (string, error) {
+	if err := d.missingTenantEnforcedLabels(lbs, tenantID); err != nil {
+		return validation.MissingEnforcedLabels, err
+	}
+
+	if err := d.missingPolicyEnforcedLabels(lbs, tenantID); err != nil {
+		return validation.MissingPolicyEnforcedLabels, err
+	}
+
+	return "", nil
+}
+
+// missingTenantEnforcedLabels returns true if the stream is missing any of the required labels for the tenant.
 //
 // It also returns the first label that is missing if any (for the case of multiple labels missing).
-func (d *Distributor) missingEnforcedLabels(lbs labels.Labels, tenantID string) (bool, []string) {
+func (d *Distributor) missingTenantEnforcedLabels(lbs labels.Labels, tenantID string) error {
 	requiredLbs := d.validator.Limits.EnforcedLabels(tenantID)
 	if len(requiredLbs) == 0 {
 		// no enforced labels configured.
-		return false, []string{}
+		return nil
 	}
 
 	missingLbs := []string{}
@@ -761,20 +810,85 @@ func (d *Distributor) missingEnforcedLabels(lbs labels.Labels, tenantID string)
 		}
 	}
 
-	return len(missingLbs) > 0, missingLbs
+	if len(missingLbs) > 0 {
+		return fmt.Errorf(validation.MissingEnforcedLabelsErrorMsg, strings.Join(missingLbs, ","), tenantID)
+	}
+
+	return nil
+}
+
+func (d *Distributor) policyForStream(lbs labels.Labels, tenantID string) string {
+	policyStreamMapping := d.validator.Limits.PolicyStreamMapping(tenantID)
+	if len(policyStreamMapping) == 0 {
+		// not configured.
+		return ""
+	}
+
+	for policy, streamSelector := range policyStreamMapping {
+		matchers, err := syntax.ParseMatchers(streamSelector, true)
+		if err != nil {
+			level.Error(d.logger).Log("msg", "failed to parse stream selector", "error", err, "stream_selector", streamSelector, "tenant_id", tenantID)
+			continue
+		}
+
+		if validation.LabelMatchesMatchers(lbs, matchers) {
+			return policy
+		}
+	}
+
+	return ""
+}
+
+func (d *Distributor) missingPolicyEnforcedLabels(lbs labels.Labels, tenantID string) error {
+	policyEnforcedLabels := d.validator.Limits.PolicyEnforcedLabels(tenantID)
+	if len(policyEnforcedLabels) == 0 {
+		// not configured.
+		return nil
+	}
+
+	policy := d.policyForStream(lbs, tenantID)
+	if policy == "" {
+		// no policy found for stream.
+		return nil
+	}
+
+	missingLbs := []string{}
+
+	for _, lb := range policyEnforcedLabels[policy] {
+		if !lbs.Has(lb) {
+			missingLbs = append(missingLbs, lb)
+		}
+	}
+
+	if len(missingLbs) > 0 {
+		return fmt.Errorf(validation.MissingPolicyEnforcedLabelsErrorMsg, strings.Join(missingLbs, ","), policy, tenantID)
+	}
+
+	return nil
+}
+
+func (d *Distributor) trackDiscardedDataFromPolicy(ctx context.Context, policy string, tenantID string, reason string, retention time.Duration, discardedStreamBytes int, lbs labels.Labels) {
+	validation.DiscardedSamplesByPolicy.WithLabelValues(reason, policy, tenantID, retention.String()).Add(float64(1))
+	validation.DiscardedBytesByPolicy.WithLabelValues(reason, policy, tenantID, retention.String()).Add(float64(discardedStreamBytes))
+
+	if d.usageTracker != nil {
+		d.usageTracker.DiscardedBytesAddByPolicy(ctx, tenantID, policy, reason, retention, lbs, float64(discardedStreamBytes))
+	}
 }
 
 func (d *Distributor) trackDiscardedData(
 	ctx context.Context,
 	req *logproto.PushRequest,
 	validationContext validationContext,
 	tenantID string,
-	validatedLineCount int,
-	validatedLineSize int,
+	validatedLineCount map[string]int,
+	validatedLineSize map[string]int,
 	reason string,
 ) {
-	validation.DiscardedSamples.WithLabelValues(reason, tenantID).Add(float64(validatedLineCount))
-	validation.DiscardedBytes.WithLabelValues(reason, tenantID).Add(float64(validatedLineSize))
+	for retention, count := range validatedLineCount {
+		validation.DiscardedSamples.WithLabelValues(reason, tenantID, retention).Add(float64(count))
+		validation.DiscardedBytes.WithLabelValues(reason, tenantID, retention).Add(float64(validatedLineSize[retention]))
+	}
 
 	if d.usageTracker != nil {
 		for _, stream := range req.Streams {