Implement maxTanant lmit on regex matcher

Signed-off-by: SungJin1212 <[email protected]>

Author: SungJin1212

CHANGELOG.md

@@ -18,6 +18,7 @@
 * [FEATURE] Distributor/Ingester: Implemented experimental feature to use gRPC stream connection for push requests. This can be enabled by setting `-distributor.use-stream-push=true`. #6580
 * [FEATURE] Compactor: Add support for percentage based sharding for compactors. #6738
 * [FEATURE] Querier: Allow choosing PromQL engine via header. #6777
+* [ENHANCEMENT] Tenant Federation: Add a # of query result limit logic when the `-tenant-federation.regex-matcher-enabled` is enabled. #6845
 * [ENHANCEMENT] Query Frontend: Change to return 400 when the tenant resolving fail. #6715
 * [ENHANCEMENT] Querier: Support query parameters to metadata api (/api/v1/metadata) to allow user to limit metadata to return. Add a `-ingester.return-all-metadata` flag to make the metadata API run when the deployment. Please set this flag to `false` to use the metadata API with the limits later. #6681 #6744
 * [ENHANCEMENT] Ingester: Add a `cortex_ingester_active_native_histogram_series` metric to track # of active NH series. #6695

pkg/cortex/modules.go

@@ -295,7 +295,7 @@ func (t *Cortex) initTenantFederation() (serv services.Service, err error) {
 				return bucket.NewClient(ctx, t.Cfg.BlocksStorage.Bucket, nil, "regex-resolver", util_log.Logger, prometheus.DefaultRegisterer)
 			}
 
-			regexResolver, err := tenantfederation.NewRegexResolver(t.Cfg.BlocksStorage.UsersScanner, prometheus.DefaultRegisterer, bucketClientFactory, t.Cfg.TenantFederation.UserSyncInterval, util_log.Logger)
+			regexResolver, err := tenantfederation.NewRegexResolver(t.Cfg.BlocksStorage.UsersScanner, t.Cfg.TenantFederation, prometheus.DefaultRegisterer, bucketClientFactory, util_log.Logger)
 			if err != nil {
 				return nil, fmt.Errorf("failed to initialize regex resolver: %v", err)
 			}

pkg/frontend/transport/handler.go

@@ -35,8 +35,6 @@ const (
 	// StatusClientClosedRequest is the status code for when a client request cancellation of a http request
 	StatusClientClosedRequest = 499
 	ServiceTimingHeaderName   = "Server-Timing"
-
-	errTooManyTenants = "too many tenants, max: %d, actual: %d"
 )
 
 var (
@@ -234,7 +232,7 @@ func (f *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 			if f.cfg.QueryStatsEnabled {
 				f.rejectedQueries.WithLabelValues(reasonTooManyTenants, source, userID).Inc()
 			}
-			http.Error(w, fmt.Errorf(errTooManyTenants, maxTenant, len(tenantIDs)).Error(), http.StatusBadRequest)
+			http.Error(w, fmt.Errorf(tenantfederation.ErrTooManyTenants, maxTenant, len(tenantIDs)).Error(), http.StatusBadRequest)
 			return
 		}
 	}

pkg/querier/tenantfederation/exemplar_merge_queryable_test.go

@@ -335,7 +335,8 @@ func Test_MergeExemplarQuerier_Select_WhenUseRegexResolver(t *testing.T) {
 	}
 
 	usersScannerConfig := cortex_tsdb.UsersScannerConfig{Strategy: cortex_tsdb.UserScanStrategyList}
-	regexResolver, err := NewRegexResolver(usersScannerConfig, reg, bucketClientFactory, time.Second, log.NewNopLogger())
+	tenantFederationConfig := Config{UserSyncInterval: time.Second}
+	regexResolver, err := NewRegexResolver(usersScannerConfig, tenantFederationConfig, reg, bucketClientFactory, log.NewNopLogger())
 	require.NoError(t, err)
 	tenant.WithDefaultResolver(regexResolver)
 	require.NoError(t, services.StartAndAwaitRunning(context.Background(), regexResolver))

