From 56bc3098f4ed42b272a45c246dcd42d28d89a69a Mon Sep 17 00:00:00 2001
From: Michael Anthony Knyszek <mknyszek@google.com>
Date: Fri, 10 Jun 2022 18:47:37 +0000
Subject: [PATCH] sync: improve linearity test robustness

This change improves the robustness of the locklinear test in the
following ways:

* It removes allocations from the timing, which may be very variable if
  we're unlucky.
* It ensures that goroutines are properly cleaned up before the test
  function returns, reducing the chance that they bleed into repeat
  attempts. It also stops timing before this cleanup.

Fixes #32986.

Change-Id: I3a8096e6922f23d899ad602e2845bdfc639ed742
Reviewed-on: https://go-review.googlesource.com/c/go/+/409894
Reviewed-by: Michael Pratt <mpratt@google.com>
---
 src/sync/mutex_test.go | 36 +++++++++++++++++++++++++++++-------
 1 file changed, 29 insertions(+), 7 deletions(-)

diff --git a/src/sync/mutex_test.go b/src/sync/mutex_test.go
index 9a4187c6726..8e34b024598 100644
--- a/src/sync/mutex_test.go
+++ b/src/sync/mutex_test.go
@@ -340,16 +340,25 @@ func TestMutexLinearOne(t *testing.T) {
 	testenv.CheckLinear(t, func(scale float64) func(*testing.B) {
 		n := int(1000 * scale)
 		return func(b *testing.B) {
-			ch := make(chan int)
+			ch := make(chan struct{})
 			locks := make([]RWMutex, runtimeSemaHashTableSize+1)
+
+			b.ResetTimer()
+
+			var wgStart, wgFinish WaitGroup
 			for i := 0; i < n; i++ {
+				wgStart.Add(1)
+				wgFinish.Add(1)
 				go func() {
+					wgStart.Done()
 					locks[0].Lock()
-					ch <- 1
+					ch <- struct{}{}
+					wgFinish.Done()
 				}()
 			}
-			time.Sleep(1 * time.Millisecond)
+			wgStart.Wait()
 
+			wgFinish.Add(1)
 			go func() {
 				for j := 0; j < n; j++ {
 					locks[1].Lock()
@@ -358,6 +367,7 @@ func TestMutexLinearOne(t *testing.T) {
 					runtime.Gosched()
 					locks[runtimeSemaHashTableSize].Unlock()
 				}
+				wgFinish.Done()
 			}()
 
 			for j := 0; j < n; j++ {
@@ -368,10 +378,14 @@ func TestMutexLinearOne(t *testing.T) {
 				locks[runtimeSemaHashTableSize].Unlock()
 			}
 
+			b.StopTimer()
+
 			for i := 0; i < n; i++ {
 				<-ch
 				locks[0].Unlock()
 			}
+
+			wgFinish.Wait()
 		}
 	})
 }
@@ -387,17 +401,21 @@ func TestMutexLinearMany(t *testing.T) {
 		return func(b *testing.B) {
 			locks := make([]RWMutex, n*runtimeSemaHashTableSize+1)
 
-			var wg WaitGroup
+			b.ResetTimer()
+
+			var wgStart, wgFinish WaitGroup
 			for i := 0; i < n; i++ {
-				wg.Add(1)
+				wgStart.Add(1)
+				wgFinish.Add(1)
 				go func(i int) {
 					locks[(i+1)*runtimeSemaHashTableSize].Lock()
-					wg.Done()
+					wgStart.Done()
 					locks[(i+1)*runtimeSemaHashTableSize].Lock()
 					locks[(i+1)*runtimeSemaHashTableSize].Unlock()
+					wgFinish.Done()
 				}(i)
 			}
-			wg.Wait()
+			wgStart.Wait()
 
 			go func() {
 				for j := 0; j < n; j++ {
@@ -417,9 +435,13 @@ func TestMutexLinearMany(t *testing.T) {
 				locks[0].Unlock()
 			}
 
+			b.StopTimer()
+
 			for i := 0; i < n; i++ {
 				locks[(i+1)*runtimeSemaHashTableSize].Unlock()
 			}
+
+			wgFinish.Wait()
 		}
 	})
 }