Apply review suggestions

Signed-off-by: Florent Poinsard <florent.poinsard@outlook.fr>

go/vt/vtgate/endtoend/connectiondrain/main_test.go → go/test/endtoend/vtgate/connectiondrain/main_test.go

@@ -21,7 +21,6 @@ import (
 	_ "embed"
 	"flag"
 	"os"
-	"sync"
 	"testing"
 	"time"
 
@@ -176,24 +175,13 @@ func TestConnectionDrainOnTermTimeout(t *testing.T) {
 	require.NoError(t, err)
 
 	// Run a busy query that returns only after the onterm_timeout is reached, this should fail when we reach the timeout
-	wg := sync.WaitGroup{}
-	wg.Add(1)
-	go func() {
-		defer wg.Done()
-		var err error
-		_, err = vtConn.ExecuteFetch("select sleep(40)", 1, false)
-		require.Error(t, err)
-	}()
-
-	// Sleeping 40 seconds here is already plenty of time, and we will for sure reach the onterm_timeout of 30s
-	time.Sleep(40 * time.Second)
+	_, err = vtConn.ExecuteFetch("select sleep(40)", 1, false)
+	require.Error(t, err)
 
 	// Running a query after we have reached the onterm_timeout should fail
 	_, err = vtConn2.ExecuteFetch("select id from t1", 1, false)
 	require.Error(t, err)
 
-	wg.Wait()
-
 	// By now vtgate will be shutdown becaused it reached its onterm_timeout, despite idle connections still being opened
 	require.True(t, clusterInstance.VtgateProcess.IsShutdown())
 }

go/vt/vtgate/plugin_mysql_server.go

@@ -626,10 +626,15 @@ func (srv *mysqlServer) shutdownMysqlProtocolAndDrain() {
 	if srv.sigChan != nil {
 		signal.Stop(srv.sigChan)
 	}
+	setListenerToNil := func() {
+		srv.tcpListener = nil
+		srv.unixListener = nil
+	}
 
 	if mysqlDrainOnTerm {
 		stopListener(srv.unixListener, false)
 		stopListener(srv.tcpListener, false)
+		setListenerToNil()
 		// We wait for connected clients to drain by themselves or to run into the onterm timeout
 		log.Infof("Starting drain loop, waiting for all clients to disconnect")
 		reported := time.Now()
@@ -645,6 +650,7 @@ func (srv *mysqlServer) shutdownMysqlProtocolAndDrain() {
 
 	stopListener(srv.unixListener, true)
 	stopListener(srv.tcpListener, true)
+	setListenerToNil()
 	if busy := srv.vtgateHandle.busyConnections.Load(); busy > 0 {
 		log.Infof("Waiting for all client connections to be idle (%d active)...", busy)
 		start := time.Now()
@@ -671,7 +677,6 @@ func stopListener(listener *mysql.Listener, shutdown bool) {
 	} else {
 		listener.Close()
 	}
-	listener = nil
 }
 
 func (srv *mysqlServer) rollbackAtShutdown() {

test/config.json

@@ -500,6 +500,15 @@
 			"RetryMax": 2,
 			"Tags": []
 		},
+		"vtgate_connectiondrain": {
+			"File": "unused.go",
+			"Args": ["vitess.io/vitess/go/test/endtoend/vtgate/connectiondrain"],
+			"Command": [],
+			"Manual": false,
+			"Shard": "vtgate_general_heavy",
+			"RetryMax": 2,
+			"Tags": []
+		},
 		"vtgate_queries_derived": {
 			"File": "unused.go",
 			"Args": ["vitess.io/vitess/go/test/endtoend/vtgate/queries/derived"],