database/sql: don't try to synctest.Wait in benchmarks

Benchmarks shouldn't use synctest. Avoid a synctest.Wait in database
shutdown when in benchmarks.

Fixes #78373

Change-Id: I14755f8eefb3ab5c69cc660b1e884bbf6a6a6964
Reviewed-on: https://go-review.googlesource.com/c/go/+/760400
Auto-Submit: Damien Neil <dneil@google.com>
Reviewed-by: Michael Pratt <mpratt@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>

1 files changed, 15 insertions, 9 deletions

diff --git a/src/database/sql/sql_test.go b/src/database/sql/sql_test.go
index 44de5e70b8..0f5268054c 100644
--- a/src/database/sql/sql_test.go
+++ b/src/database/sql/sql_test.go
@@ -184,8 +184,21 @@ func closeDB(t testing.TB, db *DB) {
 	}
 
 	// Connections close asynchronously; wait for them to finish doing so.
-	synctest.Wait()
-	if numOpen := db.numOpenConns(); numOpen != 0 {
+	numOpenConns := func() int {
+		db.mu.Lock()
+		defer db.mu.Unlock()
+		return db.numOpen
+	}
+	if _, ok := t.(*testing.B); ok {
+		// We don't use synctest in benchmarks, so just poll.
+		deadline := time.Now().Add(5 * time.Second)
+		for numOpenConns() > 0 && time.Now().Before(deadline) {
+			time.Sleep(1 * time.Millisecond)
+		}
+	} else {
+		synctest.Wait()
+	}
+	if numOpen := numOpenConns(); numOpen != 0 {
 		t.Fatalf("%d connections still open after closing DB", numOpen)
 	}
 }
@@ -213,13 +226,6 @@ func (db *DB) numFreeConns() int {
 	return len(db.freeConn)
 }
 
-func (db *DB) numOpenConns() int {
-	synctest.Wait()
-	db.mu.Lock()
-	defer db.mu.Unlock()
-	return db.numOpen
-}
-
 // clearAllConns closes all connections in db.
 func (db *DB) clearAllConns(t *testing.T) {
 	db.SetMaxIdleConns(0)