Add max tenant config to tenant federation

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

Author: SungJin1212

CHANGELOG.md

@@ -22,6 +22,7 @@
 * [FEATURE] Query Frontend: Support a metadata federated query when `-tenant-federation.enabled=true`. #6461
 * [FEATURE] Query Frontend: Support an exemplar federated query when `-tenant-federation.enabled=true`. #6455
 * [FEATURE] Ingester/StoreGateway: Add support for cache regex query matchers via `-ingester.matchers-cache-max-items` and `-blocks-storage.bucket-store.matchers-cache-max-items`. #6477 #6491
+* [ENHANCEMENT] Query Frontend: Add a flag `-tenant-federation.max-tenant` to limit the number of tenants for federated query. #6493
 * [ENHANCEMENT] Querier: Add a `-tenant-federation.max-concurrent` flags to configure the number of worker processing federated query and add a `cortex_querier_federated_tenants_per_query` histogram to track the number of tenants per query. #6449
 * [ENHANCEMENT] Query Frontend: Add a number of series in the query response to the query stat log. #6423
 * [ENHANCEMENT] Store Gateway: Add a hedged request to reduce the tail latency. #6388

docs/configuration/config-file-reference.md

@@ -161,6 +161,10 @@ tenant_federation:
   # CLI flag: -tenant-federation.max-concurrent
   [max_concurrent: <int> | default = 16]
 
+  # A maximum number of tenants to query at once. 0 means no limit.
+  # CLI flag: -tenant-federation.max-tenant
+  [max_tenant: <int> | default = 0]
+
 # The ruler_config configures the Cortex ruler.
 [ruler: <ruler_config>]
 

pkg/api/api.go

