Check if the endpoint has been removed before evicting

Signed-off-by: JmPotato <[email protected]>

pkg/utils/etcdutil/health_checker.go

@@ -273,7 +273,7 @@ func (checker *healthChecker) updateEvictedEps(lastEps, pickedEps []string) {
 	for _, ep := range pickedEps {
 		pickedSet[ep] = true
 	}
-	// Reset the count to 0 if it's in evictedEps but not in the pickedEps.
+	// Reset the count to 0 if it's in `evictedEps` but not in `pickedEps`.
 	checker.evictedEps.Range(func(key, value any) bool {
 		ep := key.(string)
 		count := value.(int)
@@ -286,18 +286,21 @@ func (checker *healthChecker) updateEvictedEps(lastEps, pickedEps []string) {
 		}
 		return true
 	})
-	// Find all endpoints which are in the lastEps but not in the pickedEps,
-	// and add them to the evictedEps.
+	// Find all endpoints which are in `lastEps` and `healthyClients` but not in `pickedEps`,
+	// and add them to the `evictedEps`.
 	for _, ep := range lastEps {
 		if pickedSet[ep] {
 			continue
 		}
+		if hc := checker.loadClient(ep); hc == nil {
+			continue
+		}
 		checker.evictedEps.Store(ep, 0)
 		log.Info("evicted etcd endpoint found",
 			zap.String("endpoint", ep),
 			zap.String("source", checker.source))
 	}
-	// Find all endpoints which are in both pickedEps and evictedEps to
+	// Find all endpoints which are in both `pickedEps` and `evictedEps` to
 	// increase their picked count.
 	for _, ep := range pickedEps {
 		if count, ok := checker.evictedEps.Load(ep); ok {

pkg/utils/etcdutil/health_checker_test.go

@@ -35,6 +35,8 @@ func check(re *require.Assertions, testCases []*testCase) {
 		probeCh := make(chan healthProbe, len(tc.healthProbes))
 		for _, probe := range tc.healthProbes {
 			probeCh <- probe
+			// Mock that all the endpoints are healthy.
+			checker.healthyClients.LoadOrStore(probe.ep, &healthyClient{})
 		}
 		close(probeCh)
 		// Pick and filter the endpoints.
@@ -361,3 +363,27 @@ func TestLatencyPick(t *testing.T) {
 	}
 	check(re, testCases)
 }
+
+func TestUpdateEvictedEpsAfterRemoval(t *testing.T) {
+	re := require.New(t)
+	var (
+		checker   = &healthChecker{}
+		lastEps   = []string{"A", "B", "C"}
+		pickedEps = []string{"A", "C"}
+	)
+	// All endpoints are healthy.
+	for _, ep := range lastEps {
+		checker.healthyClients.Store(ep, &healthyClient{})
+	}
+	checker.updateEvictedEps(lastEps, pickedEps)
+	// B should be evicted.
+	_, ok := checker.evictedEps.Load("B")
+	re.True(ok)
+	// Remove the endpoint B to mock member removal.
+	checker.healthyClients.Delete("B")
+	checker.evictedEps.Delete("B")
+	checker.updateEvictedEps(lastEps, pickedEps)
+	// B should not be evicted since it has been removed.
+	_, ok = checker.evictedEps.Load("B")
+	re.False(ok)
+}