diff --git a/rules/manager.go b/rules/manager.go index 17b8c4dca..973a47161 100644 --- a/rules/manager.go +++ b/rules/manager.go @@ -903,7 +903,7 @@ type Manager struct { // NotifyFunc sends notifications about a set of alerts generated by the given expression. type NotifyFunc func(ctx context.Context, expr string, alerts ...*Alert) -type ContextFunc func(g *Group) context.Context +type ContextWrapFunc func(ctx context.Context, g *Group) context.Context // ManagerOptions bundles options for the Manager. type ManagerOptions struct { @@ -911,16 +911,17 @@ type ManagerOptions struct { QueryFunc QueryFunc NotifyFunc NotifyFunc Context context.Context - // FederatedContextFunc will be called to obtain a context when evaluating rules with non-empty SourceTenants - FederatedContextFunc ContextFunc - Appendable storage.Appendable - Queryable storage.Queryable - Logger log.Logger - Registerer prometheus.Registerer - OutageTolerance time.Duration - ForGracePeriod time.Duration - ResendDelay time.Duration - GroupLoader GroupLoader + // GroupEvaluationContextFunc will be called to wrap Context based on the group being evaluated. + // Will be skipped if nil. + GroupEvaluationContextFunc ContextWrapFunc + Appendable storage.Appendable + Queryable storage.Queryable + Logger log.Logger + Registerer prometheus.Registerer + OutageTolerance time.Duration + ForGracePeriod time.Duration + ResendDelay time.Duration + GroupLoader GroupLoader Metrics *Metrics } @@ -1014,8 +1015,8 @@ func (m *Manager) Update(interval time.Duration, files []string, externalLabels wg.Done() ctx := m.opts.Context - if len(newg.sourceTenants) > 0 { - ctx = m.opts.FederatedContextFunc(newg) + if m.opts.GroupEvaluationContextFunc != nil { + ctx = m.opts.GroupEvaluationContextFunc(ctx, newg) } // Wait with starting evaluation until the rule manager // is told to run. This is necessary to avoid running diff --git a/rules/manager_test.go b/rules/manager_test.go index ae77caa6b..ff2f42dd4 100644 --- a/rules/manager_test.go +++ b/rules/manager_test.go @@ -718,12 +718,11 @@ func TestUpdate(t *testing.T) { } engine := promql.NewEngine(opts) ruleManager := NewManager(&ManagerOptions{ - Appendable: st, - Queryable: st, - QueryFunc: EngineQueryFunc(engine, st), - Context: context.Background(), - FederatedContextFunc: func(*Group) context.Context { return context.Background() }, - Logger: log.NewNopLogger(), + Appendable: st, + Queryable: st, + QueryFunc: EngineQueryFunc(engine, st), + Context: context.Background(), + Logger: log.NewNopLogger(), }) ruleManager.start() defer ruleManager.Stop()