@@ -32,6 +32,7 @@ import (
 	"github.com/cortexproject/cortex/pkg/ingester/client"
 	"github.com/cortexproject/cortex/pkg/purger"
 	"github.com/cortexproject/cortex/pkg/querier"
+	"github.com/cortexproject/cortex/pkg/querier/tenantfederation"
 	"github.com/cortexproject/cortex/pkg/ring"
 	"github.com/cortexproject/cortex/pkg/ruler"
 	"github.com/cortexproject/cortex/pkg/scheduler"
@@ -136,7 +137,7 @@ type API struct {
 	corsOrigin           *regexp.Regexp
 }
 
-func New(cfg Config, serverCfg server.Config, s *server.Server, logger log.Logger) (*API, error) {
+func New(cfg Config, tenantFederationCfg tenantfederation.Config, serverCfg server.Config, s *server.Server, logger log.Logger) (*API, error) {
 	// Ensure the encoded path is used. Required for the rules API
 	s.HTTP.UseEncodedPath()
 
@@ -169,6 +170,11 @@ func New(cfg Config, serverCfg server.Config, s *server.Server, logger log.Logge
 	if cfg.HTTPAuthMiddleware == nil {
 		api.AuthMiddleware = middleware.AuthenticateUser
 	}
+
+	if tenantFederationCfg.Enabled {
+		api.AuthMiddleware = middleware.Merge(api.AuthMiddleware, tenantFederationMiddleWare(tenantFederationCfg.MaxTenant))
+	}
+
 	if len(cfg.HTTPRequestHeadersToLog) > 0 {
 		api.HTTPHeaderMiddleware = &HTTPHeaderMiddleware{TargetHeaders: cfg.HTTPRequestHeadersToLog}
 	}

pkg/api/api_test.go

@@ -12,6 +12,8 @@ import (
 	"github.com/prometheus/prometheus/model/labels"
 	"github.com/stretchr/testify/require"
 	"github.com/weaveworks/common/server"
+
+	"github.com/cortexproject/cortex/pkg/querier/tenantfederation"
 )
 
 const (
@@ -34,7 +36,9 @@ func TestNewApiWithoutSourceIPExtractor(t *testing.T) {
 	server, err := server.New(serverCfg)
 	require.NoError(t, err)
 
-	api, err := New(cfg, serverCfg, server, &FakeLogger{})
+	tenantFederationCfg := tenantfederation.Config{}
+
+	api, err := New(cfg, tenantFederationCfg, serverCfg, server, &FakeLogger{})
 	require.NoError(t, err)
 	require.Nil(t, api.sourceIPs)
 }
@@ -49,7 +53,9 @@ func TestNewApiWithSourceIPExtractor(t *testing.T) {
 	server, err := server.New(serverCfg)
 	require.NoError(t, err)
 
-	api, err := New(cfg, serverCfg, server, &FakeLogger{})
+	tenantFederationCfg := tenantfederation.Config{}
+
+	api, err := New(cfg, tenantFederationCfg, serverCfg, server, &FakeLogger{})
 	require.NoError(t, err)
 	require.NotNil(t, api.sourceIPs)
 }
@@ -67,7 +73,9 @@ func TestNewApiWithInvalidSourceIPExtractor(t *testing.T) {
 		MetricsNamespace:   "with_invalid_source_ip_extractor",
 	}
 
-	api, err := New(cfg, serverCfg, &s, &FakeLogger{})
+	tenantFederationCfg := tenantfederation.Config{}
+
+	api, err := New(cfg, tenantFederationCfg, serverCfg, &s, &FakeLogger{})
 	require.Error(t, err)
 	require.Nil(t, api)
 }
@@ -83,7 +91,9 @@ func TestNewApiWithHeaderLogging(t *testing.T) {
 	server, err := server.New(serverCfg)
 	require.NoError(t, err)
 
-	api, err := New(cfg, serverCfg, server, &FakeLogger{})
+	tenantFederationCfg := tenantfederation.Config{}
+
+	api, err := New(cfg, tenantFederationCfg, serverCfg, server, &FakeLogger{})
 	require.NoError(t, err)
 	require.NotNil(t, api.HTTPHeaderMiddleware)
 
@@ -100,7 +110,9 @@ func TestNewApiWithoutHeaderLogging(t *testing.T) {
 	server, err := server.New(serverCfg)
 	require.NoError(t, err)
 
-	api, err := New(cfg, serverCfg, server, &FakeLogger{})
+	tenantFederationCfg := tenantfederation.Config{}
+
+	api, err := New(cfg, tenantFederationCfg, serverCfg, server, &FakeLogger{})
 	require.NoError(t, err)
 	require.Nil(t, api.HTTPHeaderMiddleware)
 
@@ -157,7 +169,10 @@ func Benchmark_Compression(b *testing.B) {
 
 			server, err := server.New(serverCfg)
 			require.NoError(b, err)
-			api, err := New(cfg, serverCfg, server, &FakeLogger{})
+
+			tenantFederationCfg := tenantfederation.Config{}
+
+			api, err := New(cfg, tenantFederationCfg, serverCfg, server, &FakeLogger{})
 			require.NoError(b, err)
 
 			labels := labels.ScratchBuilder{}

pkg/api/middlewares.go

@@ -2,11 +2,19 @@ package api
 
 import (
 	"context"
+	"fmt"
 	"net/http"
 
+	"github.com/weaveworks/common/middleware"
+
+	"github.com/cortexproject/cortex/pkg/tenant"
 	util_log "github.com/cortexproject/cortex/pkg/util/log"
 )
 
+const (
+	errTooManyTenants = "too many tenants, max: %d, actual: %d"
+)
+
 // HTTPHeaderMiddleware adds specified HTTPHeaders to the request context
 type HTTPHeaderMiddleware struct {
 	TargetHeaders []string
@@ -38,3 +46,24 @@ func (h HTTPHeaderMiddleware) Wrap(next http.Handler) http.Handler {
 		next.ServeHTTP(w, r.WithContext(ctx))
 	})
 }
+
+func tenantFederationMiddleWare(maxTenant int) middleware.Interface {
+	return middleware.Func(func(next http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			ctx := r.Context()
+
+			ids, err := tenant.TenantIDs(ctx)
+			if err != nil {
+				http.Error(w, err.Error(), http.StatusUnauthorized)
+				return
+			}
+
+			if maxTenant > 0 && len(ids) > maxTenant {
+				http.Error(w, fmt.Errorf(errTooManyTenants, maxTenant, len(ids)).Error(), http.StatusBadRequest)
+				return
+			}
+
+			next.ServeHTTP(w, r.WithContext(ctx))
+		})
+	})
+}

pkg/api/middlewares_test.go

@@ -2,11 +2,15 @@ package api
 
 import (
 	"context"
+	"io"
 	"net/http"
+	"net/http/httptest"
 	"testing"
 
 	"github.com/stretchr/testify/require"
+	"github.com/weaveworks/common/middleware"
 
+	"github.com/cortexproject/cortex/pkg/tenant"
 	util_log "github.com/cortexproject/cortex/pkg/util/log"
 )
 
@@ -72,3 +76,71 @@ func TestExistingHeaderInContextIsNotOverridden(t *testing.T) {
 	require.Equal(t, contentsMap, util_log.HeaderMapFromContext(ctx))
 
 }
+
+func TestTenantFederationMiddleWare(t *testing.T) {
+	// set a multi tenant resolver
+	tenant.WithDefaultResolver(tenant.NewMultiResolver())
+
+	tests := []struct {
+		name               string
+		maxTenant          int
+		orgId              string
+		expectedStatusCode int
+		expectedErrMsg     string
+	}{
+		{
+			name:               "less than max tenant",
+			maxTenant:          3,
+			orgId:              "org1|org2",
+			expectedStatusCode: http.StatusOK,
+		},
+		{
+			name:               "equal to max tenant",
+			maxTenant:          2,
+			orgId:              "org1|org2",
+			expectedStatusCode: http.StatusOK,
+		},
+		{
+			name:               "exceeds max tenant",
+			maxTenant:          2,
+			orgId:              "org1|org2|org3",
+			expectedStatusCode: http.StatusBadRequest,
+			expectedErrMsg:     "too many tenants, max: 2, actual: 3",
+		},
+		{
+			name:               "no org Id",
+			maxTenant:          0,
+			orgId:              "",
+			expectedStatusCode: http.StatusUnauthorized,
+			expectedErrMsg:     "no org id",
+		},
+		{
+			name:               "no limit",
+			maxTenant:          0,
+			orgId:              "org1|org2|org3",
+			expectedStatusCode: http.StatusOK,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})
+
+			middlerwares := middleware.Merge(middleware.AuthenticateUser, tenantFederationMiddleWare(test.maxTenant)).Wrap(handler)
+
+			req := httptest.NewRequest(http.MethodGet, "/", nil)
+			req.Header.Set("X-Scope-OrgId", test.orgId)
+			resp := httptest.NewRecorder()
+
+			middlerwares.ServeHTTP(resp, req)
+
+			body, err := io.ReadAll(resp.Body)
+			require.NoError(t, err)
+			require.Equal(t, test.expectedStatusCode, resp.Code)
+
+			if test.expectedErrMsg != "" {
+				require.Contains(t, string(body), test.expectedErrMsg)
+			}
+		})
+	}
+}

pkg/cortex/modules.go

@@ -100,7 +100,7 @@ func (t *Cortex) initAPI() (services.Service, error) {
 	t.Cfg.API.ServerPrefix = t.Cfg.Server.PathPrefix
 	t.Cfg.API.LegacyHTTPPrefix = t.Cfg.HTTPPrefix
 
-	a, err := api.New(t.Cfg.API, t.Cfg.Server, t.Server, util_log.Logger)
+	a, err := api.New(t.Cfg.API, t.Cfg.TenantFederation, t.Cfg.Server, t.Server, util_log.Logger)
 	if err != nil {
 		return nil, err
 	}

pkg/querier/tenantfederation/tenant_federation.go

@@ -9,9 +9,12 @@ type Config struct {
 	Enabled bool `yaml:"enabled"`
 	// MaxConcurrent The number of workers used for processing federated query.
 	MaxConcurrent int `yaml:"max_concurrent"`
+	// MaxTenant A maximum number of tenants to query at once.
+	MaxTenant int `yaml:"max_tenant"`
 }
 
 func (cfg *Config) RegisterFlags(f *flag.FlagSet) {
 	f.BoolVar(&cfg.Enabled, "tenant-federation.enabled", false, "If enabled on all Cortex services, queries can be federated across multiple tenants. The tenant IDs involved need to be specified separated by a `|` character in the `X-Scope-OrgID` header (experimental).")
 	f.IntVar(&cfg.MaxConcurrent, "tenant-federation.max-concurrent", defaultMaxConcurrency, "The number of workers used to process each federated query.")
+	f.IntVar(&cfg.MaxTenant, "tenant-federation.max-tenant", 0, "A maximum number of tenants to query at once. 0 means no limit.")
 }