From f1b0d532813bfa8bda4eb8c2c11b9ff7c62c2912 Mon Sep 17 00:00:00 2001
From: Florent Poinsard <florent.poinsard@outlook.fr>
Date: Tue, 2 Jul 2024 15:56:43 -0400
Subject: [PATCH] Apply review suggestions

Signed-off-by: Florent Poinsard <florent.poinsard@outlook.fr>
---
 .../vtgate}/connectiondrain/main_test.go         | 16 ++--------------
 .../endtoend/vtgate}/connectiondrain/schema.sql  |  0
 go/vt/vtgate/plugin_mysql_server.go              |  7 ++++++-
 test/config.json                                 |  9 +++++++++
 4 files changed, 17 insertions(+), 15 deletions(-)
 rename go/{vt/vtgate/endtoend => test/endtoend/vtgate}/connectiondrain/main_test.go (94%)
 rename go/{vt/vtgate/endtoend => test/endtoend/vtgate}/connectiondrain/schema.sql (100%)

diff --git a/go/vt/vtgate/endtoend/connectiondrain/main_test.go b/go/test/endtoend/vtgate/connectiondrain/main_test.go
similarity index 94%
rename from go/vt/vtgate/endtoend/connectiondrain/main_test.go
rename to go/test/endtoend/vtgate/connectiondrain/main_test.go
index 33a68f08645..32807b22bde 100644
--- a/go/vt/vtgate/endtoend/connectiondrain/main_test.go
+++ b/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())
 }
diff --git a/go/vt/vtgate/endtoend/connectiondrain/schema.sql b/go/test/endtoend/vtgate/connectiondrain/schema.sql
similarity index 100%
rename from go/vt/vtgate/endtoend/connectiondrain/schema.sql
rename to go/test/endtoend/vtgate/connectiondrain/schema.sql
diff --git a/go/vt/vtgate/plugin_mysql_server.go b/go/vt/vtgate/plugin_mysql_server.go
index 5b460ff3388..55954013862 100644
--- a/go/vt/vtgate/plugin_mysql_server.go
+++ b/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() {
diff --git a/test/config.json b/test/config.json
index 713faf97024..8b48ccc3ec0 100644
--- a/test/config.json
+++ b/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"],