o11y/pkg/alertmanager/alertmanagerserver/server.go

package alertmanagerserver

import (
	"context"
	"log/slog"
	"strings"
	"sync"
	"time"

	"github.com/prometheus/alertmanager/types"
	"golang.org/x/sync/errgroup"

	"github.com/hanzoai/o11y/pkg/errors"
	"github.com/prometheus/alertmanager/dispatch"
	"github.com/prometheus/alertmanager/featurecontrol"
	"github.com/prometheus/alertmanager/inhibit"
	"github.com/prometheus/alertmanager/nflog"
	"github.com/prometheus/alertmanager/notify"
	"github.com/prometheus/alertmanager/provider/mem"
	"github.com/prometheus/alertmanager/silence"
	"github.com/prometheus/alertmanager/template"
	"github.com/prometheus/alertmanager/timeinterval"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/common/model"

	"github.com/hanzoai/o11y/pkg/alertmanager/alertmanagernotify"
	"github.com/hanzoai/o11y/pkg/alertmanager/alertmanagertemplate"
	"github.com/hanzoai/o11y/pkg/alertmanager/nfmanager"
	"github.com/hanzoai/o11y/pkg/types/alertmanagertypes"
)

// This is not a real snapshot file and will never be used. We need this placeholder to ensure maintenance runs on shutdown.
// See https://github.com/prometheus/alertmanager/blob/3ee2cd0f1271e277295c02b6160507b4d193dde2/silence/silence.go#L435-L438
// and https://github.com/prometheus/alertmanager/blob/3b06b97af4d146e141af92885a185891eb79a5b0/nflog/nflog.go#L362.
var snapfnoop string = "snapfnoop"

type Server struct {
	// logger is the logger for the alertmanager
	logger *slog.Logger

	// registry is the prometheus registry for the alertmanager
	registry prometheus.Registerer

	// srvConfig is the server config for the alertmanager
	srvConfig Config

	// alertmanagerConfig is the config of the alertmanager
	alertmanagerConfig *alertmanagertypes.Config

	// orgID is the orgID for the alertmanager
	orgID string

	// store is the backing store for the alertmanager
	stateStore alertmanagertypes.StateStore

	// alertmanager primitives from upstream alertmanager
	alerts              *mem.Alerts
	nflog               *nflog.Log
	dispatcher          *Dispatcher
	dispatcherMetrics   *DispatcherMetrics
	inhibitor           *inhibit.Inhibitor
	silencer            *silence.Silencer
	silences            *silence.Silences
	timeIntervals       map[string][]timeinterval.TimeInterval
	pipelineBuilder     *pipelineBuilder
	muter               *MaintenanceMuter
	marker              *types.MemMarker
	tmpl                *template.Template
	templater           alertmanagertypes.Templater
	wg                  sync.WaitGroup
	stopc               chan struct{}
	notificationManager nfmanager.NotificationManager
}

