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

Code before:
		opts.MaxConcurrentEvals = int64(ruleCount) * 2
		opts.RuleConcurrencyController = nil
		ruleManager := NewManager(opts)

		groups, errs := ruleManager.LoadGroups(time.Second, labels.EmptyLabels(), "", nil, []string{"fixtures/rules_multiple_independent.yaml"}...)
		require.Empty(t, errs)
		require.Len(t, groups, 1)

		for _, group := range groups {
			require.Len(t, group.rules, ruleCount)

			start := time.Now()

			group.Eval(ctx, start)

			// Max inflight can be up to MaxConcurrentEvals concurrent evals, since there is sufficient concurrency to run all rules at once.
			require.LessOrEqual(t, int64(maxInflight.Load()), opts.MaxConcurrentEvals)
			// Some rules should execute concurrently so should complete quicker.
			require.Less(t, time.Since(start).Seconds(), (time.Duration(ruleCount) * artificialDelay).Seconds())
			// Each rule produces one vector.
			require.EqualValues(t, ruleCount, testutil.ToFloat64(group.metrics.GroupSamples))
		}
	})
}

func TestBoundedRuleEvalConcurrency(t *testing.T) {
	storage := teststorage.New(t)
	t.Cleanup(func() { storage.Close() })

	var (
		inflightQueries atomic.Int32
		maxInflight     atomic.Int32
		maxConcurrency  int64 = 3
		groupCount            = 2
	)

	files := []string{"fixtures/rules_multiple_groups.yaml"}

	ruleManager := NewManager(optsFactory(storage, &maxInflight, &inflightQueries, maxConcurrency))

	groups, errs := ruleManager.LoadGroups(time.Second, labels.EmptyLabels(), "", nil, files...)
	require.Empty(t, errs)
	require.Len(t, groups, groupCount)

	ctx, cancel := context.WithCancel(context.Background())
	t.Cleanup(cancel)

	// Evaluate groups concurrently (like they normally do).
	var wg sync.WaitGroup
	for _, group := range groups {
		group := group

Review comment (suggestion):
Thanks for the tests, nice addition.
I ran them on my machine and they took `2-4s` each, if you could make them runnable in parallel it'd be great.

Code after:
		opts.MaxConcurrentEvals = int64(ruleCount) * 2
		opts.RuleConcurrencyController = nil
		ruleManager := NewManager(opts)

		groups, errs := ruleManager.LoadGroups(time.Second, labels.EmptyLabels(), "", nil, []string{"fixtures/rules_multiple_independent.yaml"}...)
		require.Empty(t, errs)
		require.Len(t, groups, 1)

		for _, group := range groups {
			require.Len(t, group.rules, ruleCount)

			start := time.Now()

			group.Eval(ctx, start)

			// Max inflight can be up to MaxConcurrentEvals concurrent evals, since there is sufficient concurrency to run all rules at once.
			require.LessOrEqual(t, int64(maxInflight.Load()), opts.MaxConcurrentEvals)
			// Some rules should execute concurrently so should complete quicker.
			require.Less(t, time.Since(start).Seconds(), (time.Duration(ruleCount) * artificialDelay).Seconds())
			// Each rule produces one vector.
			require.EqualValues(t, ruleCount, testutil.ToFloat64(group.metrics.GroupSamples))
		}
	})
}

func TestNewRuleGroupRestoration(t *testing.T) {
	store := teststorage.New(t)
	t.Cleanup(func() { store.Close() })
	var (
		inflightQueries atomic.Int32
		maxInflight     atomic.Int32
		maxConcurrency  int64
		interval        = 60 * time.Second
	)

	waitForEvaluations := func(t *testing.T, ch <-chan int32, targetCount int32) {
		for {
			select {
			case cnt := <-ch:
				if cnt == targetCount {
					return
				}
			case <-time.After(5 * time.Second):
				return
			}
		}
	}

	files := []string{"fixtures/alert_rule.yaml"}

	option := optsFactory(store, &maxInflight, &inflightQueries, maxConcurrency)