pkg/querier/tenantfederation/merge_queryable_test.go

@@ -677,8 +677,8 @@ func TestMergeQueryable_Select(t *testing.T) {
 							}
 
 							usersScannerConfig := cortex_tsdb.UsersScannerConfig{Strategy: cortex_tsdb.UserScanStrategyList}
-
-							regexResolver, err := NewRegexResolver(usersScannerConfig, reg, bucketClientFactory, time.Second, log.NewNopLogger())
+							tenantFederationConfig := Config{UserSyncInterval: time.Second}
+							regexResolver, err := NewRegexResolver(usersScannerConfig, tenantFederationConfig, reg, bucketClientFactory, log.NewNopLogger())
 							require.NoError(t, err)
 
 							// set a regex tenant resolver
@@ -876,8 +876,8 @@ func TestMergeQueryable_LabelNames(t *testing.T) {
 						return bucketClient, nil
 					}
 					usersScannerConfig := cortex_tsdb.UsersScannerConfig{Strategy: cortex_tsdb.UserScanStrategyList}
-
-					regexResolver, err := NewRegexResolver(usersScannerConfig, reg, bucketClientFactory, time.Second, log.NewNopLogger())
+					tenantFederationConfig := Config{UserSyncInterval: time.Second}
+					regexResolver, err := NewRegexResolver(usersScannerConfig, tenantFederationConfig, reg, bucketClientFactory, log.NewNopLogger())
 					require.NoError(t, err)
 
 					// set a regex tenant resolver
@@ -1114,8 +1114,8 @@ func TestMergeQueryable_LabelValues(t *testing.T) {
 								return bucketClient, nil
 							}
 							usersScannerConfig := cortex_tsdb.UsersScannerConfig{Strategy: cortex_tsdb.UserScanStrategyList}
-
-							regexResolver, err := NewRegexResolver(usersScannerConfig, reg, bucketClientFactory, time.Second, log.NewNopLogger())
+							tenantFederationConfig := Config{UserSyncInterval: time.Second}
+							regexResolver, err := NewRegexResolver(usersScannerConfig, tenantFederationConfig, reg, bucketClientFactory, log.NewNopLogger())
 							require.NoError(t, err)
 
 							// set a regex tenant resolver

pkg/querier/tenantfederation/metadata_merge_querier_test.go

@@ -169,7 +169,8 @@ func Test_mergeMetadataQuerier_MetricsMetadata_WhenUseRegexResolver(t *testing.T
 	}
 
 	usersScannerConfig := cortex_tsdb.UsersScannerConfig{Strategy: cortex_tsdb.UserScanStrategyList}
-	regexResolver, err := NewRegexResolver(usersScannerConfig, reg, bucketClientFactory, time.Second, log.NewNopLogger())
+	tenantFederationConfig := Config{UserSyncInterval: time.Second}
+	regexResolver, err := NewRegexResolver(usersScannerConfig, tenantFederationConfig, reg, bucketClientFactory, log.NewNopLogger())
 	require.NoError(t, err)
 	tenant.WithDefaultResolver(regexResolver)
 	require.NoError(t, services.StartAndAwaitRunning(context.Background(), regexResolver))

pkg/querier/tenantfederation/regex_resolver.go

