From a9eb7d0dd6e8b8fd61e4110a107d257a0e96fc0b Mon Sep 17 00:00:00 2001
From: wooyang2018 <2201212838@stu.pku.edu.cn>
Date: Mon, 14 Oct 2024 22:51:35 +0800
Subject: [PATCH] Fix data race of outlier unit tests

---
 core/outlier/recycler_test.go | 16 +++++++++++-----
 core/outlier/retryer_test.go  | 20 ++++++++++++--------
 2 files changed, 23 insertions(+), 13 deletions(-)

diff --git a/core/outlier/recycler_test.go b/core/outlier/recycler_test.go
index ff46d4a9a..9316bf54f 100644
--- a/core/outlier/recycler_test.go
+++ b/core/outlier/recycler_test.go
@@ -32,6 +32,12 @@ func newRecycler(resource string, interval time.Duration) *Recycler {
 	}
 }
 
+func (r *Recycler) length() int {
+	r.mtx.Lock()
+	defer r.mtx.Unlock()
+	return len(r.status)
+}
+
 func addOutlierRuleForRecycler(resource string, seconds uint32) {
 	updateMux.Lock()
 	defer updateMux.Unlock()
@@ -66,7 +72,7 @@ func testRecycler(t *testing.T) {
 
 	recycler := newRecycler(resource, 4*time.Second)
 	recycler.scheduleNodes(nodes)
-	assert.Equal(t, len(nodes), len(recycler.status))
+	assert.Equal(t, len(nodes), recycler.length())
 
 	// Restore node0 after 2 seconds of simulation
 	time.AfterFunc(2*time.Second, func() {
@@ -101,9 +107,9 @@ func testRecyclerConcurrent(t *testing.T) {
 	wg.Wait()
 
 	// Check the status of nodes
-	assert.GreaterOrEqual(t, len(nodes), len(recycler.status))
+	assert.GreaterOrEqual(t, len(nodes), recycler.length())
 	recycler.scheduleNodes(nodes)
-	assert.Equal(t, len(nodes), len(recycler.status))
+	assert.Equal(t, len(nodes), recycler.length())
 
 	// Recover node0 and node1 after 2 seconds
 	time.AfterFunc(2*time.Second, func() {
@@ -162,9 +168,9 @@ func testRecyclerChConcurrent(t *testing.T) {
 	wg.Wait()
 
 	// Check the status of nodes
-	assert.GreaterOrEqual(t, len(nodes), len(recycler.status))
+	assert.GreaterOrEqual(t, len(nodes), recycler.length())
 	recycler.scheduleNodes(nodes)
-	assert.Equal(t, len(nodes), len(recycler.status))
+	assert.Equal(t, len(nodes), recycler.length())
 
 	// Recover node0 and node1 after 2 seconds
 	time.AfterFunc(2*time.Second, func() {
diff --git a/core/outlier/retryer_test.go b/core/outlier/retryer_test.go
index a04ef5e4b..3e4badb42 100644
--- a/core/outlier/retryer_test.go
+++ b/core/outlier/retryer_test.go
@@ -61,9 +61,7 @@ func (d *dummyCall) Check(address string) bool {
 		}
 		return false
 	}
-	fmt.Println(d.callCounts)
 	panic("Attempting to call an unregistered node address.")
-	return false
 }
 
 func (d *dummyCall) getRecoverCount() int {
@@ -131,6 +129,12 @@ func setNodeBreaker(resource string, node string, breaker *MockCircuitBreaker) {
 	nodeBreakers[resource][node] = breaker
 }
 
+func (r *Retryer) length() int {
+	r.mtx.Lock()
+	defer r.mtx.Unlock()
+	return len(r.counts)
+}
+
 // Construct two dummy node addresses: the first one recovers after the third check,
 // and the second one recovers after math.MaxInt32 checks. Observe the changes in the
 // circuit breaker and callCounts status for the first node before and after recovery.
@@ -154,7 +158,7 @@ func testRetryer(t *testing.T) {
 	for d.getRecoverCount() < 1 {
 		time.Sleep(minDuration)
 	}
-	assert.Equal(t, len(nodes)-1, len(retryer.counts))
+	assert.Equal(t, len(nodes)-1, retryer.length())
 	mockCB.AssertExpectations(t)
 }
 
@@ -175,7 +179,7 @@ func testRetryerConcurrent(t *testing.T) {
 			d.registerAddress(node, math.MaxInt32)
 		}
 	}
-	fmt.Println(d.callCounts)
+
 	addOutlierRuleForRetryer(resource, n, internal, d.Check)
 	retryer := getRetryerOfResource(resource)
 	numGoroutines := 10
@@ -194,15 +198,15 @@ func testRetryerConcurrent(t *testing.T) {
 	wg.Wait()
 
 	// Check the status of nodes
-	assert.GreaterOrEqual(t, len(nodes), len(retryer.counts))
+	assert.GreaterOrEqual(t, len(nodes), retryer.length())
 	retryer.scheduleNodes(nodes)
-	assert.Equal(t, len(nodes), len(retryer.counts))
+	assert.Equal(t, len(nodes), retryer.length())
 
 	minDuration := time.Duration(n * (n + 1) / 2 * internal * 1e6)
 	for d.getRecoverCount() < len(nodes)/2 {
 		time.Sleep(minDuration)
 	}
-	assert.Equal(t, len(nodes)/2, len(retryer.counts))
+	assert.Equal(t, len(nodes)/2, retryer.length())
 	for _, breaker := range mockCBs {
 		breaker.AssertExpectations(t)
 	}
@@ -229,7 +233,7 @@ func testRetryerCh(t *testing.T) {
 	for d.getRecoverCount() < 1 {
 		time.Sleep(minDuration)
 	}
-	assert.Equal(t, len(nodes)-1, len(retryer.counts))
+	assert.Equal(t, len(nodes)-1, retryer.length())
 	mockCB.AssertExpectations(t)
 }