diff --git a/pkg/receive/handler.go b/pkg/receive/handler.go index 3cce0552ab8..1359a911c78 100644 --- a/pkg/receive/handler.go +++ b/pkg/receive/handler.go @@ -10,6 +10,7 @@ import ( "fmt" "io" stdlog "log" + "math" "net" "net/http" "path" @@ -54,12 +55,16 @@ const ( DefaultTenantHeader = "THANOS-TENANT" // DefaultTenant is the default value used for when no tenant is passed via the tenant header. DefaultTenant = "default-tenant" + // DefaultStatsLimit is the default value used for limiting tenant stats. + DefaultStatsLimit = 10 // DefaultTenantLabel is the default label-name used for when no tenant is passed via the tenant header. DefaultTenantLabel = "tenant_id" // DefaultReplicaHeader is the default header used to designate the replica count of a write request. DefaultReplicaHeader = "THANOS-REPLICA" // AllTenantsQueryParam is the query parameter for getting TSDB stats for all tenants. AllTenantsQueryParam = "all_tenants" + // LimitStatsQueryParam is the query parameter for limiting the amount of returned TSDB stats. + LimitStatsQueryParam = "limit" // Labels for metrics. labelSuccess = "success" labelError = "error" @@ -280,6 +285,21 @@ func (h *Handler) testReady(f http.HandlerFunc) http.HandlerFunc { } } +func (h *Handler) getStatsLimitParameter(r *http.Request) (int, error) { + limitStatsStr := r.FormValue(LimitStatsQueryParam) + if limitStatsStr == "" { + return DefaultStatsLimit, nil + } + limitStats, err := strconv.ParseInt(limitStatsStr, 10, 0) + if err != nil { + return 0, fmt.Errorf("unable to parse '%s' parameter: %w", LimitStatsQueryParam, err) + } + if limitStats > math.MaxInt { + return 0, fmt.Errorf("'%s' parameter is larger than %d", LimitStatsQueryParam, math.MaxInt) + } + return int(limitStats), nil +} + func (h *Handler) getStats(r *http.Request, statsByLabelName string) ([]statusapi.TenantStats, *api.ApiError) { if !h.isReady() { return nil, &api.ApiError{Typ: api.ErrorInternal, Err: fmt.Errorf("service unavailable")} @@ -292,15 +312,20 @@ func (h *Handler) getStats(r *http.Request, statsByLabelName string) ([]statusap return nil, &api.ApiError{Typ: api.ErrorBadData, Err: err} } + statsLimit, err := h.getStatsLimitParameter(r) + if err != nil { + return nil, &api.ApiError{Typ: api.ErrorBadData, Err: err} + } + if getAllTenantStats { - return h.options.TSDBStats.TenantStats(statsByLabelName), nil + return h.options.TSDBStats.TenantStats(statsLimit, statsByLabelName), nil } if tenantID == "" { tenantID = h.options.DefaultTenantID } - return h.options.TSDBStats.TenantStats(statsByLabelName, tenantID), nil + return h.options.TSDBStats.TenantStats(statsLimit, statsByLabelName, tenantID), nil } // Close stops the Handler. diff --git a/pkg/receive/multitsdb.go b/pkg/receive/multitsdb.go index 2fceb1ac77b..349ceb98ef5 100644 --- a/pkg/receive/multitsdb.go +++ b/pkg/receive/multitsdb.go @@ -44,7 +44,7 @@ import ( type TSDBStats interface { // TenantStats returns TSDB head stats for the given tenants. // If no tenantIDs are provided, stats for all tenants are returned. - TenantStats(statsByLabelName string, tenantIDs ...string) []status.TenantStats + TenantStats(limit int, statsByLabelName string, tenantIDs ...string) []status.TenantStats } type MultiTSDB struct { @@ -518,7 +518,7 @@ func (t *MultiTSDB) TSDBExemplars() map[string]*exemplars.TSDB { return res } -func (t *MultiTSDB) TenantStats(statsByLabelName string, tenantIDs ...string) []status.TenantStats { +func (t *MultiTSDB) TenantStats(limit int, statsByLabelName string, tenantIDs ...string) []status.TenantStats { t.mtx.RLock() defer t.mtx.RUnlock() if len(tenantIDs) == 0 { @@ -545,7 +545,7 @@ func (t *MultiTSDB) TenantStats(statsByLabelName string, tenantIDs ...string) [] if db == nil { return } - stats := db.Head().Stats(statsByLabelName, 10) + stats := db.Head().Stats(statsByLabelName, limit) mu.Lock() defer mu.Unlock() diff --git a/pkg/receive/multitsdb_test.go b/pkg/receive/multitsdb_test.go index 9d5b0ac26d9..58241cd5a05 100644 --- a/pkg/receive/multitsdb_test.go +++ b/pkg/receive/multitsdb_test.go @@ -568,7 +568,7 @@ func TestMultiTSDBStats(t *testing.T) { testutil.Ok(t, appendSample(m, "baz", time.Now())) testutil.Equals(t, 3, len(m.TSDBLocalClients())) - stats := m.TenantStats(labels.MetricName, test.tenants...) + stats := m.TenantStats(10, labels.MetricName, test.tenants...) testutil.Equals(t, test.expectedStats, len(stats)) }) }