Restructure mutex such that manager is not holding one mutex for the entirety of sync rules

Signed-off-by: Anand Rajagopal <anrajag@amazon.com>

Author: rajagopalanand

pkg/ruler/manager.go

@@ -35,7 +35,7 @@ type DefaultMultiTenantManager struct {
 
 	// Structs for holding per-user Prometheus rules Managers
 	// and a corresponding metrics struct
-	userManagerMtx     sync.Mutex
+	userManagerMtx     sync.RWMutex
 	userManagers       map[string]RulesManager
 	userManagerMetrics *ManagerMetrics
 
@@ -50,6 +50,10 @@ type DefaultMultiTenantManager struct {
 	configUpdatesTotal            *prometheus.CounterVec
 	registry                      prometheus.Registerer
 	logger                        log.Logger
+
+	ruleCache    map[string][]*promRules.Group
+	ruleCacheMtx sync.RWMutex
+	syncRuleMtx  sync.Mutex
 }
 
 func NewDefaultMultiTenantManager(cfg Config, managerFactory ManagerFactory, evalMetrics *RuleEvalMetrics, reg prometheus.Registerer, logger log.Logger) (*DefaultMultiTenantManager, error) {
@@ -85,6 +89,7 @@ func NewDefaultMultiTenantManager(cfg Config, managerFactory ManagerFactory, eva
 		mapper:                    newMapper(cfg.RulePath, logger),
 		userManagers:              map[string]RulesManager{},
 		userManagerMetrics:        userManagerMetrics,
+		ruleCache:                 map[string][]*promRules.Group{},
 		managersTotal: promauto.With(reg).NewGauge(prometheus.GaugeOpts{
 			Namespace: "cortex",
 			Name:      "ruler_managers_total",
@@ -111,15 +116,17 @@ func NewDefaultMultiTenantManager(cfg Config, managerFactory ManagerFactory, eva
 }
 
 func (r *DefaultMultiTenantManager) SyncRuleGroups(ctx context.Context, ruleGroups map[string]rulespb.RuleGroupList) {
-	// A lock is taken to ensure if this function is called concurrently, then each call
-	// returns after the call map files and check for updates
-	r.userManagerMtx.Lock()
-	defer r.userManagerMtx.Unlock()
+	// this is a safety lock to ensure this method is executed sequentially
+	r.syncRuleMtx.Lock()
+	defer r.syncRuleMtx.Unlock()
 
 	for userID, ruleGroup := range ruleGroups {
 		r.syncRulesToManager(ctx, userID, ruleGroup)
 	}
 
+	r.userManagerMtx.Lock()
+	defer r.userManagerMtx.Unlock()
+
 	// Check for deleted users and remove them
 	for userID, mngr := range r.userManagers {
 		if _, exists := ruleGroups[userID]; !exists {
@@ -142,6 +149,18 @@ func (r *DefaultMultiTenantManager) SyncRuleGroups(ctx context.Context, ruleGrou
 	r.managersTotal.Set(float64(len(r.userManagers)))
 }
 
+func (r *DefaultMultiTenantManager) updateRuleCache(user string, rules []*promRules.Group) {
+	r.ruleCacheMtx.Lock()
+	defer r.ruleCacheMtx.Unlock()
+	r.ruleCache[user] = rules
+}
+
+func (r *DefaultMultiTenantManager) deleteRuleCache(user string) {
+	r.ruleCacheMtx.Lock()
+	defer r.ruleCacheMtx.Unlock()
+	delete(r.ruleCache, user)
+}
+
 // syncRulesToManager maps the rule files to disk, detects any changes and will create/update the
 // the users Prometheus Rules Manager.
 func (r *DefaultMultiTenantManager) syncRulesToManager(ctx context.Context, user string, groups rulespb.RuleGroupList) {
@@ -154,25 +173,20 @@ func (r *DefaultMultiTenantManager) syncRulesToManager(ctx context.Context, user
 		return
 	}
 
-	manager, exists := r.userManagers[user]
-	if !exists || update {
+	manager, existing := r.getRulesManager(user, ctx)
+
+	if manager == nil {
+		return
+	}
+
+	if !existing || update {
 		level.Debug(r.logger).Log("msg", "updating rules", "user", user)
 		r.configUpdatesTotal.WithLabelValues(user).Inc()
-		if !exists {
-			level.Debug(r.logger).Log("msg", "creating rule manager for user", "user", user)
-			manager, err = r.newManager(ctx, user)
-			if err != nil {
-				r.lastReloadSuccessful.WithLabelValues(user).Set(0)
-				level.Error(r.logger).Log("msg", "unable to create rule manager", "user", user, "err", err)
-				return
-			}
-			// manager.Run() starts running the manager and blocks until Stop() is called.
-			// Hence run it as another goroutine.
-			go manager.Run()
-			r.userManagers[user] = manager
+		if update && existing {
+			r.updateRuleCache(user, manager.RuleGroups())
 		}
-
 		err = manager.Update(r.cfg.EvaluationInterval, files, r.cfg.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)
@@ -184,6 +198,29 @@ func (r *DefaultMultiTenantManager) syncRulesToManager(ctx context.Context, user
 	}
 }
 
+func (r *DefaultMultiTenantManager) getRulesManager(user string, ctx context.Context) (RulesManager, bool) {
+	r.userManagerMtx.RLock()
+	manager, exists := r.userManagers[user]
+	r.userManagerMtx.RUnlock()
+	if exists {
+		return manager, true
+	}
+	r.userManagerMtx.Lock()
+	defer r.userManagerMtx.Unlock()
+
+	manager, err := r.newManager(ctx, user)
+	if err != nil {
+		r.lastReloadSuccessful.WithLabelValues(user).Set(0)
+		level.Error(r.logger).Log("msg", "unable to create rule manager", "user", user, "err", err)
+		return nil, false
+	}
+	// manager.Run() starts running the manager and blocks until Stop() is called.
+	// Hence run it as another goroutine.
+	go manager.Run()
+	r.userManagers[user] = manager
+	return manager, false
+}
+
 func ruleGroupIterationFunc(ctx context.Context, g *promRules.Group, evalTimestamp time.Time) {
 	logMessage := []interface{}{
 		"msg", "evaluating rule group",
@@ -269,13 +306,25 @@ func (r *DefaultMultiTenantManager) getOrCreateNotifier(userID string, userManag
 	return n.notifier, nil
 }
 
+func (r *DefaultMultiTenantManager) getCachedRules(userID string) ([]*promRules.Group, bool) {
+	r.ruleCacheMtx.RLock()
+	defer r.ruleCacheMtx.RUnlock()
+	groups, exists := r.ruleCache[userID]
+	return groups, exists
+}
+
 func (r *DefaultMultiTenantManager) GetRules(userID string) []*promRules.Group {
 	var groups []*promRules.Group
-	r.userManagerMtx.Lock()
-	if mngr, exists := r.userManagers[userID]; exists {
+	groups, cached := r.getCachedRules(userID)
+	if cached {
+		return groups
+	}
+	r.userManagerMtx.RLock()
+	mngr, exists := r.userManagers[userID]
+	r.userManagerMtx.RUnlock()
+	if exists {
 		groups = mngr.RuleGroups()
 	}
-	r.userManagerMtx.Unlock()
 	return groups
 }
 

pkg/ruler/manager_test.go

@@ -2,6 +2,7 @@ package ruler
 
 import (
 	"context"
+	"sync"
 	"testing"
 	"time"
 
@@ -21,7 +22,14 @@ import (
 func TestSyncRuleGroups(t *testing.T) {
 	dir := t.TempDir()
 
-	m, err := NewDefaultMultiTenantManager(Config{RulePath: dir}, factory, nil, nil, log.NewNopLogger())
+	waitDurations := []time.Duration{
+		1 * time.Millisecond,
+		1 * time.Millisecond,
+	}
+
+	ruleManagerFactory := RuleManagerFactory(nil, waitDurations)
+
+	m, err := NewDefaultMultiTenantManager(Config{RulePath: dir}, ruleManagerFactory, nil, nil, log.NewNopLogger())
 	require.NoError(t, err)
 
 	const user = "testUser"
@@ -97,6 +105,107 @@ func TestSyncRuleGroups(t *testing.T) {
 	})
 }
 
+func TestSlowRuleGroupSyncDoesNotSlowdownListRules(t *testing.T) {
+	dir := t.TempDir()
+	const user = "testUser"
+	userRules := map[string]rulespb.RuleGroupList{
+		user: {
+			&rulespb.RuleGroupDesc{
+				Name:      "group1",
+				Namespace: "ns",
+				Interval:  1 * time.Minute,
+				User:      user,
+			},
+		},
+	}
+
+	groupsToReturn := [][]*promRules.Group{
+		{
+			promRules.NewGroup(promRules.GroupOptions{
+				Name:     "group1",
+				File:     "ns",
+				Interval: 60,
+				Limit:    0,
+				Opts:     &promRules.ManagerOptions{},
+			}),
+		},
+		{
+			promRules.NewGroup(promRules.GroupOptions{
+				Name:     "group1",
+				File:     "ns",
+				Interval: 60,
+				Limit:    0,
+				Opts:     &promRules.ManagerOptions{},
+			}),
+			promRules.NewGroup(promRules.GroupOptions{
+				Name:     "group2",
+				File:     "ns",
+				Interval: 60,
+				Limit:    0,
+				Opts:     &promRules.ManagerOptions{},
+			}),
+		},
+	}
+
+	waitDurations := []time.Duration{
+		5 * time.Millisecond,
+		1 * time.Second,
+	}
+
+	ruleManagerFactory := RuleManagerFactory(groupsToReturn, waitDurations)
+
+	m, err := NewDefaultMultiTenantManager(Config{RulePath: dir}, ruleManagerFactory, nil, nil, log.NewNopLogger())
+	require.NoError(t, err)
+
+	m.SyncRuleGroups(context.Background(), userRules)
+
+	mgr := getManager(m, user)
+	require.NotNil(t, mgr)
+
+	test.Poll(t, 1*time.Second, true, func() interface{} {
+		return mgr.(*mockRulesManager).running.Load()
+	})
+	groups := m.GetRules(user)
+	require.Len(t, groups, len(groupsToReturn[0]), "expected %d but got %d", len(groupsToReturn[0]), len(groups))
+
+	// update rules and call list rules concurrently
+	userRules = map[string]rulespb.RuleGroupList{
+		user: {
+			&rulespb.RuleGroupDesc{
+				Name:      "group1",
+				Namespace: "ns",
+				Interval:  1 * time.Minute,
+				User:      user,
+			},
+			&rulespb.RuleGroupDesc{
+				Name:      "group2",
+				Namespace: "ns",
+				Interval:  1 * time.Minute,
+				User:      user,
+			},
+		},
+	}
+	go m.SyncRuleGroups(context.Background(), userRules)
+
+	groups = m.GetRules(user)
+	require.Len(t, groups, len(groupsToReturn[0]), "expected %d but got %d", len(groupsToReturn[0]), len(groups))
+
+	test.Poll(t, 5*time.Second, len(groupsToReturn[1]), func() interface{} {
+		groups = m.GetRules(user)
+		return len(groups)
+	})
+
+	test.Poll(t, 1*time.Second, true, func() interface{} {
+		return mgr.(*mockRulesManager).running.Load()
+	})
+
+	m.Stop()
+
+	test.Poll(t, 1*time.Second, false, func() interface{} {
+		return mgr.(*mockRulesManager).running.Load()
+	})
+}
+
 func TestSyncRuleGroupsCleanUpPerUserMetrics(t *testing.T) {
 	dir := t.TempDir()
 	reg := prometheus.NewPedanticRegistry()
@@ -139,8 +248,8 @@ func TestSyncRuleGroupsCleanUpPerUserMetrics(t *testing.T) {
 }
 
 func getManager(m *DefaultMultiTenantManager, user string) RulesManager {
-	m.userManagerMtx.Lock()
-	defer m.userManagerMtx.Unlock()
+	m.userManagerMtx.RLock()
+	defer m.userManagerMtx.RUnlock()
 
 	return m.userManagers[user]
 }
@@ -149,9 +258,46 @@ func factory(_ context.Context, _ string, _ *notifier.Manager, _ log.Logger, _ p
 	return &mockRulesManager{done: make(chan struct{})}
 }
 
+func RuleManagerFactory(groupsToReturn [][]*promRules.Group, waitDurations []time.Duration) ManagerFactory {
+	return func(_ context.Context, _ string, _ *notifier.Manager, _ log.Logger, _ prometheus.Registerer) RulesManager {
+		return &mockRulesManager{
+			done:           make(chan struct{}),
+			groupsToReturn: groupsToReturn,
+			waitDurations:  waitDurations,
+			iteration:      -1,
+		}
+	}
+}
+
 type mockRulesManager struct {
-	running atomic.Bool
-	done    chan struct{}
+	mtx            sync.Mutex
+	groupsToReturn [][]*promRules.Group
+	iteration      int
+	waitDurations  []time.Duration
+	running        atomic.Bool
+	done           chan struct{}
+}
+
+func (m *mockRulesManager) Update(_ time.Duration, _ []string, _ labels.Labels, _ string, _ promRules.GroupEvalIterationFunc) error {
+	m.mtx.Lock()
+	defer m.mtx.Unlock()
+	ticker := time.NewTicker(m.waitDurations[m.iteration+1])
+	select {
+	case <-ticker.C:
+		m.iteration = m.iteration + 1
+		return nil
+	case <-m.done:
+		return nil
+	}
+}
+
+func (m *mockRulesManager) RuleGroups() []*promRules.Group {
+	m.mtx.Lock()
+	m.mtx.Unlock()
+	if m.iteration < 0 {
+		return nil
+	}
+	return m.groupsToReturn[m.iteration]
 }
 
 func (m *mockRulesManager) Run() {
@@ -163,11 +309,3 @@ func (m *mockRulesManager) Stop() {
 	m.running.Store(false)
 	close(m.done)
 }
-
-func (m *mockRulesManager) Update(_ time.Duration, _ []string, _ labels.Labels, _ string, _ promRules.GroupEvalIterationFunc) error {
-	return nil
-}
-
-func (m *mockRulesManager) RuleGroups() []*promRules.Group {
-	return nil
-}