Use routeID in aggrGroup and mute stages

Signed-off-by: George Robinson <george.robinson@grafana.com>

Author: grobinson-grafana

dispatch/dispatch_test.go

@@ -725,14 +725,14 @@ func TestDispatcher_DoMaintenance(t *testing.T) {
 	go aggrGroup1.run(func(context.Context, ...*types.Alert) bool { return true })
 
 	// Insert a marker for the aggregation group's group key.
-	marker.SetMuted(aggrGroup1.GroupKey(), []string{"weekends"})
-	mutedBy, isMuted := marker.Muted(aggrGroup1.GroupKey())
+	marker.SetMuted(route.ID(), aggrGroup1.GroupKey(), []string{"weekends"})
+	mutedBy, isMuted := marker.Muted(route.ID(), aggrGroup1.GroupKey())
 	require.True(t, isMuted)
 	require.Equal(t, []string{"weekends"}, mutedBy)
 
 	// Run the maintenance and the marker should be removed.
 	dispatcher.doMaintenance()
-	mutedBy, isMuted = marker.Muted(aggrGroup1.GroupKey())
+	mutedBy, isMuted = marker.Muted(route.ID(), aggrGroup1.GroupKey())
 	require.False(t, isMuted)
 	require.Empty(t, mutedBy)
 }

notify/notify.go

@@ -119,6 +119,7 @@ const (
 	keyNow
 	keyMuteTimeIntervals
 	keyActiveTimeIntervals
+	keyRouteID
 )
 
 // WithReceiverName populates a context with a receiver name.
@@ -165,6 +166,10 @@ func WithActiveTimeIntervals(ctx context.Context, at []string) context.Context {
 	return context.WithValue(ctx, keyActiveTimeIntervals, at)
 }
 
