Receive: make tsdb stats limit configurable

Signed-off-by: Michael Hoffmann <[email protected]>

pkg/receive/handler.go

@@ -54,12 +54,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 +284,18 @@ 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, 64)
+	if err != nil {
+		return 0, fmt.Errorf("unable to parse '%s' parameter: %w", LimitStatsQueryParam, err)
+	}
+	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 +308,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.

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()

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))
 		})
 	}