Repurpose FederatedContextFunc into GroupEvaluationContextFunc

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,8 +911,9 @@ 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
+	// 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
@@ -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

rules/manager_test.go

@@ -722,7 +722,6 @@ func TestUpdate(t *testing.T) {
 		Queryable:  st,
 		QueryFunc:  EngineQueryFunc(engine, st),
 		Context:    context.Background(),
-		FederatedContextFunc: func(*Group) context.Context { return context.Background() },
 		Logger:     log.NewNopLogger(),
 	})
 	ruleManager.start()