test: fix TestHashKVWhenCompacting: ensure all goroutine finished

Signed-off-by: Benjamin Wang <[email protected]>

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,16 +600,29 @@ 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
 			}
 		}
 	}()
 
 	// 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