Ruler: Add support for per-user external labels

Signed-off-by: Xiaochao Dong (@damnever) <the.xcdong@gmail.com>

Author: damnever

CHANGELOG.md

@@ -15,6 +15,7 @@
 * [FEATURE] Ruler: Minimize chances of missed rule group evaluations that can occur due to OOM kills, bad underlying nodes, or due to an unhealthy ruler that appears in the ring as healthy. This feature is enabled via `-ruler.enable-ha-evaluation` flag. #6129
 * [FEATURE] Store Gateway: Add an in-memory chunk cache. #6245
 * [FEATURE] Chunk Cache: Support multi level cache and add metrics. #6249
+* [FEATURE] Ruler: Add support for per-user external labels #6340
 * [ENHANCEMENT] S3 Bucket Client: Add a list objects version configs to configure list api object version. #6280
 * [ENHANCEMENT] OpenStack Swift: Add application credential configs for Openstack swift object storage backend. #6255
 * [ENHANCEMENT] Query Frontend: Add new query stats metrics `cortex_query_samples_scanned_total` and `cortex_query_peak_samples` to track scannedSamples and peakSample per user. #6228

docs/configuration/config-file-reference.md

@@ -3604,6 +3604,9 @@ query_rejection:
 
 # list of rule groups to disable
 [disabled_rule_groups: <list of DisabledRuleGroup> | default = []]
+
+# external labels for alerting rules
+[external_labels: <map of string (labelName) to string (labelValue)> | default = []]
 ```
 
 ### `memberlist_config`

pkg/cortex/modules.go

@@ -154,7 +154,8 @@ func (t *Cortex) initRuntimeConfig() (services.Service, error) {
 		// no need to initialize module if load path is empty
 		return nil, nil
 	}
-	t.Cfg.RuntimeConfig.Loader = loadRuntimeConfig
+	runtimeConfigLoader := runtimeConfigLoader{cfg: t.Cfg}
+	t.Cfg.RuntimeConfig.Loader = runtimeConfigLoader.load
 
 	// make sure to set default limits before we start loading configuration into memory
 	validation.SetDefaultLimitsForYAMLUnmarshalling(t.Cfg.LimitsConfig)
@@ -611,14 +612,14 @@ func (t *Cortex) initRuler() (serv services.Service, err error) {
 		}
 
 		managerFactory := ruler.DefaultTenantManagerFactory(t.Cfg.Ruler, t.Cfg.ExternalPusher, t.Cfg.ExternalQueryable, queryEngine, t.Overrides, metrics, prometheus.DefaultRegisterer)
-		manager, err = ruler.NewDefaultMultiTenantManager(t.Cfg.Ruler, managerFactory, metrics, prometheus.DefaultRegisterer, util_log.Logger)
+		manager, err = ruler.NewDefaultMultiTenantManager(t.Cfg.Ruler, t.Overrides, managerFactory, metrics, prometheus.DefaultRegisterer, util_log.Logger)
 	} else {
 		rulerRegisterer := prometheus.WrapRegistererWith(prometheus.Labels{"engine": "ruler"}, prometheus.DefaultRegisterer)
 		// TODO: Consider wrapping logger to differentiate from querier module logger
 		queryable, _, engine := querier.New(t.Cfg.Querier, t.Overrides, t.Distributor, t.StoreQueryables, rulerRegisterer, util_log.Logger)
 
 		managerFactory := ruler.DefaultTenantManagerFactory(t.Cfg.Ruler, t.Distributor, queryable, engine, t.Overrides, metrics, prometheus.DefaultRegisterer)
-		manager, err = ruler.NewDefaultMultiTenantManager(t.Cfg.Ruler, managerFactory, metrics, prometheus.DefaultRegisterer, util_log.Logger)
+		manager, err = ruler.NewDefaultMultiTenantManager(t.Cfg.Ruler, t.Overrides, managerFactory, metrics, prometheus.DefaultRegisterer, util_log.Logger)
 	}
 
 	if err != nil {

pkg/cortex/runtime_config.go

@@ -58,7 +58,11 @@ func (l *runtimeConfigTenantLimits) AllByUserID() map[string]*validation.Limits
 	return nil
 }
 
-func loadRuntimeConfig(r io.Reader) (interface{}, error) {
+type runtimeConfigLoader struct {
+	cfg Config
+}
+
+func (l runtimeConfigLoader) load(r io.Reader) (interface{}, error) {
 	var overrides = &RuntimeConfigValues{}
 
 	decoder := yaml.NewDecoder(r)
@@ -74,6 +78,12 @@ func loadRuntimeConfig(r io.Reader) (interface{}, error) {
 		return nil, errMultipleDocuments
 	}
 
+	for _, ul := range overrides.TenantLimits {
+		if err := ul.Validate(l.cfg.Distributor.ShardByAllLabels); err != nil {
+			return nil, err
+		}
+	}
+
 	return overrides, nil
 }
 

pkg/cortex/runtime_config_test.go

@@ -7,6 +7,7 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
+	"github.com/cortexproject/cortex/pkg/distributor"
 	"github.com/cortexproject/cortex/pkg/util/validation"
 )
 
@@ -28,7 +29,8 @@ overrides:
   '1235': *id001
   '1236': *id001
 `)
