From 6b774527c3b987c57580e45838359f7a13732405 Mon Sep 17 00:00:00 2001
From: Benjamin Wang <benjamin.ahrtr@gmail.com>
Date: Thu, 25 Jan 2024 10:19:25 +0000
Subject: [PATCH] test: fix TestHashKVWhenCompacting: ensure all goroutine
 finished

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

diff --git a/mvcc/kvstore_test.go b/mvcc/kvstore_test.go
index 44baee21c97..f776b7fb691 100644
--- a/mvcc/kvstore_test.go
+++ b/mvcc/kvstore_test.go
@@ -546,6 +546,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++ {
@@ -558,6 +559,8 @@ func TestHashKVWhenCompacting(t *testing.T) {
 					t.Error(err)
 				}
 				select {
+				case <-stopc:
+					return
 				case <-donec:
 					return
 				case hashCompactc <- hashKVResult{hash, compactRev}:
@@ -581,6 +584,8 @@ func TestHashKVWhenCompacting(t *testing.T) {
 			}
 
 			select {
+			case <-stopc:
+				return
 			case <-donec:
 				return
 			default:
@@ -589,9 +594,20 @@ func TestHashKVWhenCompacting(t *testing.T) {
 	}()
 
 	// Compact the store in a goroutine, using revision 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)
@@ -602,12 +618,17 @@ func TestHashKVWhenCompacting(t *testing.T) {
 			time.Sleep(10 * time.Millisecond)
 		}
 	}()
+
 	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