@@ -2,6 +2,8 @@ package tenantfederation
 
 import (
 	"context"
+	"fmt"
+	"net/http"
 	"sort"
 	"sync"
 	"time"
@@ -14,6 +16,7 @@ import (
 	"github.com/prometheus/prometheus/model/labels"
 	"github.com/thanos-io/objstore"
 	"github.com/thanos-io/thanos/pkg/extprom"
+	"github.com/weaveworks/common/httpgrpc"
 	"github.com/weaveworks/common/user"
 
 	"github.com/cortexproject/cortex/pkg/storage/tsdb"
@@ -24,6 +27,8 @@ import (
 
 var (
 	errInvalidRegex = errors.New("invalid regex present")
+
+	ErrTooManyTenants = "too many tenants, max: %d, actual: %d"
 )
 
 // RegexResolver resolves tenantIDs matched given regex.
@@ -32,6 +37,7 @@ type RegexResolver struct {
 
 	knownUsers       []string
 	userSyncInterval time.Duration
+	maxTenant        int
 	userScanner      users.Scanner
 	logger           log.Logger
 	sync.Mutex
@@ -42,7 +48,7 @@ type RegexResolver struct {
 	discoveredUsers prometheus.Gauge
 }
 
-func NewRegexResolver(cfg tsdb.UsersScannerConfig, reg prometheus.Registerer, bucketClientFactory func(ctx context.Context) (objstore.InstrumentedBucket, error), userSyncInterval time.Duration, logger log.Logger) (*RegexResolver, error) {
+func NewRegexResolver(cfg tsdb.UsersScannerConfig, tenantFederationCfg Config, reg prometheus.Registerer, bucketClientFactory func(ctx context.Context) (objstore.InstrumentedBucket, error), logger log.Logger) (*RegexResolver, error) {
 	bucketClient, err := bucketClientFactory(context.Background())
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to create the bucket client")
@@ -54,7 +60,8 @@ func NewRegexResolver(cfg tsdb.UsersScannerConfig, reg prometheus.Registerer, bu
 	}
 
 	r := &RegexResolver{
-		userSyncInterval: userSyncInterval,
+		userSyncInterval: tenantFederationCfg.UserSyncInterval,
+		maxTenant:        tenantFederationCfg.MaxTenant,
 		userScanner:      userScanner,
 		logger:           logger,
 	}
@@ -160,6 +167,10 @@ func (r *RegexResolver) getRegexMatchedOrgIds(orgID string) ([]string, error) {
 		return []string{"fake"}, nil
 	}
 
+	if r.maxTenant > 0 && len(matched) > r.maxTenant {
+		return nil, httpgrpc.Errorf(http.StatusBadRequest, "%s", fmt.Errorf(ErrTooManyTenants, r.maxTenant, len(matched)).Error())
+	}
+
 	return matched, nil
 }
 

pkg/querier/tenantfederation/regex_resolver_test.go

@@ -27,6 +27,7 @@ func Test_RegexResolver(t *testing.T) {
 		orgID           string
 		expectedErr     error
 		expectedOrgIDs  []string
+		maxTenants      int
 	}{
 		{
 			description:     "invalid regex",
@@ -64,6 +65,13 @@ func Test_RegexResolver(t *testing.T) {
 			orgID:           "user-1|user-2",
 			expectedOrgIDs:  []string{"user-1"},
 		},
+		{
+			description:     "adjust maxTenant",
+			existingTenants: []string{"user-1", "user-2", "user-3"},
+			orgID:           "user-.+",
+			maxTenants:      2,
+			expectedErr:     errors.New("too many tenants, max: 2, actual: 3"),
+		},
 	}
 
 	for _, tc := range tests {
@@ -82,7 +90,8 @@ func Test_RegexResolver(t *testing.T) {
 			}
 
 			usersScannerConfig := cortex_tsdb.UsersScannerConfig{Strategy: cortex_tsdb.UserScanStrategyList}
-			regexResolver, err := NewRegexResolver(usersScannerConfig, reg, bucketClientFactory, time.Second, log.NewNopLogger())
+			tenantFederationConfig := Config{UserSyncInterval: time.Second, MaxTenant: tc.maxTenants}
+			regexResolver, err := NewRegexResolver(usersScannerConfig, tenantFederationConfig, reg, bucketClientFactory, log.NewNopLogger())
 			require.NoError(t, err)
 			require.NoError(t, services.StartAndAwaitRunning(context.Background(), regexResolver))
 
@@ -97,7 +106,7 @@ func Test_RegexResolver(t *testing.T) {
 			orgIDs, err := regexResolver.TenantIDs(ctx)
 
 			if tc.expectedErr != nil {
-				require.Error(t, err)
+				require.Contains(t, err.Error(), tc.expectedErr.Error())
 			} else {
 				require.NoError(t, err)
 				require.Equal(t, tc.expectedOrgIDs, orgIDs)