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,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

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()