Add ResourceMonitor module in Cortex, and add ResourceBasedLimiter in Ingesters and StoreGateways

CHANGELOG.md

@@ -5,6 +5,8 @@
 * [FEATURE] Query Frontend: Add dynamic interval size for query splitting. This is enabled by configuring experimental flags `querier.max-shards-per-query` and/or `querier.max-fetched-data-duration-per-query`. The split interval size is dynamically increased to maintain a number of shards and total duration fetched below the configured values. #6458
 * [FEATURE] Querier/Ruler: Add `query_partial_data` and `rules_partial_data` limits to allow queries/rules to be evaluated with data from a single zone, if other zones are not available. #6526
 * [FEATURE] Update prometheus alertmanager version to v0.28.0 and add new integration msteamsv2, jira, and rocketchat. #6590
+* [FEATURE] Ingester: Add a `-ingester.enable-ooo-native-histograms` flag to enable out-of-order native histogram ingestion per tenant. It only takes effect when `-blocks-storage.tsdb.enable-native-histograms=true` and `-ingester.out-of-order-time-window` > 0. It is applied after the restart if it is changed at runtime through the runtime config. #6626
+* [FEATURE] Ingester/StoreGateway: Add `resource-thresholds` in ingesters and store gateways to throttle query requests when the pods are under resource pressure. #6674
 * [FEATURE] Ingester: Support out-of-order native histogram ingestion. It automatically enabled when `-ingester.out-of-order-time-window > 0` and `-blocks-storage.tsdb.enable-native-histograms=true`. #6626 #6663
 * [ENHANCEMENT] Alertmanager: Add nflog and silences maintenance metrics. #6659
 * [ENHANCEMENT] Querier: limit label APIs to query only ingesters if `start` param is not been specified. #6618

cmd/cortex/main.go