+func WithRouteID(ctx context.Context, routeID string) context.Context {
+	return context.WithValue(ctx, keyRouteID, routeID)
+}
+
 // RepeatInterval extracts a repeat interval from the context. Iff none exists, the
 // second argument is false.
 func RepeatInterval(ctx context.Context) (time.Duration, bool) {
@@ -228,6 +233,13 @@ func ActiveTimeIntervalNames(ctx context.Context) ([]string, bool) {
 	return v, ok
 }
 
+// RouteID extracts a RouteID from the context. Iff none exists, the
+// // second argument is false.
+func RouteID(ctx context.Context) (string, bool) {
+	v, ok := ctx.Value(keyRouteID).(string)
+	return v, ok
+}
+
 // A Stage processes alerts under the constraints of the given context.
 type Stage interface {
 	Exec(ctx context.Context, l log.Logger, alerts ...*types.Alert) (context.Context, []*types.Alert, error)
@@ -937,6 +949,11 @@ func NewTimeMuteStage(muter types.TimeMuter, marker types.GroupMarker, metrics *
 // Exec implements the stage interface for TimeMuteStage.
 // TimeMuteStage is responsible for muting alerts whose route is not in an active time.
 func (tms TimeMuteStage) Exec(ctx context.Context, l log.Logger, alerts ...*types.Alert) (context.Context, []*types.Alert, error) {
+	routeID, ok := RouteID(ctx)
+	if !ok {
+		return ctx, nil, errors.New("route ID missing")
+	}
+
 	gkey, ok := GroupKey(ctx)
 	if !ok {
 		return ctx, nil, errors.New("group key missing")
@@ -961,7 +978,7 @@ func (tms TimeMuteStage) Exec(ctx context.Context, l log.Logger, alerts ...*type
 		return ctx, alerts, err
 	}
 	// If muted is false then mutedBy is nil and the muted marker is removed.
-	tms.marker.SetMuted(gkey, mutedBy)
+	tms.marker.SetMuted(routeID, gkey, mutedBy)
 
 	// If the current time is inside a mute time, all alerts are removed from the pipeline.
 	if muted {
@@ -982,6 +999,11 @@ func NewTimeActiveStage(muter types.TimeMuter, marker types.GroupMarker, metrics
 // Exec implements the stage interface for TimeActiveStage.
 // TimeActiveStage is responsible for muting alerts whose route is not in an active time.
 func (tas TimeActiveStage) Exec(ctx context.Context, l log.Logger, alerts ...*types.Alert) (context.Context, []*types.Alert, error) {
+	routeID, ok := RouteID(ctx)
+	if !ok {
+		return ctx, nil, errors.New("route ID missing")
+	}
+
 	gkey, ok := GroupKey(ctx)
 	if !ok {
 		return ctx, nil, errors.New("group key missing")
@@ -1014,7 +1036,7 @@ func (tas TimeActiveStage) Exec(ctx context.Context, l log.Logger, alerts ...*ty
 		// to be active.
 		mutedBy = activeTimeIntervalNames
 	}
-	tas.marker.SetMuted(gkey, mutedBy)
+	tas.marker.SetMuted(routeID, gkey, mutedBy)
 
 	// If the current time is not inside an active time, all alerts are removed from the pipeline
 	if !active {

notify/notify_test.go

@@ -906,6 +906,7 @@ func TestTimeMuteStage(t *testing.T) {
 			ctx = WithGroupKey(ctx, "group1")
 			ctx = WithActiveTimeIntervals(ctx, nil)
 			ctx = WithMuteTimeIntervals(ctx, muteTimeIntervalNames)
+			ctx = WithRouteID(ctx, "route1")
 
 			_, active, err := st.Exec(ctx, log.NewNopLogger(), test.alerts...)
 			require.NoError(t, err)
@@ -914,7 +915,7 @@ func TestTimeMuteStage(t *testing.T) {
 				// All alerts should be active.
 				require.Equal(t, len(test.alerts), len(active))
 				// The group should not be marked.
-				mutedBy, isMuted := marker.Muted("group1")
+				mutedBy, isMuted := marker.Muted("route1", "group1")
 				require.False(t, isMuted)
 				require.Empty(t, mutedBy)
 				// The metric for total suppressed notifications should not
@@ -930,7 +931,7 @@ alertmanager_marked_alerts{state="unprocessed"} 0
 				// All alerts should be muted.
 				require.Empty(t, active)
 				// The group should be marked as muted.
-				mutedBy, isMuted := marker.Muted("group1")
+				mutedBy, isMuted := marker.Muted("route1", "group1")
 				require.True(t, isMuted)
 				require.Equal(t, test.mutedBy, mutedBy)
 				// Gets the metric for total suppressed notifications.
@@ -1023,6 +1024,7 @@ func TestTimeActiveStage(t *testing.T) {
 			ctx = WithGroupKey(ctx, "group1")
 			ctx = WithActiveTimeIntervals(ctx, activeTimeIntervalNames)
 			ctx = WithMuteTimeIntervals(ctx, nil)
+			ctx = WithRouteID(ctx, "route1")
 
 			_, active, err := st.Exec(ctx, log.NewNopLogger(), test.alerts...)
 			require.NoError(t, err)
@@ -1031,7 +1033,7 @@ func TestTimeActiveStage(t *testing.T) {
 				// All alerts should be active.
 				require.Equal(t, len(test.alerts), len(active))
 				// The group should not be marked.
-				mutedBy, isMuted := marker.Muted("group1")
+				mutedBy, isMuted := marker.Muted("route1", "group1")
 				require.False(t, isMuted)
 				require.Empty(t, mutedBy)
 				// The metric for total suppressed notifications should not
@@ -1047,7 +1049,7 @@ alertmanager_marked_alerts{state="unprocessed"} 0
 				// All alerts should be muted.
 				require.Empty(t, active)
 				// The group should be marked as muted.
-				mutedBy, isMuted := marker.Muted("group1")
+				mutedBy, isMuted := marker.Muted("route1", "group1")
 				require.True(t, isMuted)
 				require.Equal(t, test.mutedBy, mutedBy)
 				// Gets the metric for total suppressed notifications.