Repository: prometheus/prometheus
File: rules/manager.go (Go)

Code before:

// AnalyseRules implements RuleDependencyController.
func (c ruleDependencyController) AnalyseRules(rules []Rule) {
	depMap := buildDependencyMap(rules)

	if depMap == nil {
		return
	}

	for _, r := range rules {
		r.SetNoDependentRules(depMap.dependents(r) == 0)
		r.SetNoDependencyRules(depMap.dependencies(r) == 0)
	}
}

// ConcurrentRules represents a slice of indexes of rules that can be evaluated concurrently.
type ConcurrentRules []int

// RuleConcurrencyController controls concurrency for rules that are safe to be evaluated concurrently.
// Its purpose is to bound the amount of concurrency in rule evaluations to avoid overwhelming the Prometheus
// server with additional query load. Concurrency is controlled globally, not on a per-group basis.
type RuleConcurrencyController interface {
	// Sort returns a slice of ConcurrentRules, which are batches of rules that can be evaluated concurrently.
	// The rules are represented by their index in the original list of rules of the group.
	// We need the original index to update the group's state and we don't want to mutate the group's rules.
	Sort(ctx context.Context, group *Group) []ConcurrentRules

	// Allow determines if the given rule is allowed to be evaluated concurrently.
	// If Allow() returns true, then Done() must be called to release the acquired slot and corresponding cleanup is done.
	// It is important that both *Group and Rule are not retained and only be used for the duration of the call.
	Allow(ctx context.Context, group *Group, rule Rule) bool

	// Done releases a concurrent evaluation slot.
	Done(ctx context.Context)
}

// concurrentRuleEvalController holds a weighted semaphore which controls the concurrent evaluation of rules.
type concurrentRuleEvalController struct {
	sema *semaphore.Weighted
}

func newRuleConcurrencyController(maxConcurrency int64) RuleConcurrencyController {
	return &concurrentRuleEvalController{
		sema: semaphore.NewWeighted(maxConcurrency),
	}
}

func (c *concurrentRuleEvalController) Allow(_ context.Context, _ *Group, rule Rule) bool {
	return c.sema.TryAcquire(1)
}


Review comment (suggestion):
``` suggestion
	PartitionRules(ctx context.Context, group *Group) []int
```

Code after:

// AnalyseRules implements RuleDependencyController.
func (c ruleDependencyController) AnalyseRules(rules []Rule) {
	depMap := buildDependencyMap(rules)

	if depMap == nil {
		return
	}

	for _, r := range rules {
		r.SetNoDependentRules(depMap.dependents(r) == 0)
		r.SetNoDependencyRules(depMap.dependencies(r) == 0)
	}
}

// ConcurrentRules represents a slice of indexes of rules that can be evaluated concurrently.
type ConcurrentRules []int

// RuleConcurrencyController controls concurrency for rules that are safe to be evaluated concurrently.
// Its purpose is to bound the amount of concurrency in rule evaluations to avoid overwhelming the Prometheus
// server with additional query load. Concurrency is controlled globally, not on a per-group basis.
type RuleConcurrencyController interface {
	// SplitGroupIntoBatches returns an ordered slice of of ConcurrentRules, which are batches of rules that can be evaluated concurrently.
	// The rules are represented by their index from the input rule group.
	SplitGroupIntoBatches(ctx context.Context, group *Group) []ConcurrentRules

	// Allow determines if the given rule is allowed to be evaluated concurrently.
	// If Allow() returns true, then Done() must be called to release the acquired slot and corresponding cleanup is done.
	// It is important that both *Group and Rule are not retained and only be used for the duration of the call.
	Allow(ctx context.Context, group *Group, rule Rule) bool

	// Done releases a concurrent evaluation slot.
	Done(ctx context.Context)
}

// concurrentRuleEvalController holds a weighted semaphore which controls the concurrent evaluation of rules.
type concurrentRuleEvalController struct {
	sema *semaphore.Weighted
}

func newRuleConcurrencyController(maxConcurrency int64) RuleConcurrencyController {
	return &concurrentRuleEvalController{
		sema: semaphore.NewWeighted(maxConcurrency),
	}
}

func (c *concurrentRuleEvalController) Allow(_ context.Context, _ *Group, rule Rule) bool {
	return c.sema.TryAcquire(1)
}

func (c *concurrentRuleEvalController) SplitGroupIntoBatches(_ context.Context, g *Group) []ConcurrentRules {