@@ -18,6 +18,7 @@ import (
 	"github.com/prometheus/client_golang/prometheus"
 	collectorversion "github.com/prometheus/client_golang/prometheus/collectors/version"
 	"github.com/prometheus/common/version"
+	_ "go.uber.org/automaxprocs"
 	"gopkg.in/yaml.v2"
 
 	"github.com/cortexproject/cortex/pkg/cortex"

docs/configuration/config-file-reference.md

@@ -271,6 +271,15 @@ query_scheduler:
     # CLI flag: -query-scheduler.grpc-client-config.connect-timeout
     [connect_timeout: <duration> | default = 5s]
 
+resource_thresholds:
+  # Utilization threshold for CPU in percentage, between 0 and 1. 0 to disable.
+  # CLI flag: -resource-thresholds.cpu
+  [cpu: <float> | default = 0]
+
+  # Utilization threshold for heap in percentage, between 0 and 1. 0 to disable.
+  # CLI flag: -resource-thresholds.heap
+  [heap: <float> | default = 0]
+
 # The tracing_config configures backends cortex uses.
 [tracing: <tracing_config>]
 ```

go.mod

@@ -81,9 +81,11 @@ require (
 	github.com/google/go-cmp v0.7.0
 	github.com/hashicorp/golang-lru/v2 v2.0.7
 	github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822
+	github.com/prometheus/procfs v0.15.1
 	github.com/sercand/kuberesolver/v5 v5.1.1
 	github.com/tjhop/slog-gokit v0.1.2
 	go.opentelemetry.io/collector/pdata v1.24.0
+	go.uber.org/automaxprocs v1.6.0
 	google.golang.org/protobuf v1.36.4
 )
 
@@ -199,7 +201,6 @@ require (
 	github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect
 	github.com/prometheus-community/prom-label-proxy v0.8.1-0.20240127162815-c1195f9aabc0 // indirect
 	github.com/prometheus/exporter-toolkit v0.13.2 // indirect
-	github.com/prometheus/procfs v0.15.1 // indirect
 	github.com/prometheus/sigv4 v0.1.1 // indirect
 	github.com/redis/rueidis v1.0.45-alpha.1 // indirect
 	github.com/rs/cors v1.11.1 // indirect

go.sum

@@ -1560,6 +1560,8 @@ github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRI
 github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/posener/complete v1.1.1/go.mod h1:em0nMJCgc9GFtwrmVmEMR/ZL6WyhyjMBndrE9hABlRI=
 github.com/posener/complete v1.2.3/go.mod h1:WZIdtGGp+qx0sLrYKtIRAruyNpv6hFCicSgv7Sy7s/s=
+github.com/prashantv/gostub v1.1.0 h1:BTyx3RfQjRHnUWaGF9oQos79AlQ5k8WNktv7VGvVH4g=
+github.com/prashantv/gostub v1.1.0/go.mod h1:A5zLQHz7ieHGG7is6LLXLz7I8+3LZzsrV0P1IAHhP5U=
 github.com/prometheus-community/prom-label-proxy v0.8.1-0.20240127162815-c1195f9aabc0 h1:owfYHh79h8Y5HvNMGyww+DaVwo10CKiRW1RQrrZzIwg=
 github.com/prometheus-community/prom-label-proxy v0.8.1-0.20240127162815-c1195f9aabc0/go.mod h1:rT989D4UtOcfd9tVqIZRVIM8rkg+9XbreBjFNEKXvVI=
 github.com/prometheus/alertmanager v0.28.1 h1:BK5pCoAtaKg01BYRUJhEDV1tqJMEtYBGzPw8QdvnnvA=
@@ -1811,6 +1813,8 @@ go.opentelemetry.io/proto/otlp v1.5.0/go.mod h1:keN8WnHxOy8PG0rQZjJJ5A2ebUoafqWp
 go.uber.org/atomic v1.7.0/go.mod h1:fEN4uk6kAWBTFdckzkM89CLk9XfWZrxpCo0nPH17wJc=
 go.uber.org/atomic v1.11.0 h1:ZvwS0R+56ePWxUNi+Atn9dWONBPp/AUETXlHW0DxSjE=
 go.uber.org/atomic v1.11.0/go.mod h1:LUxbIzbOniOlMKjJjyPfpl4v+PKK2cNJn91OQbhoJI0=
+go.uber.org/automaxprocs v1.6.0 h1:O3y2/QNTOdbF+e/dpXNNW7Rx2hZ4sTIPyybbxyNqTUs=
+go.uber.org/automaxprocs v1.6.0/go.mod h1:ifeIMSnPZuznNm6jmdzmU3/bfk01Fe2fotchwEFJ8r8=
 go.uber.org/goleak v1.1.10/go.mod h1:8a7PlsEVH3e/a/GLqe5IIrQx6GzcnRmZEufDUTk4A7A=
 go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto=
 go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE=

pkg/configs/resources.go

@@ -0,0 +1,26 @@
+package configs
+
+import (
+	"errors"
+	"flag"
+)
+
+var ErrInvalidResourceThreshold = errors.New("invalid resource utilization threshold, it must be between 0 and 1")
+
+type Resources struct {
+	CPU  float64 `yaml:"cpu"`
+	Heap float64 `yaml:"heap"`
+}
+
+func (cfg *Resources) RegisterFlags(f *flag.FlagSet) {
+	f.Float64Var(&cfg.CPU, "resource-thresholds.cpu", 0, "Utilization threshold for CPU in percentage, between 0 and 1. 0 to disable.")
+	f.Float64Var(&cfg.Heap, "resource-thresholds.heap", 0, "Utilization threshold for heap in percentage, between 0 and 1. 0 to disable.")
+}
+
+func (cfg *Resources) Validate() error {
+	if cfg.CPU > 1 || cfg.CPU < 0 || cfg.Heap > 1 || cfg.Heap < 0 {
+		return ErrInvalidResourceThreshold
+	}
+
+	return nil
+}

pkg/cortex/cortex.go

@@ -22,6 +22,7 @@ import (
 	"gopkg.in/yaml.v2"
 
 	"github.com/cortexproject/cortex/pkg/util/grpcclient"
+	"github.com/cortexproject/cortex/pkg/util/resource"
 
 	"github.com/cortexproject/cortex/pkg/alertmanager"
 	"github.com/cortexproject/cortex/pkg/alertmanager/alertstore"
@@ -122,6 +123,7 @@ type Config struct {
 	RuntimeConfig       runtimeconfig.Config                       `yaml:"runtime_config"`
 	MemberlistKV        memberlist.KVConfig                        `yaml:"memberlist"`
 	QueryScheduler      scheduler.Config                           `yaml:"query_scheduler"`
+	ResourceThresholds  configs.Resources                          `yaml:"resource_thresholds"`
 
 	Tracing tracing.Config `yaml:"tracing"`
 }
@@ -170,6 +172,7 @@ func (c *Config) RegisterFlags(f *flag.FlagSet) {
 	c.MemberlistKV.RegisterFlags(f)
 	c.QueryScheduler.RegisterFlags(f)
 	c.Tracing.RegisterFlags(f)
+	c.ResourceThresholds.RegisterFlags(f)
 }
 
 // Validate the cortex config and returns an error if the validation
@@ -237,6 +240,10 @@ func (c *Config) Validate(log log.Logger) error {
 		return errors.Wrap(err, "invalid tracing config")
 	}
 
+	if err := c.ResourceThresholds.Validate(); err != nil {
+		return errors.Wrap(err, "invalid resource_thresholds config")
+	}
+
 	return nil
 }
 
@@ -315,6 +322,7 @@ type Cortex struct {
 	MetadataQuerier          querier.MetadataQuerier
 	QuerierEngine            promql.QueryEngine
 	QueryFrontendTripperware tripperware.Tripperware
+	ResourceMonitor          *resource.Monitor
 
 	Ruler        *ruler.Ruler
 	RulerStorage rulestore.RuleStore

pkg/cortex/cortex_test.go

@@ -22,6 +22,7 @@ import (
 
 	"github.com/cortexproject/cortex/pkg/alertmanager"
 	"github.com/cortexproject/cortex/pkg/alertmanager/alertstore"
+	"github.com/cortexproject/cortex/pkg/configs"
 	"github.com/cortexproject/cortex/pkg/cortex/storage"
 	"github.com/cortexproject/cortex/pkg/frontend/v1/frontendv1pb"
 	"github.com/cortexproject/cortex/pkg/ingester"
@@ -165,11 +166,20 @@ func TestConfigValidation(t *testing.T) {
 			},
 			expectedError: errInvalidHTTPPrefix,
 		},
+		{
+			name: "should fail validation for resource thresholds",
+			getTestConfig: func() *Config {
+				configuration := newDefaultConfig()
+				configuration.ResourceThresholds.CPU = 10
+				return configuration
+			},
+			expectedError: configs.ErrInvalidResourceThreshold,
+		},
 	} {
 		t.Run(tc.name, func(t *testing.T) {
 			err := tc.getTestConfig().Validate(nil)
 			if tc.expectedError != nil {
-				require.Equal(t, tc.expectedError, err)
+				require.ErrorContains(t, err, tc.expectedError.Error())
 			} else {
 				require.NoError(t, err)
 			}

pkg/cortex/modules.go

@@ -6,6 +6,8 @@ import (
 	"fmt"
 	"log/slog"
 	"net/http"
+	"runtime"
+	"runtime/debug"
 
 	"github.com/go-kit/log/level"
 	"github.com/opentracing-contrib/go-stdlib/nethttp"
@@ -27,6 +29,7 @@ import (
 	"github.com/cortexproject/cortex/pkg/alertmanager/alertstore"
 	"github.com/cortexproject/cortex/pkg/api"
 	"github.com/cortexproject/cortex/pkg/compactor"
+	"github.com/cortexproject/cortex/pkg/configs"
 	configAPI "github.com/cortexproject/cortex/pkg/configs/api"
 	"github.com/cortexproject/cortex/pkg/configs/db"
 	"github.com/cortexproject/cortex/pkg/distributor"
@@ -51,6 +54,7 @@ import (
 	"github.com/cortexproject/cortex/pkg/util/grpcclient"
 	util_log "github.com/cortexproject/cortex/pkg/util/log"
 	"github.com/cortexproject/cortex/pkg/util/modules"
+	"github.com/cortexproject/cortex/pkg/util/resource"
 	"github.com/cortexproject/cortex/pkg/util/runtimeconfig"
 	"github.com/cortexproject/cortex/pkg/util/services"
 	"github.com/cortexproject/cortex/pkg/util/validation"
@@ -86,6 +90,7 @@ const (
 	Purger                   string = "purger"
 	QueryScheduler           string = "query-scheduler"
 	TenantFederation         string = "tenant-federation"
+	ResourceMonitor          string = "resource-monitor"
 	All                      string = "all"
 )
 
@@ -441,7 +446,7 @@ func (t *Cortex) initIngesterService() (serv services.Service, err error) {
 	t.Cfg.Ingester.QueryIngestersWithin = t.Cfg.Querier.QueryIngestersWithin
 	t.tsdbIngesterConfig()
 
-	t.Ingester, err = ingester.New(t.Cfg.Ingester, t.Overrides, prometheus.DefaultRegisterer, util_log.Logger)
+	t.Ingester, err = ingester.New(t.Cfg.Ingester, t.Overrides, prometheus.DefaultRegisterer, util_log.Logger, t.ResourceMonitor)
 	if err != nil {
 		return
 	}
@@ -705,7 +710,7 @@ func (t *Cortex) initCompactor() (serv services.Service, err error) {
 func (t *Cortex) initStoreGateway() (serv services.Service, err error) {
 	t.Cfg.StoreGateway.ShardingRing.ListenPort = t.Cfg.Server.GRPCListenPort
 
-	t.StoreGateway, err = storegateway.NewStoreGateway(t.Cfg.StoreGateway, t.Cfg.BlocksStorage, t.Overrides, t.Cfg.Server.LogLevel, util_log.Logger, prometheus.DefaultRegisterer)
+	t.StoreGateway, err = storegateway.NewStoreGateway(t.Cfg.StoreGateway, t.Cfg.BlocksStorage, t.Overrides, t.Cfg.Server.LogLevel, util_log.Logger, prometheus.DefaultRegisterer, t.ResourceMonitor)
 	if err != nil {
 		return nil, err
 	}
@@ -765,11 +770,32 @@ func (t *Cortex) initQueryScheduler() (services.Service, error) {
 	return s, nil
 }
 
+func (t *Cortex) initResourceMonitor() (services.Service, error) {
+	if t.Cfg.ResourceThresholds.CPU <= 0 && t.Cfg.ResourceThresholds.Heap <= 0 {
+		return nil, nil
+	}
+
+	scanner, err := resource.NewScanner()
+	if err != nil {
+		return nil, err
+	}
+
+	limits := configs.Resources{
+		CPU:  float64(runtime.GOMAXPROCS(0)),
+		Heap: float64(debug.SetMemoryLimit(-1)),
+	}
+
+	t.ResourceMonitor, err = resource.NewMonitor(t.Cfg.ResourceThresholds, limits, scanner, prometheus.DefaultRegisterer)
+
+	return t.ResourceMonitor, err
+}
+
 func (t *Cortex) setupModuleManager() error {
 	mm := modules.NewManager(util_log.Logger)
 
 	// Register all modules here.
 	// RegisterModule(name string, initFn func()(services.Service, error))
+	mm.RegisterModule(ResourceMonitor, t.initResourceMonitor)
 	mm.RegisterModule(Server, t.initServer, modules.UserInvisibleModule)
 	mm.RegisterModule(API, t.initAPI, modules.UserInvisibleModule)
 	mm.RegisterModule(RuntimeConfig, t.initRuntimeConfig, modules.UserInvisibleModule)
@@ -811,7 +837,7 @@ func (t *Cortex) setupModuleManager() error {
 		Distributor:              {DistributorService, API, GrpcClientService},
 		DistributorService:       {Ring, Overrides},
 		Ingester:                 {IngesterService, Overrides, API},
-		IngesterService:          {Overrides, RuntimeConfig, MemberlistKV},
+		IngesterService:          {Overrides, RuntimeConfig, MemberlistKV, ResourceMonitor},
 		Flusher:                  {Overrides, API},
 		Queryable:                {Overrides, DistributorService, Overrides, Ring, API, StoreQueryable, MemberlistKV},
 		Querier:                  {TenantFederation},
@@ -824,7 +850,7 @@ func (t *Cortex) setupModuleManager() error {
 		Configs:                  {API},
 		AlertManager:             {API, MemberlistKV, Overrides},
 		Compactor:                {API, MemberlistKV, Overrides},
-		StoreGateway:             {API, Overrides, MemberlistKV},
+		StoreGateway:             {API, Overrides, MemberlistKV, ResourceMonitor},
 		TenantDeletion:           {API, Overrides},
 		Purger:                   {TenantDeletion},
 		TenantFederation:         {Queryable},

pkg/frontend/transport/handler.go

@@ -28,6 +28,7 @@ import (
 	"github.com/cortexproject/cortex/pkg/util"
 	util_api "github.com/cortexproject/cortex/pkg/util/api"
 	util_log "github.com/cortexproject/cortex/pkg/util/log"
+	"github.com/cortexproject/cortex/pkg/util/resource"
 )
 
 const (
@@ -49,6 +50,7 @@ const (
 	reasonRequestBodySizeExceeded  = "request_body_size_exceeded"
 	reasonResponseBodySizeExceeded = "response_body_size_exceeded"
 	reasonTooManyRequests          = "too_many_requests"
+	reasonResourceExhausted        = "resource_exhausted"
 	reasonTimeRangeExceeded        = "time_range_exceeded"
 	reasonTooManySamples           = "too_many_samples"
 	reasonSeriesFetched            = "series_fetched"
@@ -489,30 +491,38 @@ func (f *Handler) reportQueryStats(r *http.Request, source, userID string, query
 	}
 
 	var reason string
-	if statusCode == http.StatusTooManyRequests {
-		reason = reasonTooManyRequests
-	} else if statusCode == http.StatusRequestEntityTooLarge {
+	if statusCode == http.StatusRequestEntityTooLarge {
 		reason = reasonResponseBodySizeExceeded
-	} else if statusCode == http.StatusUnprocessableEntity {
+	} else if error != nil {
+		// We are unable to use errors.As to compare since body string from the http response is wrapped as an error
 		errMsg := error.Error()
-		if strings.Contains(errMsg, limitTooManySamples) {
-			reason = reasonTooManySamples
-		} else if strings.Contains(errMsg, limitTimeRangeExceeded) {
-			reason = reasonTimeRangeExceeded
-		} else if strings.Contains(errMsg, limitSeriesFetched) {
-			reason = reasonSeriesFetched
-		} else if strings.Contains(errMsg, limitChunksFetched) {
-			reason = reasonChunksFetched
-		} else if strings.Contains(errMsg, limitChunkBytesFetched) {
-			reason = reasonChunkBytesFetched
-		} else if strings.Contains(errMsg, limitDataBytesFetched) {
-			reason = reasonDataBytesFetched
-		} else if strings.Contains(errMsg, limitSeriesStoreGateway) {
-			reason = reasonSeriesLimitStoreGateway
-		} else if strings.Contains(errMsg, limitChunksStoreGateway) {
-			reason = reasonChunksLimitStoreGateway
-		} else if strings.Contains(errMsg, limitBytesStoreGateway) {
-			reason = reasonBytesLimitStoreGateway
+		if statusCode == http.StatusTooManyRequests {
+			var resourceExhaustedErr resource.ExhaustedError
+			if strings.Contains(errMsg, resourceExhaustedErr.Error()) {
+				reason = reasonResourceExhausted
+			} else {
+				reason = reasonTooManyRequests
+			}
+		} else if statusCode == http.StatusUnprocessableEntity {
+			if strings.Contains(errMsg, limitTooManySamples) {
+				reason = reasonTooManySamples
+			} else if strings.Contains(errMsg, limitTimeRangeExceeded) {
+				reason = reasonTimeRangeExceeded
+			} else if strings.Contains(errMsg, limitSeriesFetched) {
+				reason = reasonSeriesFetched
+			} else if strings.Contains(errMsg, limitChunksFetched) {
+				reason = reasonChunksFetched
+			} else if strings.Contains(errMsg, limitChunkBytesFetched) {
+				reason = reasonChunkBytesFetched
+			} else if strings.Contains(errMsg, limitDataBytesFetched) {
+				reason = reasonDataBytesFetched
+			} else if strings.Contains(errMsg, limitSeriesStoreGateway) {
+				reason = reasonSeriesLimitStoreGateway
+			} else if strings.Contains(errMsg, limitChunksStoreGateway) {
+				reason = reasonChunksLimitStoreGateway
+			} else if strings.Contains(errMsg, limitBytesStoreGateway) {
+				reason = reasonBytesLimitStoreGateway
+			}
 		}
 	}
 	if len(reason) > 0 {

pkg/frontend/transport/handler_test.go

@@ -30,6 +30,7 @@ import (
 	"github.com/cortexproject/cortex/pkg/tenant"
 	util_api "github.com/cortexproject/cortex/pkg/util/api"
 	util_log "github.com/cortexproject/cortex/pkg/util/log"
+	"github.com/cortexproject/cortex/pkg/util/resource"
 )
 
 type roundTripperFunc func(*http.Request) (*http.Response, error)
@@ -380,6 +381,23 @@ func TestHandler_ServeHTTP(t *testing.T) {
 			},
 			expectedStatusCode: http.StatusUnprocessableEntity,
 		},
+		{
+			name:            "test handler with reasonResourceExhausted",
+			cfg:             HandlerConfig{QueryStatsEnabled: true},
+			expectedMetrics: 6,
+			roundTripperFunc: roundTripperFunc(func(req *http.Request) (*http.Response, error) {
+				resourceExhaustedErr := &resource.ExhaustedError{}
+				return &http.Response{
+					StatusCode: http.StatusTooManyRequests,
+					Body:       io.NopCloser(strings.NewReader(resourceExhaustedErr.Error())),
+				}, nil
+			}),
+			additionalMetricsCheckFunc: func(h *Handler) {
+				v := promtest.ToFloat64(h.rejectedQueries.WithLabelValues(reasonResourceExhausted, tripperware.SourceAPI, userID))
+				assert.Equal(t, float64(1), v)
+			},
+			expectedStatusCode: http.StatusTooManyRequests,
+		},
 	} {
 		t.Run(tt.name, func(t *testing.T) {
 			reg := prometheus.NewPedanticRegistry()