Implement ReadFullStateForUser in MultitenantAlertmanager.

Signed-off-by: Steve Simpson <[email protected]>

Author: stevesg

pkg/alertmanager/alertmanager.go

@@ -161,13 +161,7 @@ func New(cfg *Config, reg *prometheus.Registry) (*Alertmanager, error) {
 		am.state = cfg.Peer
 	} else if cfg.ShardingEnabled {
 		level.Debug(am.logger).Log("msg", "starting tenant alertmanager with ring-based replication")
-		state := newReplicatedStates(cfg.UserID, cfg.ReplicationFactor, cfg.Replicator, am.logger, am.registry)
-
-		if err := state.Service.StartAsync(context.Background()); err != nil {
-			return nil, errors.Wrap(err, "failed to start ring-based replication service")
-		}
-
-		am.state = state
+		am.state = newReplicatedStates(cfg.UserID, cfg.ReplicationFactor, cfg.Replicator, am.logger, am.registry)
 	} else {
 		level.Debug(am.logger).Log("msg", "starting tenant alertmanager without replication")
 		am.state = &NilPeer{}
@@ -205,6 +199,13 @@ func New(cfg *Config, reg *prometheus.Registry) (*Alertmanager, error) {
 	c = am.state.AddState("sil:"+cfg.UserID, am.silences, am.registry)
 	am.silences.SetBroadcast(c.Broadcast)
 
+	// State replication needs to be started after the state keys are defined.
+	if service, ok := am.state.(services.Service); ok {
+		if err := service.StartAsync(context.Background()); err != nil {
+			return nil, errors.Wrap(err, "failed to start ring-based replication service")
+		}
+	}
+
 	am.pipelineBuilder = notify.NewPipelineBuilder(am.registry)
 
 	am.wg.Add(1)

pkg/alertmanager/multitenant.go

@@ -35,6 +35,7 @@ import (
 	"github.com/cortexproject/cortex/pkg/ring/kv"
 	"github.com/cortexproject/cortex/pkg/tenant"
 	"github.com/cortexproject/cortex/pkg/util"
+	"github.com/cortexproject/cortex/pkg/util/concurrency"
 	"github.com/cortexproject/cortex/pkg/util/flagext"
 	util_log "github.com/cortexproject/cortex/pkg/util/log"
 	"github.com/cortexproject/cortex/pkg/util/services"
@@ -1011,8 +1012,77 @@ func (am *MultitenantAlertmanager) ReplicateStateForUser(ctx context.Context, us
 	return err
 }
 
-func (am *MultitenantAlertmanager) ReadFullStateForUser(context.Context, string) ([]*clusterpb.FullState, error) {
-	return nil, nil
+// ReadFullStateForUser attempts to read the full state from each replica for user. Note that it will try to obtain and return
+// state from all replicas, but will consider it a success if state is obtained from at least one replica.
+func (am *MultitenantAlertmanager) ReadFullStateForUser(ctx context.Context, userID string) ([]*clusterpb.FullState, error) {
+
+	// Only get the set of replicas which contain the specified user.
+	key := shardByUser(userID)
+	replicationSet, err := am.ring.Get(key, RingOp, nil, nil, nil)
+	if err != nil {
+		return nil, err
+	}
+
+	// We should only query state from other replicas, and not our own state.
+	addrs := replicationSet.GetAddressesWithout(am.ringLifecycler.GetInstanceAddr())
+
+	var (
+		resultsMtx      sync.Mutex
+		results         []*clusterpb.FullState
+		numUserNotFound int
+	)
+
+	// Note that the jobs swallow the errors - this is because we want to give each replica a chance to respond.
+	jobs := concurrency.CreateJobsFromStrings(addrs)
+	err = concurrency.ForEach(ctx, jobs, len(jobs), func(ctx context.Context, job interface{}) error {
+		addr := job.(string)
+		level.Debug(am.logger).Log("msg", "contacting replica for full state", "user", userID, "addr", addr)
+
+		c, err := am.alertmanagerClientsPool.GetClientFor(addr)
+		if err != nil {
+			level.Error(am.logger).Log("msg", "failed to get rpc client", "err", err)
+			return nil
+		}
+
+		resp, err := c.ReadState(user.InjectOrgID(ctx, userID), &alertmanagerpb.ReadStateRequest{})
+		if err != nil {
+			level.Error(am.logger).Log("msg", "rpc reading state from replica failed", "addr", addr, "user", userID, "err", err)
+			return nil
+		}
+
+		switch resp.Status {
+		case alertmanagerpb.READ_OK:
+			resultsMtx.Lock()
+			results = append(results, resp.State)
+			resultsMtx.Unlock()
+		case alertmanagerpb.READ_ERROR:
+			level.Error(am.logger).Log("msg", "error trying to read state", "addr", addr, "user", userID, "err", resp.Error)
+		case alertmanagerpb.READ_USER_NOT_FOUND:
+			level.Debug(am.logger).Log("msg", "user not found while trying to read state", "addr", addr, "user", userID)
+			resultsMtx.Lock()
+			numUserNotFound++
+			resultsMtx.Unlock()
+		default:
+			level.Error(am.logger).Log("msg", "unknown response trying to read state", "addr", addr, "user", userID)
+		}
+		return nil
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	// If all replicas reply with user not found, then we're not going to improve by retrying.
+	if numUserNotFound == len(addrs) {
+		level.Debug(am.logger).Log("msg", "user unknown in all replicas", "user", userID)
+		return []*clusterpb.FullState{}, nil
+	}
+
+	// We only require the state from a single replica, though we return as many as we were able to obtain.
+	if len(results) == 0 {
+		return nil, fmt.Errorf("failed to read state from any replica")
+	}
+
+	return results, nil
 }
 
 // UpdateState implements the Alertmanager service.

pkg/alertmanager/multitenant_test.go

@@ -1205,6 +1205,185 @@ func TestAlertmanager_StateReplicationWithSharding(t *testing.T) {
 	}
 }
 
+func TestAlertmanager_StateReplicationWithSharding_InitialSyncFromPeers(t *testing.T) {
+	tc := []struct {
+		name              string
+		replicationFactor int
+	}{
+		{
+			name:              "RF = 2",
+			replicationFactor: 2,
+		},
+		{
+			name:              "RF = 3",
+			replicationFactor: 3,
+		},
+	}
+
+	for _, tt := range tc {
+		t.Run(tt.name, func(t *testing.T) {
+			ctx := context.Background()
+			ringStore := consul.NewInMemoryClient(ring.GetCodec())
+			mockStore := prepareInMemoryAlertStore()
+			clientPool := newPassthroughAlertmanagerClientPool()
+			externalURL := flagext.URLValue{}
+			err := externalURL.Set("http://localhost:8080/alertmanager")
+			require.NoError(t, err)
+
+			var instances []*MultitenantAlertmanager
+			var instanceIDs []string
+			registries := util.NewUserRegistries()
+
+			// Create only two users - no need for more for these test cases.
+			for i := 1; i <= 2; i++ {
+				u := fmt.Sprintf("u-%d", i)
+				require.NoError(t, mockStore.SetAlertConfig(ctx, alertspb.AlertConfigDesc{
+					User:      u,
+					RawConfig: simpleConfigOne,
+					Templates: []*alertspb.TemplateDesc{},
+				}))
+			}
+
+			createInstance := func(i int) *MultitenantAlertmanager {
+				instanceIDs = append(instanceIDs, fmt.Sprintf("alertmanager-%d", i))
+				instanceID := fmt.Sprintf("alertmanager-%d", i)
+
+				amConfig := mockAlertmanagerConfig(t)
+				amConfig.ExternalURL = externalURL
+				amConfig.ShardingRing.ReplicationFactor = tt.replicationFactor
+				amConfig.ShardingRing.InstanceID = instanceID
+				amConfig.ShardingRing.InstanceAddr = fmt.Sprintf("127.0.0.%d", i)
+
+				// Do not check the ring topology changes or poll in an interval in this test (we explicitly sync alertmanagers).
+				amConfig.PollInterval = time.Hour
+				amConfig.ShardingRing.RingCheckPeriod = time.Hour
+
+				amConfig.ShardingEnabled = true
+
+				reg := prometheus.NewPedanticRegistry()
+				am, err := createMultitenantAlertmanager(amConfig, nil, nil, mockStore, ringStore, log.NewNopLogger(), reg)
+				require.NoError(t, err)
+
+				clientPool.servers[amConfig.ShardingRing.InstanceAddr+":0"] = am
+				am.alertmanagerClientsPool = clientPool
+
+				require.NoError(t, services.StartAndAwaitRunning(ctx, am))
+				t.Cleanup(func() {
+					services.StopAndAwaitTerminated(ctx, am) //nolint:errcheck
+				})
+
+				instances = append(instances, am)
+				instanceIDs = append(instanceIDs, instanceID)
+				registries.AddUserRegistry(instanceID, reg)
+
+				// Make sure the ring is settled.
+				{
+					ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
+					defer cancel()
+
+					// The alertmanager is ready to be tested once all instances are ACTIVE and the ring settles.
+					for _, am := range instances {
+						for _, id := range instanceIDs {
+							require.NoError(t, ring.WaitInstanceState(ctx, am.ring, id, ring.ACTIVE))
+						}
+					}
+				}
+
+				// Now that the ring has settled, sync configs with the instances.
+				require.NoError(t, am.loadAndSyncConfigs(ctx, reasonRingChange))
+
+				return am
+			}
+
+			writeSilence := func(i *MultitenantAlertmanager, userID string) {
+				silence := types.Silence{
+					Matchers: labels.Matchers{
+						{Name: "instance", Value: "prometheus-one"},
+					},
+					Comment:  "Created for a test case.",
+					StartsAt: time.Now(),
+					EndsAt:   time.Now().Add(time.Hour),
+				}
+				data, err := json.Marshal(silence)
+				require.NoError(t, err)
+
+				req := httptest.NewRequest(http.MethodPost, externalURL.String()+"/api/v2/silences", bytes.NewReader(data))
+				req.Header.Set("content-type", "application/json")
+				reqCtx := user.InjectOrgID(req.Context(), userID)
+				{
+					w := httptest.NewRecorder()
+					i.ServeHTTP(w, req.WithContext(reqCtx))
+
+					resp := w.Result()
+					body, _ := ioutil.ReadAll(resp.Body)
+					assert.Equal(t, http.StatusOK, w.Code)
+					require.Regexp(t, regexp.MustCompile(`{"silenceID":".+"}`), string(body))
+				}
+			}
+
+			checkSilence := func(i *MultitenantAlertmanager, userID string) {
+				req := httptest.NewRequest(http.MethodGet, externalURL.String()+"/api/v2/silences", nil)
+				req.Header.Set("content-type", "application/json")
+				reqCtx := user.InjectOrgID(req.Context(), userID)
+				{
+					w := httptest.NewRecorder()
+					i.ServeHTTP(w, req.WithContext(reqCtx))
+
+					resp := w.Result()
+					body, _ := ioutil.ReadAll(resp.Body)
+					assert.Equal(t, http.StatusOK, w.Code)
+					require.Regexp(t, regexp.MustCompile(`"comment":"Created for a test case."`), string(body))
+				}
+			}
+
+			// 1. Create the first instance and load the user configurations.
+			i1 := createInstance(1)
+
+			// 2. Create a silence in the first alertmanager instance and check we can read it.
+			writeSilence(i1, "u-1")
+			// 2.a. Check the silence was created (paranoia).
+			checkSilence(i1, "u-1")
+			// 2.b. Check the relevant metrics were updated.
+			{
+				metrics := registries.BuildMetricFamiliesPerUser()
+				assert.Equal(t, float64(1), metrics.GetSumOfGauges("cortex_alertmanager_silences"))
+				assert.Equal(t, float64(1), metrics.GetSumOfCounters("cortex_alertmanager_state_replication_total"))
+				assert.Equal(t, float64(0), metrics.GetSumOfCounters("cortex_alertmanager_state_replication_failed_total"))
+			}
+
+			// 3. Create a second instance. This should attempt to fetch the silence from the first.
+			i2 := createInstance(2)
+
+			// 3.a. Check the silence was fetched from the first instance successfully.
+			checkSilence(i2, "u-1")
+
+			// 3.b. Check the metrics: We should see the additional silences without any replication activity.
+			{
+				metrics := registries.BuildMetricFamiliesPerUser()
+				assert.Equal(t, float64(2), metrics.GetSumOfGauges("cortex_alertmanager_silences"))
+				assert.Equal(t, float64(1), metrics.GetSumOfCounters("cortex_alertmanager_state_replication_total"))
+				assert.Equal(t, float64(0), metrics.GetSumOfCounters("cortex_alertmanager_state_replication_failed_total"))
+			}
+
+			if tt.replicationFactor >= 3 {
+				// 4. When testing RF = 3, create a third instance, to test obtaining state from multiple places.
+				i3 := createInstance(3)
+
+				// 4.a. Check the silence was fetched one or both of the instances successfully.
+				checkSilence(i3, "u-1")
+
+				// 4.b. Check the metrics one more time. We should have three replicas of the silence.
+				{
+					metrics := registries.BuildMetricFamiliesPerUser()
+					assert.Equal(t, float64(3), metrics.GetSumOfGauges("cortex_alertmanager_silences"))
+					assert.Equal(t, float64(1), metrics.GetSumOfCounters("cortex_alertmanager_state_replication_total"))
+					assert.Equal(t, float64(0), metrics.GetSumOfCounters("cortex_alertmanager_state_replication_failed_total"))
+				}
+			}
+		})
+	}
+}
+
 // prepareInMemoryAlertStore builds and returns an in-memory alert store.
 func prepareInMemoryAlertStore() alertstore.AlertStore {
 	return bucketclient.NewBucketAlertStore(objstore.NewInMemBucket(), nil, log.NewNopLogger())
@@ -1251,6 +1430,10 @@ func (am *passthroughAlertmanagerClient) UpdateState(ctx context.Context, in *cl
 	return am.server.UpdateState(ctx, in)
 }
 
+func (am *passthroughAlertmanagerClient) ReadState(ctx context.Context, in *alertmanagerpb.ReadStateRequest, opts ...grpc.CallOption) (*alertmanagerpb.ReadStateResponse, error) {
+	return am.server.ReadState(ctx, in)
+}
+
 func (am *passthroughAlertmanagerClient) HandleRequest(context.Context, *httpgrpc.HTTPRequest, ...grpc.CallOption) (*httpgrpc.HTTPResponse, error) {
 	return nil, fmt.Errorf("unexpected call to HandleRequest")
 }

pkg/ring/replication_set.go

@@ -115,6 +115,18 @@ func (r ReplicationSet) GetAddresses() []string {
 	return addrs
 }
 
+// GetAddressesWithout returns the addresses of all instances within the replication set while
+// excluding the specified addresses.Returned slice order is not guaranteed.
+func (r ReplicationSet) GetAddressesWithout(exclude string) []string {
+	addrs := make([]string, 0, len(r.Instances))
+	for _, desc := range r.Instances {
+		if desc.Addr != exclude {
+			addrs = append(addrs, desc.Addr)
+		}
+	}
+	return addrs
+}
+
 // HasReplicationSetChanged returns true if two replications sets are the same (with possibly different timestamps),
 // false if they differ in any way (number of instances, instance states, tokens, zones, ...).
 func HasReplicationSetChanged(before, after ReplicationSet) bool {

pkg/ring/replication_set_test.go

@@ -39,6 +39,48 @@ func TestReplicationSet_GetAddresses(t *testing.T) {
 	}
 }
 
+func TestReplicationSet_GetAddressesWithout(t *testing.T) {
+	tests := map[string]struct {
+		rs       ReplicationSet
+		expected []string
+		exclude  string
+	}{
+		"should return an empty slice on empty replication set": {
+			rs:       ReplicationSet{},
+			expected: []string{},
+			exclude:  "127.0.0.1",
+		},
+		"non-matching exclusion, should return all addresses": {
+			rs: ReplicationSet{
+				Instances: []InstanceDesc{
+					{Addr: "127.0.0.1"},
+					{Addr: "127.0.0.2"},
+					{Addr: "127.0.0.3"},
+				},
+			},
+			expected: []string{"127.0.0.1", "127.0.0.2", "127.0.0.3"},
+			exclude:  "127.0.0.4",
+		},
+		"matching exclusion, should return non-excluded addresses": {
+			rs: ReplicationSet{
+				Instances: []InstanceDesc{
+					{Addr: "127.0.0.1"},
+					{Addr: "127.0.0.2"},
+					{Addr: "127.0.0.3"},
+				},
+			},
+			expected: []string{"127.0.0.1", "127.0.0.3"},
+			exclude:  "127.0.0.2",
+		},
+	}
+
+	for testName, testData := range tests {
+		t.Run(testName, func(t *testing.T) {
+			assert.ElementsMatch(t, testData.expected, testData.rs.GetAddressesWithout(testData.exclude))
+		})
+	}
+}
+
 var (
 	errFailure     = errors.New("failed")
 	errZoneFailure = errors.New("zone failed")