-	runtimeCfg, err := loadRuntimeConfig(yamlFile)
+	loader := runtimeConfigLoader{cfg: Config{Distributor: distributor.Config{ShardByAllLabels: true}}}
+	runtimeCfg, err := loader.load(yamlFile)
 	require.NoError(t, err)
 
 	limits := validation.Limits{
@@ -51,7 +53,7 @@ func TestLoadRuntimeConfig_ShouldLoadEmptyFile(t *testing.T) {
 	yamlFile := strings.NewReader(`
 # This is an empty YAML.
 `)
-	actual, err := loadRuntimeConfig(yamlFile)
+	actual, err := runtimeConfigLoader{}.load(yamlFile)
 	require.NoError(t, err)
 	assert.Equal(t, &RuntimeConfigValues{}, actual)
 }
@@ -60,7 +62,7 @@ func TestLoadRuntimeConfig_MissingPointerFieldsAreNil(t *testing.T) {
 	yamlFile := strings.NewReader(`
 # This is an empty YAML.
 `)
-	actual, err := loadRuntimeConfig(yamlFile)
+	actual, err := runtimeConfigLoader{}.load(yamlFile)
 	require.NoError(t, err)
 
 	actualCfg, ok := actual.(*RuntimeConfigValues)
@@ -102,7 +104,7 @@ overrides:
 	}
 
 	for _, tc := range cases {
-		actual, err := loadRuntimeConfig(strings.NewReader(tc))
+		actual, err := runtimeConfigLoader{}.load(strings.NewReader(tc))
 		assert.Equal(t, errMultipleDocuments, err)
 		assert.Nil(t, actual)
 	}

pkg/ruler/compat.go

@@ -153,6 +153,7 @@ type RulesLimits interface {
 	RulerMaxRulesPerRuleGroup(userID string) int
 	RulerQueryOffset(userID string) time.Duration
 	DisabledRuleGroups(userID string) validation.DisabledRuleGroups
+	ExternalLabels(userID string) labels.Labels
 }
 
 // EngineQueryFunc returns a new engine query function validating max queryLength.

pkg/ruler/external_labels.go

@@ -0,0 +1,68 @@
+package ruler
+
+import (
+	"sync"
+
+	"github.com/prometheus/prometheus/model/labels"
+)
+
+// userExternalLabels checks and merges per-user external labels with global external labels.
+type userExternalLabels struct {
+	global  labels.Labels
+	limits  RulesLimits
+	builder *labels.Builder
+
+	mtx   sync.Mutex
+	users map[string]labels.Labels
+}
+
+func newUserExternalLabels(global labels.Labels, limits RulesLimits) *userExternalLabels {
+	return &userExternalLabels{
+		global:  global,
+		limits:  limits,
+		builder: labels.NewBuilder(nil),
+
+		mtx:   sync.Mutex{},
+		users: map[string]labels.Labels{},
+	}
+}
+
+func (e *userExternalLabels) get(userID string) (labels.Labels, bool) {
+	e.mtx.Lock()
+	defer e.mtx.Unlock()
+	lset, ok := e.users[userID]
+	return lset, ok
+}
+
+func (e *userExternalLabels) update(userID string) (labels.Labels, bool) {
+	lset := e.limits.ExternalLabels(userID)
+
+	e.mtx.Lock()
+	defer e.mtx.Unlock()
+
+	e.builder.Reset(e.global)
+	for _, l := range lset {
+		e.builder.Set(l.Name, l.Value)
+	}
+	lset = e.builder.Labels()
+
+	if !labels.Equal(e.users[userID], lset) {
+		e.users[userID] = lset
+		return lset, true
+	}
+	return lset, false
+}
+
+func (e *userExternalLabels) remove(user string) {
+	e.mtx.Lock()
+	defer e.mtx.Unlock()
+	delete(e.users, user)
+}
+
+func (e *userExternalLabels) cleanup() {
+	e.mtx.Lock()
+	defer e.mtx.Unlock()
+	for user := range e.users {
+		delete(e.users, user)
+	}
+}

pkg/ruler/manager.go

@@ -17,6 +17,7 @@ import (
 	"github.com/prometheus/client_golang/prometheus/promauto"
 	"github.com/prometheus/prometheus/config"
 	"github.com/prometheus/prometheus/discovery"
+	"github.com/prometheus/prometheus/model/labels"
 	"github.com/prometheus/prometheus/model/rulefmt"
 	"github.com/prometheus/prometheus/notifier"
 	promRules "github.com/prometheus/prometheus/rules"
@@ -47,6 +48,9 @@ type DefaultMultiTenantManager struct {
 	notifiers                 map[string]*rulerNotifier
 	notifiersDiscoveryMetrics map[string]discovery.DiscovererMetrics
 
+	// Per-user externalLabels.
+	userExternalLabels *userExternalLabels
+
 	// rules backup
 	rulesBackupManager *rulesBackupManager
 
@@ -62,7 +66,7 @@ type DefaultMultiTenantManager struct {
 	syncRuleMtx  sync.Mutex
 }
 
-func NewDefaultMultiTenantManager(cfg Config, managerFactory ManagerFactory, evalMetrics *RuleEvalMetrics, reg prometheus.Registerer, logger log.Logger) (*DefaultMultiTenantManager, error) {
+func NewDefaultMultiTenantManager(cfg Config, limits RulesLimits, managerFactory ManagerFactory, evalMetrics *RuleEvalMetrics, reg prometheus.Registerer, logger log.Logger) (*DefaultMultiTenantManager, error) {
 	ncfg, err := buildNotifierConfig(&cfg)
 	if err != nil {
 		return nil, err
@@ -92,6 +96,7 @@ func NewDefaultMultiTenantManager(cfg Config, managerFactory ManagerFactory, eva
 		frontendPool:              newFrontendPool(cfg, logger, reg),
 		ruleEvalMetrics:           evalMetrics,
 		notifiers:                 map[string]*rulerNotifier{},
+		userExternalLabels:        newUserExternalLabels(cfg.ExternalLabels, limits),
 		notifiersDiscoveryMetrics: notifiersDiscoveryMetrics,
 		mapper:                    newMapper(cfg.RulePath, logger),
 		userManagers:              map[string]RulesManager{},
@@ -146,6 +151,7 @@ func (r *DefaultMultiTenantManager) SyncRuleGroups(ctx context.Context, ruleGrou
 
 			r.removeNotifier(userID)
 			r.mapper.cleanupUser(userID)
+			r.userExternalLabels.remove(userID)
 			r.lastReloadSuccessful.DeleteLabelValues(userID)
 			r.lastReloadSuccessfulTimestamp.DeleteLabelValues(userID)
 			r.configUpdatesTotal.DeleteLabelValues(userID)
@@ -183,12 +189,13 @@ func (r *DefaultMultiTenantManager) BackUpRuleGroups(ctx context.Context, ruleGr
 func (r *DefaultMultiTenantManager) syncRulesToManager(ctx context.Context, user string, groups rulespb.RuleGroupList) {
 	// Map the files to disk and return the file names to be passed to the users manager if they
 	// have been updated
-	update, files, err := r.mapper.MapRules(user, groups.Formatted())
+	rulesUpdated, files, err := r.mapper.MapRules(user, groups.Formatted())
 	if err != nil {
 		r.lastReloadSuccessful.WithLabelValues(user).Set(0)
 		level.Error(r.logger).Log("msg", "unable to map rule files", "user", user, "err", err)
 		return
 	}
+	externalLabels, externalLabelsUpdated := r.userExternalLabels.update(user)
 
 	existing := true
 	manager := r.getRulesManager(user, ctx)
@@ -201,19 +208,26 @@ func (r *DefaultMultiTenantManager) syncRulesToManager(ctx context.Context, user
 		return
 	}
 
-	if !existing || update {
+	if !existing || rulesUpdated || externalLabelsUpdated {
 		level.Debug(r.logger).Log("msg", "updating rules", "user", user)
 		r.configUpdatesTotal.WithLabelValues(user).Inc()
-		if update && existing {
+		if rulesUpdated && existing {
 			r.updateRuleCache(user, manager.RuleGroups())
 		}
-		err = manager.Update(r.cfg.EvaluationInterval, files, r.cfg.ExternalLabels, r.cfg.ExternalURL.String(), ruleGroupIterationFunc)
+		err = manager.Update(r.cfg.EvaluationInterval, files, externalLabels, r.cfg.ExternalURL.String(), ruleGroupIterationFunc)
 		r.deleteRuleCache(user)
 		if err != nil {
 			r.lastReloadSuccessful.WithLabelValues(user).Set(0)
 			level.Error(r.logger).Log("msg", "unable to update rule manager", "user", user, "err", err)
 			return
 		}
+		if externalLabelsUpdated {
+			if err = r.notifierApplyExternalLabels(user, externalLabels); err != nil {
+				r.lastReloadSuccessful.WithLabelValues(user).Set(0)
+				level.Error(r.logger).Log("msg", "unable to update notifier", "user", user, "err", err)
+				return
+			}
+		}
 
 		r.lastReloadSuccessful.WithLabelValues(user).Set(1)
 		r.lastReloadSuccessfulTimestamp.WithLabelValues(user).SetToCurrentTime()
@@ -348,6 +362,19 @@ func (r *DefaultMultiTenantManager) getOrCreateNotifier(userID string, userManag
 	return n.notifier, nil
 }
 
+func (r *DefaultMultiTenantManager) notifierApplyExternalLabels(userID string, externalLabels labels.Labels) error {
+	r.notifiersMtx.Lock()
+	defer r.notifiersMtx.Unlock()
+
+	n, ok := r.notifiers[userID]
+	if !ok {
+		return fmt.Errorf("notifier not found")
+	}
+	cfg := *r.notifierCfg // Copy it
+	cfg.GlobalConfig.ExternalLabels = externalLabels
+	return n.applyConfig(&cfg)
+}
+
 func (r *DefaultMultiTenantManager) getCachedRules(userID string) ([]*promRules.Group, bool) {
 	r.ruleCacheMtx.RLock()
 	defer r.ruleCacheMtx.RUnlock()
@@ -402,6 +429,7 @@ func (r *DefaultMultiTenantManager) Stop() {
 
 	// cleanup user rules directories
 	r.mapper.cleanup()
+	r.userExternalLabels.cleanup()
 }
 
 func (*DefaultMultiTenantManager) ValidateRuleGroup(g rulefmt.RuleGroup) []error {

pkg/ruler/manager_test.go

@@ -29,8 +29,9 @@ func TestSyncRuleGroups(t *testing.T) {
 	}
 
 	ruleManagerFactory := RuleManagerFactory(nil, waitDurations)
+	limits := ruleLimits{externalLabels: labels.FromStrings("from", "cortex")}
 
-	m, err := NewDefaultMultiTenantManager(Config{RulePath: dir}, ruleManagerFactory, nil, nil, log.NewNopLogger())
+	m, err := NewDefaultMultiTenantManager(Config{RulePath: dir}, limits, ruleManagerFactory, nil, nil, log.NewNopLogger())
 	require.NoError(t, err)
 
 	const user = "testUser"
@@ -61,6 +62,9 @@ func TestSyncRuleGroups(t *testing.T) {
 		require.NoError(t, err)
 		require.Equal(t, []string{user}, users)
 		require.True(t, ok)
+		lset, ok := m.userExternalLabels.get(user)
+		require.True(t, ok)
+		require.Equal(t, limits.externalLabels, lset)
 	}
 
 	// Passing empty map / nil stops all managers.
@@ -79,6 +83,8 @@ func TestSyncRuleGroups(t *testing.T) {
 		require.NoError(t, err)
 		require.Equal(t, []string(nil), users)
 		require.False(t, ok)
+		_, ok = m.userExternalLabels.get(user)
+		require.False(t, ok)
 	}
 
 	// Resync same rules as before. Previously this didn't restart the manager.
@@ -154,7 +160,7 @@ func TestSlowRuleGroupSyncDoesNotSlowdownListRules(t *testing.T) {
 	}
 
 	ruleManagerFactory := RuleManagerFactory(groupsToReturn, waitDurations)
-	m, err := NewDefaultMultiTenantManager(Config{RulePath: dir}, ruleManagerFactory, nil, prometheus.NewRegistry(), log.NewNopLogger())
+	m, err := NewDefaultMultiTenantManager(Config{RulePath: dir}, ruleLimits{}, ruleManagerFactory, nil, prometheus.NewRegistry(), log.NewNopLogger())
 	require.NoError(t, err)
 
 	m.SyncRuleGroups(context.Background(), userRules)
@@ -217,7 +223,7 @@ func TestSyncRuleGroupsCleanUpPerUserMetrics(t *testing.T) {
 
 	ruleManagerFactory := RuleManagerFactory(nil, waitDurations)
 
-	m, err := NewDefaultMultiTenantManager(Config{RulePath: dir}, ruleManagerFactory, evalMetrics, reg, log.NewNopLogger())
+	m, err := NewDefaultMultiTenantManager(Config{RulePath: dir}, ruleLimits{}, ruleManagerFactory, evalMetrics, reg, log.NewNopLogger())
 	require.NoError(t, err)
 
 	const user = "testUser"
@@ -265,7 +271,7 @@ func TestBackupRules(t *testing.T) {
 	ruleManagerFactory := RuleManagerFactory(nil, waitDurations)
 	config := Config{RulePath: dir}
 	config.Ring.ReplicationFactor = 3
-	m, err := NewDefaultMultiTenantManager(config, ruleManagerFactory, evalMetrics, reg, log.NewNopLogger())
+	m, err := NewDefaultMultiTenantManager(config, ruleLimits{}, ruleManagerFactory, evalMetrics, reg, log.NewNopLogger())
 	require.NoError(t, err)
 
 	const user1 = "testUser"