SH-470 api rate limiter per host (#2054)

* SH-470 api rate limiter per host

* fix test

* fix test

* disable option

* builtin

* builtin

* default fix

---------

Co-authored-by: ru.nazarov <[email protected]>

Author: Rustamchuk

cmd/statshouse-api/statshouse-api.go

@@ -175,6 +175,8 @@ func run() int {
 			ShardByMetricShards: argv.ShardByMetricShards,
 			MaxShardConnsRatio:  argv.CHMaxShardConnsRatio,
 			DialTimeout:         chDialTimeout,
+			SelectTimeout:       argv.QuerySelectTimeout,
+			RateLimitConfig:     argv.RateLimitConfig,
 			ConnLimits: chutil.ConnLimits{
 				FastLightMaxConns: argv.chV1MaxConns,
 				FastHeavyMaxConns: argv.chV1MaxConns,
@@ -196,6 +198,8 @@ func run() int {
 		ShardByMetricShards: argv.ShardByMetricShards,
 		MaxShardConnsRatio:  argv.CHMaxShardConnsRatio,
 		DialTimeout:         chDialTimeout,
+		SelectTimeout:       argv.QuerySelectTimeout,
+		RateLimitConfig:     argv.RateLimitConfig,
 		ConnLimits: chutil.ConnLimits{
 			FastLightMaxConns:    argv.chV2MaxLightFastConns,
 			FastHeavyMaxConns:    argv.chV2MaxHeavyFastConns,

internal/api/config.go

@@ -38,6 +38,7 @@ type Config struct {
 	AvailableShardsStr     string
 	AvailableShards        []uint32
 	CHMaxShardConnsRatio   int
+	chutil.RateLimitConfig
 }
 
 func (argv *Config) ValidateConfig() error {
@@ -116,13 +117,37 @@ func (argv *Config) Bind(f *flag.FlagSet, defaultI config.Config) {
 	f.StringVar(&argv.BlockedUsersS, "blocked-users", "", "comma-separated list of users that are blocked")
 	f.StringVar(&argv.AvailableShardsStr, "available-shards", default_.AvailableShardsStr, "comma-separated list of default shards for metrics when namespace doesn't specify shards")
 	f.IntVar(&argv.CHMaxShardConnsRatio, "clickhouse-max-shard-conns-ratio", default_.CHMaxShardConnsRatio, "maximum number of ClickHouse connections per shard (%)")
+
+	f.BoolVar(&argv.RateLimitDisable, "rate-limit-disable", false, "disable rate limiting")
+	f.DurationVar(&argv.WindowDuration, "rate-limit-window-duration", default_.WindowDuration, "time window for analyzing ClickHouse requests")
+	f.Uint64Var(&argv.MaxErrorRate, "rate-limit-max-error-rate", default_.MaxErrorRate, "error rate threshold (%) to trigger sleep stage")
+	f.Uint64Var(&argv.MaxInflightWeight, "rate-limit-max-inflight-weight", default_.MaxInflightWeight, "maximum weight per inflight request (%)")
+	f.Uint64Var(&argv.RecoverWeightStep, "rate-limit-recover-weight-step", default_.RecoverWeightStep, "recover weight step (%) per recovery cycle")
+	f.DurationVar(&argv.RecoverGapDuration, "rate-limit-recover-gap-duration", default_.RecoverGapDuration, "recover time gap per recovery cycle")
+	f.Float64Var(&argv.SleepMultiplier, "rate-limit-sleep-multiplier", default_.SleepMultiplier, "sleep duration multiplier after failed check")
+	f.DurationVar(&argv.MaxSleepDuration, "rate-limit-max-sleep-duration", default_.MaxSleepDuration, "maximum sleep duration limit")
+	f.DurationVar(&argv.MinSleepDuration, "rate-limit-min-sleep-duration", default_.MinSleepDuration, "minimum sleep duration when entering sleep")
+	f.Uint64Var(&argv.CheckCount, "rate-limit-check-count", default_.CheckCount, "successful checks needed to exit check stage")
+	f.DurationVar(&argv.RecalcInterval, "recalc-interval-sleep-duration", default_.RecalcInterval, "rate limit state recalculation interval")
 }
 
 func DefaultConfig() *Config {
 	return &Config{
 		ApproxCacheMaxSize:   1_000_000,
 		AvailableShardsStr:   "1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16",
 		CHMaxShardConnsRatio: 20,
+		RateLimitConfig: chutil.RateLimitConfig{
+			WindowDuration:     2 * time.Minute,
+			MaxErrorRate:       20,
+			MaxInflightWeight:  10,
+			RecoverWeightStep:  1,
+			RecoverGapDuration: 6 * time.Second,
+			SleepMultiplier:    1.5,
+			MaxSleepDuration:   2 * time.Minute,
+			MinSleepDuration:   30 * time.Second,
+			CheckCount:         9,
+			RecalcInterval:     1 * time.Second,
+		},
 	}
 }
 
@@ -154,7 +179,7 @@ type HandlerOptions struct {
 	timezone                 string
 	protectedMetricPrefixesS string
 	protectedMetricPrefixes  []string
-	querySelectTimeout       time.Duration
+	QuerySelectTimeout       time.Duration
 	weekStartAt              int
 	location                 *time.Location
 	utcOffset                int64
@@ -167,7 +192,7 @@ func (argv *HandlerOptions) Bind(f *flag.FlagSet) {
 	f.BoolVar(&argv.querySequential, "query-sequential", false, "disables query parallel execution")
 	f.BoolVar(&argv.readOnly, "readonly", false, "read only mode")
 	f.BoolVar(&argv.verbose, "verbose", false, "verbose logging")
-	f.DurationVar(&argv.querySelectTimeout, "query-select-timeout", QuerySelectTimeoutDefault, "query select timeout")
+	f.DurationVar(&argv.QuerySelectTimeout, "query-select-timeout", QuerySelectTimeoutDefault, "query select timeout")
 	f.StringVar(&argv.protectedMetricPrefixesS, "protected-metric-prefixes", "", "comma-separated list of metric prefixes that require access bits set")
 	f.StringVar(&argv.timezone, "timezone", "Europe/Moscow", "location of the desired timezone")
 	f.IntVar(&argv.weekStartAt, "week-start", int(time.Monday), "week day of beginning of the week (from sunday=0 to saturday=6)")

internal/api/endpoint_stat.go

@@ -186,6 +186,18 @@ func ChSelectActiveQueries(client *statshouse.Client, versionTag string, tagLane
 	}
 }
 
+func ChRateLimit(client *statshouse.Client, versionTag string, stats []chutil.RateLimitMetric) {
+	for _, stat := range stats {
+		metric := client.MetricRef(format.BuiltinMetricMetaAPIRateLimit.Name, statshouse.Tags{
+			2: versionTag,
+			3: stat.Stage,
+			4: strconv.FormatUint(stat.InflightWeight, 10),
+			5: strconv.Itoa(stat.ShardKey),
+			6: strconv.Itoa(stat.ReplicaKey)})
+		metric.Value(float64(stat.InflightCnt))
+	}
+}
+
 func ChRequestsMetric(shard int, aggHost string, table string, ok bool) {
 	status := "1"
 	if !ok {

internal/api/handler.go

@@ -686,7 +686,7 @@ func NewHandler(staticDir fs.FS, jsSettings JSSettings, showInvisible bool, chV1
 		h.Version3Start.Store(cfg.Version3Start)
 		h.Version3Prob.Store(cfg.Version3Prob)
 		h.Version3StrcmpOff.Store(cfg.Version3StrcmpOff)
-		chV2.SetLimits(cfg.UserLimits, cfg.CHMaxShardConnsRatio)
+		chV2.SetLimits(cfg.UserLimits, cfg.CHMaxShardConnsRatio, cfg.RateLimitConfig)
 		h.NewShardingStart.Store(cfg.NewShardingStart)
 		h.ConfigMu.Lock()
 		h.DisableCHAddr = cfg.DisableCHAddr
@@ -739,6 +739,7 @@ func NewHandler(staticDir fs.FS, jsSettings JSSettings, showInvisible bool, chV1
 				ChSelectActiveQueries(client, versionTag, format.TagValueIDAPILaneSlowHeavy, ch.SemaphoreCountSlowHeavy(), ch.ShardSemaphoreCountSlowHeavy())
 				ChSelectActiveQueries(client, versionTag, format.TagValueIDAPILaneFastHardware, ch.SemaphoreCountFastHardware(), ch.ShardSemaphoreCountFastHardware())
 				ChSelectActiveQueries(client, versionTag, format.TagValueIDAPILaneSlowHardware, ch.SemaphoreCountSlowHardware(), ch.ShardSemaphoreCountSlowHardware())
+				ChRateLimit(client, versionTag, ch.RateLimitStatistics())
 			}
 		}
 		writeActiveQuieries(chV1, "1")
@@ -1855,7 +1856,7 @@ func (h *Handler) handlePostMetric(ctx context.Context, ai accessInfo, _ string,
 }
 
 func HandleGetMetricTagValues(r *httpRequestHandler) {
-	ctx, cancel := context.WithTimeout(r.Context(), r.querySelectTimeout)
+	ctx, cancel := context.WithTimeout(r.Context(), r.QuerySelectTimeout)
 	defer cancel()
 
 	_ = r.ParseForm() // (*http.Request).FormValue ignores parse errors, too
@@ -2113,7 +2114,7 @@ func HandleGetTable(r *httpRequestHandler) {
 		return
 	}
 
-	ctx, cancel := context.WithTimeout(r.Context(), r.querySelectTimeout)
+	ctx, cancel := context.WithTimeout(r.Context(), r.QuerySelectTimeout)
 	defer cancel()
 
 	if req.numResults <= 0 || maxTableRowsPage < req.numResults {
@@ -2163,7 +2164,7 @@ func HandleBadgesQuery(r *httpRequestHandler) {
 		respondJSON(r, nil, 0, 0, err)
 		return
 	}
-	ctx, cancel := context.WithTimeout(r.Context(), r.querySelectTimeout)
+	ctx, cancel := context.WithTimeout(r.Context(), r.QuerySelectTimeout)
 	defer cancel()
 	query := promql.Query{
 		Start: req.from.Unix(),
@@ -2367,7 +2368,7 @@ func HandlePointQuery(r *httpRequestHandler) {
 }
 
 func HandleGetRender(r *httpRequestHandler) {
-	ctx, cancel := context.WithTimeout(r.Context(), r.querySelectTimeout)
+	ctx, cancel := context.WithTimeout(r.Context(), r.QuerySelectTimeout)
 	defer cancel()
 	s, err := r.parseSeriesRequestS(12)
 	if err != nil {
@@ -2614,7 +2615,7 @@ func (h *requestHandler) handleSeriesRequestS(ctx context.Context, req seriesReq
 	var cancelT func()
 	var freeBadges func()
 	var freeRes func()
-	ctx, cancelT = context.WithTimeout(ctx, h.querySelectTimeout)
+	ctx, cancelT = context.WithTimeout(ctx, h.QuerySelectTimeout)
 	cancel := func() {
 		cancelT()
 		if freeBadges != nil {

internal/api/promql.go

@@ -62,7 +62,7 @@ func HandleInstantQuery(r *httpRequestHandler) {
 	}
 	q.End = q.Start + 1 // handler expects half open interval [start, end)
 	// execute query
-	ctx, cancel := context.WithTimeout(r.Context(), r.querySelectTimeout)
+	ctx, cancel := context.WithTimeout(r.Context(), r.QuerySelectTimeout)
 	defer cancel()
 	res, dispose, err := r.promEngine.Exec(ctx, r, q)
 	if err != nil {
@@ -103,7 +103,7 @@ func HandleRangeQuery(r *httpRequestHandler) {
 		return
 	}
 	// execute query
-	ctx, cancel := context.WithTimeout(r.Context(), r.querySelectTimeout)
+	ctx, cancel := context.WithTimeout(r.Context(), r.QuerySelectTimeout)
 	defer cancel()
 	res, dispose, err := r.promEngine.Exec(ctx, r, q)
 	if err != nil {