metricreader: fix that sometimes the metric history is not purged (#639)

Author: djshow832

pkg/balance/metricsreader/backend_reader.go

@@ -49,7 +49,7 @@ const (
 )
 
 var (
-	errReadFromOwner = errors.New("read metrics from owner failed")
+	errReadMetrics = errors.New("read backend metrics failed")
 )
 
 type backendHistory struct {
@@ -85,17 +85,18 @@ type BackendReader struct {
 func NewBackendReader(lg *zap.Logger, cfgGetter config.ConfigGetter, httpCli *http.Client, etcdCli *clientv3.Client,
 	backendFetcher TopologyFetcher, cfg *config.HealthCheck) *BackendReader {
 	return &BackendReader{
-		queryRules:     make(map[string]QueryRule),
-		queryResults:   make(map[string]QueryResult),
-		history:        make(map[string]map[string]backendHistory),
-		lg:             lg,
-		cfgGetter:      cfgGetter,
-		backendFetcher: backendFetcher,
-		cfg:            cfg,
-		wgp:            waitgroup.NewWaitGroupPool(goPoolSize, goMaxIdle),
-		electionCfg:    elect.DefaultElectionConfig(sessionTTL),
-		etcdCli:        etcdCli,
-		httpCli:        httpCli,
+		queryRules:        make(map[string]QueryRule),
+		queryResults:      make(map[string]QueryResult),
+		history:           make(map[string]map[string]backendHistory),
+		lg:                lg,
+		cfgGetter:         cfgGetter,
+		backendFetcher:    backendFetcher,
+		cfg:               cfg,
+		wgp:               waitgroup.NewWaitGroupPool(goPoolSize, goMaxIdle),
+		electionCfg:       elect.DefaultElectionConfig(sessionTTL),
+		etcdCli:           etcdCli,
+		httpCli:           httpCli,
+		marshalledHistory: []byte{},
 	}
 }
 
@@ -167,30 +168,39 @@ func (br *BackendReader) ReadMetrics(ctx context.Context) error {
 	if err != nil {
 		return err
 	}
+
+	// If self is a owner, read the backends that are not read by any other owners.
 	var errs []error
-	for _, owner := range owners {
-		if owner == br.election.ID() {
-			continue
+	var backendLabels []string
+	if br.isOwner.Load() {
+		if idx := slices.Index(zones, zone); idx >= 0 {
+			zones = slices.Delete(zones, idx, idx+1)
 		}
-		if err = br.readFromOwner(ctx, owner); err != nil {
+		backendLabels, err = br.readFromBackends(ctx, zones)
+		if err != nil {
 			errs = append(errs, err)
 		}
 	}
 
-	// If self is a owner, read the backends that are not read by any other owners.
-	if br.isOwner.Load() {
-		if idx := slices.Index(zones, zone); idx >= 0 {
-			zones = slices.Delete(zones, idx, idx+1)
+	for _, owner := range owners {
+		if owner == br.election.ID() {
+			continue
 		}
-		if err := br.readFromBackends(ctx, zones); err != nil {
-			return err
+		if err = br.readFromOwner(ctx, owner); err != nil {
+			errs = append(errs, err)
 		}
 	}
 
 	// Purge expired history.
 	br.purgeHistory()
+	// Marshal backend history for other members to query.
+	if err := br.marshalHistory(backendLabels); err != nil {
+		br.lg.Error("marshal backend history failed", zap.Any("addrs", backendLabels), zap.Error(err))
+	}
+	// Generate query result for all backends.
+	br.history2QueryResult()
 	if len(errs) > 0 {
-		return errors.Collect(errReadFromOwner, errs...)
+		return errors.Collect(errReadMetrics, errs...)
 	}
 	return nil
 }
@@ -259,17 +269,17 @@ func (br *BackendReader) queryAllOwners(ctx context.Context) (zones, owners []st
 // 1. In k8s, the zone is not set at startup and then is set by HTTP API, so there may temporarily exist both global and zonal owners.
 // 2. Some backends may not be in the same zone with any owner. E.g. there are only 2 TiProxy in a 3-AZ cluster.
 // In any way, the owner queries the backends that are not queried by other owners.
-func (br *BackendReader) readFromBackends(ctx context.Context, excludeZones []string) error {
+func (br *BackendReader) readFromBackends(ctx context.Context, excludeZones []string) ([]string, error) {
 	addrs, err := br.getBackendAddrs(ctx, excludeZones)
 	if err != nil {
-		return err
+		return nil, err
 	}
 	if len(addrs) == 0 {
-		return nil
+		return nil, nil
 	}
 	allNames := br.collectAllNames()
 	if len(allNames) == 0 {
-		return nil
+		return nil, nil
 	}
 
 	backendLabels := make([]string, 0, len(addrs))
@@ -298,12 +308,7 @@ func (br *BackendReader) readFromBackends(ctx context.Context, excludeZones []st
 		}(addrs[i], backendLabels[i])
 	}
 	br.wgp.Wait()
-
-	br.history2QueryResult()
-	if err := br.marshalHistory(backendLabels); err != nil {
-		br.lg.Error("marshal backend history failed", zap.Any("addrs", backendLabels), zap.Error(err))
-	}
-	return nil
+	return backendLabels, nil
 }
 
 func (br *BackendReader) collectAllNames() []string {
@@ -473,9 +478,6 @@ func (br *BackendReader) readFromOwner(ctx context.Context, ownerAddr string) er
 
 	// If this instance becomes the owner in the next round, it can reuse the history.
 	br.mergeHistory(newHistory)
-
-	// Generate query result for all backends.
-	br.history2QueryResult()
 	return nil
 }
 
@@ -506,8 +508,6 @@ func (br *BackendReader) mergeHistory(newHistory map[string]map[string]backendHi
 			ruleHistory[backend] = backendHistory
 		}
 	}
-	// avoid that the stale history is returned to other members when it just becomes the owner
-	br.marshalledHistory = nil
 }
 
 // marshalHistory marshals the backends that are read by this owner. The marshaled data will be returned to other members.
@@ -516,12 +516,14 @@ func (br *BackendReader) marshalHistory(backends []string) error {
 	defer br.Unlock()
 
 	filteredHistory := make(map[string]map[string]backendHistory, len(br.queryRules))
-	for ruleKey, ruleHistory := range br.history {
-		filteredRuleHistory := make(map[string]backendHistory, len(backends))
-		filteredHistory[ruleKey] = filteredRuleHistory
-		for backend, backendHistory := range ruleHistory {
-			if slices.Contains(backends, backend) {
-				filteredRuleHistory[backend] = backendHistory
+	if len(backends) > 0 {
+		for ruleKey, ruleHistory := range br.history {
+			filteredRuleHistory := make(map[string]backendHistory, len(backends))
+			filteredHistory[ruleKey] = filteredRuleHistory
+			for backend, backendHistory := range ruleHistory {
+				if slices.Contains(backends, backend) {
+					filteredRuleHistory[backend] = backendHistory
+				}
 			}
 		}
 	}

pkg/balance/metricsreader/backend_reader_test.go

@@ -886,16 +886,19 @@ func TestQueryBackendConcurrently(t *testing.T) {
 		addRule(i)
 	}
 
-	// start a goroutine to query metrics
+	// start a goroutine to query metrics from backends
 	var wg waitgroup.WaitGroup
 	childCtx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
 	// fill the initial query result to ensure the result is always non-empty
-	err := br.readFromBackends(childCtx, nil)
+	backends, err := br.readFromBackends(childCtx, nil)
 	require.NoError(t, err)
+	require.Len(t, backends, initialBackends)
+	br.history2QueryResult()
 	wg.Run(func() {
 		for childCtx.Err() == nil {
-			err := br.readFromBackends(childCtx, nil)
+			backends, err := br.readFromBackends(childCtx, nil)
 			require.NoError(t, err)
+			require.Len(t, backends, initialBackends)
 			br.purgeHistory()
 		}
 	})
@@ -951,7 +954,7 @@ func TestQueryBackendConcurrently(t *testing.T) {
 		}
 	})
 
-	// start a goroutine to marshal history
+	// start a goroutine to query marshalled history
 	wg.Run(func() {
 		for childCtx.Err() == nil {
 			select {
@@ -1223,7 +1226,7 @@ func TestElection(t *testing.T) {
 			},
 		},
 	}
-	hitoryText, err := json.Marshal(history)
+	historyText, err := json.Marshal(history)
 	require.NoError(t, err)
 
 	// setup rule
@@ -1244,7 +1247,7 @@ func TestElection(t *testing.T) {
 	ownerPort := ownerHttpHandler.Start()
 	addr := net.JoinHostPort("127.0.0.1", strconv.Itoa(ownerPort))
 	ownerFunc := func(_ string) string {
-		return string(hitoryText)
+		return string(historyText)
 	}
 	ownerHttpHandler.getRespBody.Store(&ownerFunc)
 	t.Cleanup(ownerHttpHandler.Close)
@@ -1276,11 +1279,20 @@ func TestElection(t *testing.T) {
 	ts := time.Now()
 	err = br.ReadMetrics(context.Background())
 	require.NoError(t, err)
+	// check that the query result is updated
 	qr := br.GetQueryResult("rule_id1")
 	require.False(t, qr.Empty())
 	require.Equal(t, model.SampleValue(100.0), qr.Value.(model.Vector)[0].Value)
 	require.GreaterOrEqual(t, qr.UpdateTime, ts)
 	ts = qr.UpdateTime
+	// check that the marshalled history is empty
+	var unmarshalled map[string]map[string]backendHistory
+	marshalledHistory := br.GetBackendMetrics()
+	if len(marshalledHistory) > 0 {
+		err = json.Unmarshal(marshalledHistory, &unmarshalled)
+		require.NoError(t, err)
+		require.Len(t, unmarshalled, 0)
+	}
 
 	// test owner
 	suite.delKV(ownerKey)
@@ -1290,10 +1302,16 @@ func TestElection(t *testing.T) {
 	}, 3*time.Second, 10*time.Millisecond)
 	err = br.ReadMetrics(context.Background())
 	require.NoError(t, err)
+	// check that the query result is updated
 	qr = br.GetQueryResult("rule_id1")
 	require.False(t, qr.Empty())
 	require.Equal(t, model.SampleValue(80.0), qr.Value.(model.Vector)[0].Value)
 	require.GreaterOrEqual(t, qr.UpdateTime, ts)
+	// check that the marshalled history is updated
+	marshalledHistory = br.GetBackendMetrics()
+	err = json.Unmarshal(marshalledHistory, &unmarshalled)
+	require.NoError(t, err)
+	require.Equal(t, br.history, unmarshalled)
 }
 
 func setupTypicalBackendListener(t *testing.T, respBody string) (backendPort int, infos map[string]*infosync.TiDBTopologyInfo) {