Implement Memory Limiter scenario and validation tests

cmd/prometheus/main.go

@@ -83,6 +83,7 @@ import (
 	"github.com/prometheus/prometheus/util/documentcli"
 	"github.com/prometheus/prometheus/util/features"
 	"github.com/prometheus/prometheus/util/logging"
+	"github.com/prometheus/prometheus/util/memorylimiter"
 	"github.com/prometheus/prometheus/util/notifications"
 	prom_runtime "github.com/prometheus/prometheus/util/runtime"
 	"github.com/prometheus/prometheus/web"
@@ -215,6 +216,7 @@ type flagConfig struct {
 	enablePerStepStats       bool
 	enableConcurrentRuleEval bool
 	useStartTimestamps       bool
+	enableMemoryLimiter      bool
 
 	prometheusURL   string
 	corsRegexString string
@@ -332,6 +334,9 @@ func (c *flagConfig) setFeatureListOptions(logger *slog.Logger) error {
 			case "fast-startup":
 				c.tsdb.EnableFastStartup = true
 				logger.Info("Experimental fast startup is enabled.")
+			case "memory-limiter":
+				c.enableMemoryLimiter = true
+				logger.Info("Experimental memory limiter is enabled.")
 			default:
 				logger.Warn("Unknown option for --enable-feature", "option", o)
 			}
@@ -365,6 +370,7 @@ func main() {
 			collectors.NewGoCollector(
 				collectors.WithGoCollectorRuntimeMetrics(
 					collectors.MetricsGC,
+					collectors.MetricsMemory,
 					collectors.MetricsScheduler,
 					collectors.GoRuntimeMetricsRule{Matcher: goregexp.MustCompile(`^/sync/mutex/wait/total:seconds$`)},
 				),
@@ -620,7 +626,7 @@ func main() {
 	a.Flag("scrape.discovery-reload-interval", "Interval used by scrape manager to throttle target groups updates.").
 		Hidden().Default("5s").SetValue(&cfg.scrape.DiscoveryReloadInterval)
 
-	a.Flag("enable-feature", "Comma separated feature names to enable. Valid options: concurrent-rule-eval, created-timestamp-zero-ingestion, delayed-compaction, exemplar-storage, extra-scrape-metrics, memory-snapshot-on-shutdown, metadata-wal-records, old-ui, otlp-deltatocumulative, otlp-native-delta-ingestion, promql-binop-fill-modifiers, promql-delayed-name-removal, promql-duration-expr, promql-experimental-functions, promql-extended-range-selectors, promql-per-step-stats, st-storage, type-and-unit-labels, use-start-timestamps, use-uncached-io, xor2-encoding. See https://prometheus.io/docs/prometheus/latest/feature_flags/ for more details.").
+	a.Flag("enable-feature", "Comma separated feature names to enable. Valid options: concurrent-rule-eval, created-timestamp-zero-ingestion, delayed-compaction, exemplar-storage, extra-scrape-metrics, memory-limiter, memory-snapshot-on-shutdown, metadata-wal-records, old-ui, otlp-deltatocumulative, otlp-native-delta-ingestion, promql-binop-fill-modifiers, promql-delayed-name-removal, promql-duration-expr, promql-experimental-functions, promql-extended-range-selectors, promql-per-step-stats, st-storage, type-and-unit-labels, use-start-timestamps, use-uncached-io, xor2-encoding. See https://prometheus.io/docs/prometheus/latest/feature_flags/ for more details.").
 		StringsVar(&cfg.featureList)
 
 	a.Flag("agent", "Run Prometheus in 'Agent mode'.").BoolVar(&agentMode)
@@ -938,6 +944,18 @@ func main() {
 		os.Exit(1)
 	}
 
+	var memoryLimiter memorylimiter.MemoryLimiter
+	if cfg.enableMemoryLimiter {
+		ml, err := memorylimiter.NewManager(&cfgFile.Runtime.MemoryLimiter, logger.With("component", "memory limiter"), prometheus.DefaultRegisterer)
+		if err != nil {
+			logger.Error("failed to create memory limiter", "err", err)
+			os.Exit(1)
+		}
+		memoryLimiter = ml
+		cfg.scrape.MemoryLimiter = ml
+		cfg.web.MemoryLimiter = ml
+	}
+
 	scrapeManager, err := scrape.NewManager(
 		&cfg.scrape,
 		logger.With("component", "scrape manager"),
@@ -1000,6 +1018,7 @@ func main() {
 			},
 			FeatureRegistry: features.DefaultRegistry,
 			Parser:          promqlParser,
+			MemoryLimiter:   memoryLimiter,
 		})
 	}
 
@@ -1164,6 +1183,14 @@ func main() {
 		}, {
 			name:     "tracing",
 			reloader: tracingManager.ApplyConfig,
+		}, {
+			name: "memory_limiter",
+			reloader: func(cfg *config.Config) error {
+				if memoryLimiter != nil {
+					return memoryLimiter.ApplyConfig(&cfg.Runtime.MemoryLimiter)
+				}
+				return nil
+			},
 		},
 	}
 
@@ -1257,6 +1284,21 @@ func main() {
 			},
 		)
 	}
+	if memoryLimiter != nil {
+		ctxML, cancelML := context.WithCancel(context.Background())
+		g.Add(
+			func() error {
+				memoryLimiter.Start(ctxML)
+				<-ctxML.Done()
+				return nil
+			},
+			func(error) {
+				logger.Info("Stopping memory limiter...")
+				cancelML()
+				memoryLimiter.Stop()
+			},
+		)
+	}
 	if !agentMode {
 		// Rule manager.
 		g.Add(