func New(
	ctx context.Context,
	logger *slog.Logger,
	registry prometheus.Registerer,
	srvConfig Config,
	orgID string,
	stateStore alertmanagertypes.StateStore,
	nfManager nfmanager.NotificationManager,
	maintenanceStore alertmanagertypes.MaintenanceStore,
) (*Server, error) {
	server := &Server{
		logger:              logger.With(slog.String("pkg", "go.o11y.io/pkg/alertmanager/alertmanagerserver")),
		registry:            registry,
		srvConfig:           srvConfig,
		orgID:               orgID,
		stateStore:          stateStore,
		stopc:               make(chan struct{}),
		notificationManager: nfManager,
	}
	o11yRegisterer := prometheus.WrapRegistererWithPrefix("o11y_", registry)
	o11yRegisterer = prometheus.WrapRegistererWith(prometheus.Labels{"org_id": server.orgID}, o11yRegisterer)
	// initialize marker
	server.marker = types.NewMarker(o11yRegisterer)

	// get silences for initial state
	state, err := server.stateStore.Get(ctx, server.orgID)
	if err != nil && !errors.Ast(err, errors.TypeNotFound) {
		return nil, err
	}

	silencesSnapshot := ""
	if state != nil {
		silencesSnapshot, err = state.Get(alertmanagertypes.SilenceStateName)
		if err != nil && !errors.Ast(err, errors.TypeNotFound) {
			return nil, err
		}
	}
	// Initialize silences
	server.silences, err = silence.New(silence.Options{
		SnapshotReader: strings.NewReader(silencesSnapshot),
		Retention:      srvConfig.Silences.Retention,
		Limits: silence.Limits{
			MaxSilences:         func() int { return srvConfig.Silences.Max },
			MaxSilenceSizeBytes: func() int { return srvConfig.Silences.MaxSizeBytes },
		},
		Metrics: o11yRegisterer,
		Logger:  server.logger,
	})
	if err != nil {
		return nil, err
	}

	nflogSnapshot := ""
	if state != nil {
		nflogSnapshot, err = state.Get(alertmanagertypes.NFLogStateName)
		if err != nil && !errors.Ast(err, errors.TypeNotFound) {
			return nil, err
		}
	}

	// Initialize notification log
	server.nflog, err = nflog.New(nflog.Options{
		SnapshotReader: strings.NewReader(nflogSnapshot),
		Retention:      server.srvConfig.NFLog.Retention,
		Metrics:        o11yRegisterer,
		Logger:         server.logger,
	})
	if err != nil {
		return nil, err
	}

	// Start maintenance for silences
	server.wg.Add(1)
	go func() {
		defer server.wg.Done()
		server.silences.Maintenance(server.srvConfig.Silences.MaintenanceInterval, snapfnoop, server.stopc, func() (int64, error) {
			// Delete silences older than the retention period.
			if _, err := server.silences.GC(); err != nil {
				server.logger.ErrorContext(ctx, "silence garbage collection", errors.Attr(err))
				// Don't return here - we need to snapshot our state first.
			}

			storableSilences, err := server.stateStore.Get(ctx, server.orgID)
			if err != nil && !errors.Ast(err, errors.TypeNotFound) {
				return 0, err
			}

			if storableSilences == nil {
				storableSilences = alertmanagertypes.NewStoreableState(server.orgID)
			}

			c, err := storableSilences.Set(alertmanagertypes.SilenceStateName, server.silences)
			if err != nil {
				return 0, err
			}

			return c, server.stateStore.Set(ctx, storableSilences)
		})
	}()

	// Start maintenance for notification logs
	server.wg.Add(1)
	go func() {
		defer server.wg.Done()
		server.nflog.Maintenance(server.srvConfig.NFLog.MaintenanceInterval, snapfnoop, server.stopc, func() (int64, error) {
			if _, err := server.nflog.GC(); err != nil {
				server.logger.ErrorContext(ctx, "notification log garbage collection", errors.Attr(err))
				// Don't return without saving the current state.
			}

			storableNFLog, err := server.stateStore.Get(ctx, server.orgID)
			if err != nil && !errors.Ast(err, errors.TypeNotFound) {
				return 0, err
			}

			if storableNFLog == nil {
				storableNFLog = alertmanagertypes.NewStoreableState(server.orgID)
			}

			c, err := storableNFLog.Set(alertmanagertypes.NFLogStateName, server.nflog)
			if err != nil {
				return 0, err
			}

			return c, server.stateStore.Set(ctx, storableNFLog)
		})
	}()

	server.alerts, err = mem.NewAlerts(ctx, server.marker, server.srvConfig.Alerts.GCInterval, 0, nil, server.logger, o11yRegisterer, nil)
	if err != nil {
		return nil, err
	}

	server.muter = NewMaintenanceMuter(maintenanceStore, orgID, server.logger)
	server.pipelineBuilder = newPipelineBuilder(o11yRegisterer, featurecontrol.NoopFlags{})
	server.dispatcherMetrics = NewDispatcherMetrics(false, o11yRegisterer)

	return server, nil
}

func (server *Server) GetAlerts(ctx context.Context, params alertmanagertypes.GettableAlertsParams) (alertmanagertypes.GettableAlerts, error) {
	return alertmanagertypes.NewGettableAlertsFromAlertProvider(
		server.alerts, server.alertmanagerConfig, server.marker.Status,
		func(labels model.LabelSet) {
			server.inhibitor.Mutes(ctx, labels)
			server.silencer.Mutes(ctx, labels)
		},
		func(labels model.LabelSet) []string {
			return server.muter.MutedBy(ctx, labels)
		},
		params,
	)
}

