From c647614b6ab8c08d37aa2f1ff5a66c7f5d4b0adc Mon Sep 17 00:00:00 2001
From: Benjamin Wang <benjamin.ahrtr@gmail.com>
Date: Wed, 24 Jan 2024 13:32:02 +0000
Subject: [PATCH] test: fix TestHashKVWhenCompacting: ensure all goroutine
 finished

Signed-off-by: Benjamin Wang <benjamin.ahrtr@gmail.com>
---
 server/storage/mvcc/kvstore_test.go | 24 ++++++++++++++++++++++--
 1 file changed, 22 insertions(+), 2 deletions(-)

diff --git a/server/storage/mvcc/kvstore_test.go b/server/storage/mvcc/kvstore_test.go
index 3e413596c..58e23249d 100644
--- a/server/storage/mvcc/kvstore_test.go
+++ b/server/storage/mvcc/kvstore_test.go
@@ -562,6 +562,7 @@ func TestHashKVWhenCompacting(t *testing.T) {
 	hashCompactc := make(chan hashKVResult, 1)
 	var wg sync.WaitGroup
 	donec := make(chan struct{})
+	stopc := make(chan struct{})
 
 	// Call HashByRev(10000) in multiple goroutines until donec is closed
 	for i := 0; i < 10; i++ {
@@ -574,6 +575,8 @@ func TestHashKVWhenCompacting(t *testing.T) {
 					t.Error(err)
 				}
 				select {
+				case <-stopc:
+					return
 				case <-donec:
 					return
 				case hashCompactc <- hashKVResult{hash.Hash, hash.CompactRevision}:
@@ -597,6 +600,8 @@ func TestHashKVWhenCompacting(t *testing.T) {
 				if r.hash != revHash[r.compactRev] {
 					t.Errorf("Hashes differ (current %v) != (saved %v)", r.hash, revHash[r.compactRev])
 				}
+			case <-stopc:
+				return
 			case <-donec:
 				return
 			}
@@ -604,9 +609,20 @@ func TestHashKVWhenCompacting(t *testing.T) {
 	}()
 
 	// Compact the store in a goroutine, using RevisionTombstone 9900 to 10000 and close donec when finished
+	wg.Add(1)
 	go func() {
-		defer close(donec)
+		defer func() {
+			close(donec)
+			wg.Done()
+		}()
+
 		for i := 100; i >= 0; i-- {
+			select {
+			case <-stopc:
+				return
+			default:
+			}
+
 			_, err := s.Compact(traceutil.TODO(), int64(rev-i))
 			if err != nil {
 				t.Error(err)
@@ -620,10 +636,14 @@ func TestHashKVWhenCompacting(t *testing.T) {
 
 	select {
 	case <-donec:
-		wg.Wait()
 	case <-time.After(10 * time.Second):
+		close(stopc)
+		wg.Wait()
 		testutil.FatalStack(t, "timeout")
 	}
+
+	close(stopc)
+	wg.Wait()
 }
 
 // TestHashKVWithCompactedAndFutureRevisions ensures that HashKV returns a correct hash when called