Remote Alertmanager: Fetch full state from Mimir (#107905)

* Remote Alertmanager: Add method to fetch the full state

* decode and parse the remote Alertmanager state string

* test

pkg/services/ngalert/remote/alertmanager.go

@@ -389,6 +389,48 @@ func (am *Alertmanager) sendConfiguration(ctx context.Context, cfg *remoteClient
 	return nil
 }
 
+// RemoteState represents the state (silences, nflog) in use by a remote Alertmanager.
+type RemoteState struct {
+	Silences []byte
+	Nflog    []byte
+}
+
+// GetRemoteState gets the remote Alertmanager's internal state.
+func (am *Alertmanager) GetRemoteState(ctx context.Context) (RemoteState, error) {
+	var rs RemoteState
+
+	s, err := am.mimirClient.GetFullState(ctx)
+	if err != nil {
+		return rs, fmt.Errorf("failed to pull remote state: %w", err)
+	}
+
+	// Decode and unmarshal the base64-encoded state we got from Mimir.
+	decoded, err := base64.StdEncoding.DecodeString(s.State)
+	if err != nil {
+		return rs, fmt.Errorf("failed to base64-decode remote state: %w", err)
+	}
+	protoState := &alertingClusterPB.FullState{}
+	if err := protoState.Unmarshal(decoded); err != nil {
+		return rs, fmt.Errorf("failed to unmarshal remote state: %w", err)
+	}
+
+	// Mimir state has two parts:
+	// - "sil:<tenantID>": silences
+	// - "nfl:<tenantID>": notification log entries
+	for _, p := range protoState.Parts {
+		switch p.Key {
+		case "sil:" + am.tenantID:
+			rs.Silences = p.Data
+		case "nfl:" + am.tenantID:
+			rs.Nflog = p.Data
+		default:
+			return rs, fmt.Errorf("unknown part key %q", p.Key)
+		}
+	}
+
+	return rs, nil
+}
+
 // SendState gets the Alertmanager's internal state and sends it to the remote Alertmanager.
 func (am *Alertmanager) SendState(ctx context.Context) error {
 	am.metrics.StateSyncsTotal.Inc()

pkg/services/ngalert/remote/alertmanager_test.go

@@ -25,6 +25,7 @@ import (
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/stretchr/testify/require"
 
+	alertingClusterPB "github.com/grafana/alerting/cluster/clusterpb"
 	"github.com/grafana/alerting/definition"
 	alertingModels "github.com/grafana/alerting/models"
 	"github.com/grafana/alerting/notify"
@@ -136,6 +137,125 @@ func TestNewAlertmanager(t *testing.T) {
 	}
 }
 
+func TestGetRemoteState(t *testing.T) {
+	const tenantID = "test"
+	ctx := context.Background()
+	store := ngfakes.NewFakeKVStore(t)
+	fstore := notifier.NewFileStore(1, store)
+	secretsService := secretsManager.SetupTestService(t, database.ProvideSecretsStore(db.InitTestDB(t)))
+	tc := notifier.NewCrypto(secretsService, nil, log.NewNopLogger())
+	m := metrics.NewRemoteAlertmanagerMetrics(prometheus.NewRegistry())
+
+	// getOkHandler allows us to specify a full state the test server is going to respond with.
+	getOkHandler := func(state string) http.HandlerFunc {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			require.Equal(t, tenantID, r.Header.Get(client.MimirTenantHeader))
+			require.Equal(t, "true", r.Header.Get(client.RemoteAlertmanagerHeader))
+
+			res := map[string]any{
+				"status": "success",
+				"data": map[string]any{
+					"state": state,
+				},
+			}
+			w.Header().Add("content-type", "application/json")
+			require.NoError(t, json.NewEncoder(w).Encode(res))
+		})
+	}
+
+	// errorHandler makes the test server return a 500 status code and a non-JSON response.
+	errorHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Add("content-type", "text/html; charset=utf-8")
+		w.WriteHeader(http.StatusInternalServerError)
+	})
+
+	// Test full states:
+	// - One with unknown part keys
+	// - One with the expected part keys
+	badState := alertingClusterPB.FullState{
+		Parts: []alertingClusterPB.Part{
+			{Key: "unknown", Data: []byte("data")},
+		},
+	}
+	rawBadState, err := badState.Marshal()
+	require.NoError(t, err)
+
+	state := alertingClusterPB.FullState{
+		Parts: []alertingClusterPB.Part{
+			{Key: "nfl:test", Data: []byte("test-nflog")},
+			{Key: "sil:test", Data: []byte("test-silences")},
+		},
+	}
+	rawState, err := state.Marshal()
+	require.NoError(t, err)
+
+	tests := []struct {
+		name        string
+		handler     http.Handler
+		expNflog    []byte
+		expSilences []byte
+		expErr      string
+	}{
+		{
+			name:    "non base64-encoded state",
+			handler: getOkHandler("invalid state"),
+			expErr:  "failed to base64-decode remote state: illegal base64 data at input byte 7",
+		},
+		{
+			name:    "error from the Mimir API",
+			handler: errorHandler,
+			expErr:  "failed to pull remote state: Response content-type is not application/json: text/html; charset=utf-8",
+		},
+		{
+			name:    "invalid state, base64-encoded",
+			handler: getOkHandler(base64.StdEncoding.EncodeToString([]byte("invalid state"))),
+			expErr:  "failed to unmarshal remote state: proto: FullState: wiretype end group for non-group",
+		},
+		{
+			name:    "unknown part key",
+			handler: getOkHandler(base64.StdEncoding.EncodeToString(rawBadState)),
+			expErr:  "unknown part key \"unknown\"",
+		},
+		{
+			name:        "success",
+			handler:     getOkHandler(base64.StdEncoding.EncodeToString(rawState)),
+			expNflog:    []byte("test-nflog"),
+			expSilences: []byte("test-silences"),
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(tt *testing.T) {
+			server := httptest.NewServer(test.handler)
+			cfg := AlertmanagerConfig{
+				OrgID:         1,
+				TenantID:      tenantID,
+				URL:           server.URL,
+				DefaultConfig: defaultGrafanaConfig,
+			}
+			am, err := NewAlertmanager(ctx,
+				cfg,
+				fstore,
+				tc,
+				NoopAutogenFn,
+				m,
+				tracing.InitializeTracerForTest(),
+			)
+			require.NoError(t, err)
+
+			s, err := am.GetRemoteState(ctx)
+			if test.expErr != "" {
+				require.Error(t, err)
+				require.Equal(t, test.expErr, err.Error())
+				return
+			}
+
+			require.NoError(t, err)
+			require.Equal(t, test.expNflog, s.Nflog)
+			require.Equal(t, test.expSilences, s.Silences)
+		})
+	}
+}
+
 func TestIntegrationApplyConfig(t *testing.T) {
 	const tenantID = "test"
 	// errorHandler returns an error response for the readiness check and state sync.

pkg/services/ngalert/remote/client/alertmanager_state.go

@@ -9,6 +9,7 @@ import (
 )
 
 const (
+	fullStatePath                = "/api/v1/grafana/full_state"
 	grafanaAlertmanagerStatePath = "/api/v1/grafana/state"
 )
 
@@ -17,7 +18,15 @@ type UserState struct {
 	State string `json:"state"`
 }
 
+func (mc *Mimir) GetFullState(ctx context.Context) (*UserState, error) {
+	return mc.getState(ctx, fullStatePath)
+}
+
 func (mc *Mimir) GetGrafanaAlertmanagerState(ctx context.Context) (*UserState, error) {
+	return mc.getState(ctx, grafanaAlertmanagerStatePath)
+}
+
+func (mc *Mimir) getState(ctx context.Context, path string) (*UserState, error) {
 	gs := &UserState{}
 	response := successResponse{
 		Data: gs,

pkg/services/ngalert/remote/client/mimir.go

@@ -24,6 +24,7 @@ import (
 
 // MimirClient contains all the methods to query the migration critical endpoints of Mimir instance, it's an interface to allow multiple implementations.
 type MimirClient interface {
+	GetFullState(ctx context.Context) (*UserState, error)
 	GetGrafanaAlertmanagerState(ctx context.Context) (*UserState, error)
 	CreateGrafanaAlertmanagerState(ctx context.Context, state string) error
 	DeleteGrafanaAlertmanagerState(ctx context.Context) error