func (server *Server) PutAlerts(ctx context.Context, postableAlerts alertmanagertypes.PostableAlerts) error {
	alerts, err := alertmanagertypes.NewAlertsFromPostableAlerts(ctx, postableAlerts, time.Duration(server.srvConfig.Global.ResolveTimeout), time.Now())
	// Notification sending alert takes precedence over validation errors.
	if err := server.alerts.Put(ctx, alerts...); err != nil {
		return err
	}

	if err != nil {
		return errors.Join(err...)
	}

	return nil
}

func (server *Server) SetConfig(ctx context.Context, alertmanagerConfig *alertmanagertypes.Config) error {
	config := alertmanagerConfig.AlertmanagerConfig()

	var err error
	// Load O11y's alertmanager notification templates from the configured
	// globs. The upstream default templates (default.tmpl, email.tmpl) are
	// always loaded from the embedded alertmanager assets inside FromGlobs, so
	// only O11y's own templates (e.g. the email.o11y.html layout) are listed
	// here. The upstream config.Templates field is not used: O11y never
	// populates it (there is no per-org template configuration).
	server.tmpl, err = alertmanagertypes.FromGlobs(server.srvConfig.Templates)
	if err != nil {
		return err
	}

	server.tmpl.ExternalURL = server.srvConfig.ExternalURL

	server.templater = alertmanagertemplate.New(server.tmpl, server.logger)

	// Build the routing tree and record which receivers are used.
	routes := dispatch.NewRoute(config.Route, nil)
	activeReceivers := make(map[string]struct{})
	routes.Walk(func(r *dispatch.Route) {
		activeReceivers[r.RouteOpts.Receiver] = struct{}{}
	})

	// Build the map of receiver to integrations.
	receivers := make(map[string][]notify.Integration, len(activeReceivers))
	var integrationsNum int
	for _, rcv := range config.Receivers {
		if _, found := activeReceivers[rcv.Name]; !found {
			// No need to build a receiver if no route is using it.
			server.logger.InfoContext(ctx, "skipping creation of receiver not referenced by any route", slog.String("receiver", rcv.Name))
			continue
		}
		extendedRcv, err := alertmanagerConfig.GetReceiver(rcv.Name)
		if err != nil {
			return err
		}
		integrations, err := alertmanagernotify.NewReceiverIntegrations(extendedRcv, server.tmpl, server.logger, server.templater)
		if err != nil {
			return err
		}
		// rcv.Name is guaranteed to be unique across all receivers.
		receivers[rcv.Name] = integrations
		integrationsNum += len(integrations)
	}

	// Build the map of time interval names to time interval definitions.
	timeIntervals := make(map[string][]timeinterval.TimeInterval, len(config.MuteTimeIntervals)+len(config.TimeIntervals))
	for _, ti := range config.MuteTimeIntervals {
		timeIntervals[ti.Name] = ti.TimeIntervals
	}

	for _, ti := range config.TimeIntervals {
		timeIntervals[ti.Name] = ti.TimeIntervals
	}

	intervener := timeinterval.NewIntervener(timeIntervals)

	if server.inhibitor != nil {
		server.inhibitor.Stop()
	}
	if server.dispatcher != nil {
		server.dispatcher.Stop()
	}

	server.inhibitor = inhibit.NewInhibitor(server.alerts, config.InhibitRules, server.marker, server.logger)
	server.timeIntervals = timeIntervals
	server.silencer = silence.NewSilencer(server.silences, server.marker, server.logger)

	var pipelinePeer notify.Peer
	pipeline := server.pipelineBuilder.New(
		receivers,
		func() time.Duration { return 0 },
		server.inhibitor,
		server.silencer,
		intervener,
		server.marker,
		server.muter,
		server.nflog,
		pipelinePeer,
	)

	timeoutFunc := func(d time.Duration) time.Duration {
		if d < notify.MinTimeout {
			d = notify.MinTimeout
		}
		return d
	}

	server.dispatcher = NewDispatcher(
		server.alerts,
		routes,
		pipeline,
		server.marker,
		timeoutFunc,
		nil,
		server.logger,
		server.dispatcherMetrics,
		server.notificationManager,
		server.orgID,
	)

	// Do not try to add these to server.wg as there seems to be a race condition if
	// we call Start() and Stop() in quick succession.
	// Both these goroutines will run indefinitely.
	go server.dispatcher.Run()
	go server.inhibitor.Run()

	server.alertmanagerConfig = alertmanagerConfig
	return nil
}

func (server *Server) TestReceiver(ctx context.Context, receiver *alertmanagertypes.Receiver) error {
	testAlert := alertmanagertypes.NewTestAlert(receiver, time.Now(), time.Now())
	return alertmanagertypes.TestReceiver(ctx, receiver, alertmanagernotify.NewReceiverIntegrations, server.alertmanagerConfig, server.tmpl, server.logger, server.templater, testAlert.Labels, testAlert)
}

func (server *Server) TestAlert(ctx context.Context, receiversMap map[*alertmanagertypes.PostableAlert][]string, config *alertmanagertypes.NotificationConfig) error {
	if len(receiversMap) == 0 {
		return errors.Newf(errors.TypeInvalidInput, errors.CodeInvalidInput,
			"expected at least 1 alert, got 0")
	}

	postableAlerts := make(alertmanagertypes.PostableAlerts, 0, len(receiversMap))
	for alert := range receiversMap {
		postableAlerts = append(postableAlerts, alert)
	}

	alerts, err := alertmanagertypes.NewAlertsFromPostableAlerts(
		ctx,
		postableAlerts,
		time.Duration(server.srvConfig.Global.ResolveTimeout),
		time.Now(),
	)
	if err != nil {
		return errors.Newf(errors.TypeInvalidInput, errors.CodeInvalidInput,
			"failed to construct alerts from postable alerts: %v", err)
	}

	type alertGroup struct {
		groupLabels model.LabelSet
		alerts      []*types.Alert
		receivers   map[string]struct{}
	}

	groupMap := make(map[model.Fingerprint]*alertGroup)

	for i, alert := range alerts {
		labels := getGroupLabels(alert, config.NotificationGroup, config.GroupByAll)
		fp := labels.Fingerprint()

		postableAlert := postableAlerts[i]
		alertReceivers := receiversMap[postableAlert]

		if group, exists := groupMap[fp]; exists {
			group.alerts = append(group.alerts, alert)
			for _, r := range alertReceivers {
				group.receivers[r] = struct{}{}
			}
		} else {
			receiverSet := make(map[string]struct{})
			for _, r := range alertReceivers {
				receiverSet[r] = struct{}{}
			}
			groupMap[fp] = &alertGroup{
				groupLabels: labels,
				alerts:      []*types.Alert{alert},
				receivers:   receiverSet,
			}
		}
	}

	var mu sync.Mutex
	var errs []error

	g, gCtx := errgroup.WithContext(ctx)
	for _, group := range groupMap {
		for receiverName := range group.receivers {
			group := group
			receiverName := receiverName

			g.Go(func() error {
				receiver, err := server.alertmanagerConfig.GetReceiver(receiverName)
				if err != nil {
					mu.Lock()
					errs = append(errs, errors.WrapInternalf(err, errors.CodeInternal, "failed to get receiver %q", receiverName))
					mu.Unlock()
					return nil // Return nil to continue processing other goroutines
				}

				err = alertmanagertypes.TestReceiver(
					gCtx,
					receiver,
					alertmanagernotify.NewReceiverIntegrations,
					server.alertmanagerConfig,
					server.tmpl,
					server.logger,
					server.templater,
					group.groupLabels,
					group.alerts...,
				)
				if err != nil {
					mu.Lock()
					errs = append(errs, errors.WrapInternalf(err, errors.CodeInternal, "receiver %q test failed", receiverName))
					mu.Unlock()
				}
				return nil // Return nil to continue processing other goroutines
			})
		}
	}
	_ = g.Wait()

	if len(errs) > 0 {
		return errors.Join(errs...)
	}

	return nil
}

func (server *Server) Hash() string {
	if server.alertmanagerConfig == nil {
		return ""
	}

	return server.alertmanagerConfig.StoreableConfig().Hash
}

func (server *Server) Stop(ctx context.Context) error {
	if server.dispatcher != nil {
		server.dispatcher.Stop()
	}

	if server.inhibitor != nil {
		server.inhibitor.Stop()
	}

	// Close the alert provider.
	server.alerts.Close()

	// Signals maintenance goroutines of server states to stop.
	close(server.stopc)

	// Wait for all goroutines to finish.
	server.wg.Wait()

	return nil
}