Enforce rollback for each connection when shutdown occurs

go/test/endtoend/cluster/vtgate_process.go

@@ -244,7 +244,6 @@ func (vtgate *VtgateProcess) GetVars() (map[string]interface{}, error) {
 			return nil, fmt.Errorf("not able to parse response body")
 		}
 		return resultMap, nil
-	} else {
-		return nil, fmt.Errorf("unsuccessful response")
 	}
+	return nil, fmt.Errorf("unsuccessful response")
 }

go/test/endtoend/vtgate/transaction/rollback/txn_rollback_shutdown_test.go

@@ -0,0 +1,142 @@
+/*
+Copyright 2019 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package rollback
+
+import (
+	"context"
+	"flag"
+	"fmt"
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"vitess.io/vitess/go/mysql"
+	"vitess.io/vitess/go/sqltypes"
+	"vitess.io/vitess/go/test/endtoend/cluster"
+)
+
+var (
+	clusterInstance *cluster.LocalProcessCluster
+	vtParams        mysql.ConnParams
+	keyspaceName    = "ks"
+	cell            = "zone1"
+	hostname        = "localhost"
+	sqlSchema       = `
+	create table buffer(
+		id BIGINT NOT NULL,
+		msg VARCHAR(64) NOT NULL,
+		PRIMARY KEY (id)
+	) Engine=InnoDB;`
+)
+
+func TestMain(m *testing.M) {
+	flag.Parse()
+
+	exitcode, err := func() (int, error) {
+		clusterInstance = cluster.NewCluster(cell, hostname)
+		defer clusterInstance.Teardown()
+
+		// Reserve vtGate port in order to pass it to vtTablet
+		clusterInstance.VtgateGrpcPort = clusterInstance.GetAndReservePort()
+
+		// Start topo server
+		if err := clusterInstance.StartTopo(); err != nil {
+			return 1, err
+		}
+
+		// Start keyspace
+		keyspace := &cluster.Keyspace{
+			Name:      keyspaceName,
+			SchemaSQL: sqlSchema,
+		}
+		if err := clusterInstance.StartUnshardedKeyspace(*keyspace, 1, false); err != nil {
+			return 1, err
+		}
+
+		// Set a short onterm timeout so the test goes faster.
+		clusterInstance.VtGateExtraArgs = []string{"-onterm_timeout", "1s"}
+		if err := clusterInstance.StartVtgate(); err != nil {
+			return 1, err
+		}
+		vtParams = mysql.ConnParams{
+			Host: clusterInstance.Hostname,
+			Port: clusterInstance.VtgateMySQLPort,
+		}
+
+		return m.Run(), nil
+	}()
+	if err != nil {
+		fmt.Printf("%v\n", err)
+		os.Exit(1)
+	} else {
+		os.Exit(exitcode)
+	}
+}
+
+func exec(t *testing.T, conn *mysql.Conn, query string) *sqltypes.Result {
+	t.Helper()
+	qr, err := conn.ExecuteFetch(query, 1000, true)
+	if err != nil {
+		t.Fatal(err)
+	}
+	return qr
+}
+
+func TestTransactionRollBackWhenShutDown(t *testing.T) {
+	ctx := context.Background()
+	conn, err := mysql.Connect(ctx, &vtParams)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer conn.Close()
+
+	exec(t, conn, "insert into buffer(id, msg) values(3,'mark')")
+	exec(t, conn, "insert into buffer(id, msg) values(4,'doug')")
+
+	// start an incomplete transaction
+	exec(t, conn, "begin")
+	exec(t, conn, "insert into buffer(id, msg) values(33,'mark')")
+
+	// Enforce a restart to enforce rollback
+	if err = clusterInstance.ReStartVtgate(); err != nil {
+		t.Errorf("Fail to re-start vtgate: %v", err)
+	}
+
+	want := ""
+
+	// Make a new mysql connection to vtGate
+	vtParams = mysql.ConnParams{
+		Host: clusterInstance.Hostname,
+		Port: clusterInstance.VtgateMySQLPort,
+	}
+	conn2, err := mysql.Connect(ctx, &vtParams)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer conn2.Close()
+
+	vtParams = mysql.ConnParams{
+		Host: clusterInstance.Hostname,
+		Port: clusterInstance.VtgateMySQLPort,
+	}
+
+	// Verify that rollback worked
+	qr := exec(t, conn2, "select id from buffer where msg='mark'")
+	got := fmt.Sprintf("%v", qr.Rows)
+	want = `[[INT64(3)]]`
+	assert.Equal(t, want, got)
+}

go/vt/vtgate/plugin_mysql_server.go

@@ -22,6 +22,7 @@ import (
 	"net"
 	"os"
 	"regexp"
+	"sync"
 	"sync/atomic"
 	"syscall"
 	"time"
@@ -72,16 +73,29 @@ var (
 // vtgateHandler implements the Listener interface.
 // It stores the Session in the ClientData of a Connection.
 type vtgateHandler struct {
-	vtg *VTGate
+	mu sync.Mutex
+
+	vtg         *VTGate
+	connections map[*mysql.Conn]bool
 }
 
 func newVtgateHandler(vtg *VTGate) *vtgateHandler {
 	return &vtgateHandler{
-		vtg: vtg,
+		vtg:         vtg,
+		connections: make(map[*mysql.Conn]bool),
 	}
 }
 
 func (vh *vtgateHandler) NewConnection(c *mysql.Conn) {
+	vh.mu.Lock()
+	defer vh.mu.Unlock()
+	vh.connections[c] = true
+}
+
+func (vh *vtgateHandler) numConnections() int {
+	vh.mu.Lock()
+	defer vh.mu.Unlock()
+	return len(vh.connections)
 }
 
 func (vh *vtgateHandler) ComResetConnection(c *mysql.Conn) {
@@ -98,6 +112,12 @@ func (vh *vtgateHandler) ComResetConnection(c *mysql.Conn) {
 
 func (vh *vtgateHandler) ConnectionClosed(c *mysql.Conn) {
 	// Rollback if there is an ongoing transaction. Ignore error.
+	defer func() {
+		vh.mu.Lock()
+		defer vh.mu.Unlock()
+		delete(vh.connections, c)
+	}()
+
 	var ctx context.Context
 	var cancel context.CancelFunc
 	if *mysqlQueryTimeout != 0 {
@@ -310,6 +330,8 @@ func (vh *vtgateHandler) session(c *mysql.Conn) *vtgatepb.Session {
 var mysqlListener *mysql.Listener
 var mysqlUnixListener *mysql.Listener
 
+var vtgateHandle *vtgateHandler
+
 // initiMySQLProtocol starts the mysql protocol.
 // It should be called only once in a process.
 func initMySQLProtocol() {
@@ -338,9 +360,9 @@ func initMySQLProtocol() {
 
 	// Create a Listener.
 	var err error
-	vh := newVtgateHandler(rpcVTGate)
+	vtgateHandle = newVtgateHandler(rpcVTGate)
 	if *mysqlServerPort >= 0 {
-		mysqlListener, err = mysql.NewListener(*mysqlTCPVersion, net.JoinHostPort(*mysqlServerBindAddress, fmt.Sprintf("%v", *mysqlServerPort)), authServer, vh, *mysqlConnReadTimeout, *mysqlConnWriteTimeout, *mysqlProxyProtocol)
+		mysqlListener, err = mysql.NewListener(*mysqlTCPVersion, net.JoinHostPort(*mysqlServerBindAddress, fmt.Sprintf("%v", *mysqlServerPort)), authServer, vtgateHandle, *mysqlConnReadTimeout, *mysqlConnWriteTimeout, *mysqlProxyProtocol)
 		if err != nil {
 			log.Exitf("mysql.NewListener failed: %v", err)
 		}
@@ -369,7 +391,7 @@ func initMySQLProtocol() {
 		// Let's create this unix socket with permissions to all users. In this way,
 		// clients can connect to vtgate mysql server without being vtgate user
 		oldMask := syscall.Umask(000)
-		mysqlUnixListener, err = newMysqlUnixSocket(*mysqlServerSocketPath, authServer, vh)
+		mysqlUnixListener, err = newMysqlUnixSocket(*mysqlServerSocketPath, authServer, vtgateHandle)
 		_ = syscall.Umask(oldMask)
 		if err != nil {
 			log.Exitf("mysql.NewListener failed: %v", err)
@@ -436,9 +458,36 @@ func shutdownMysqlProtocolAndDrain() {
 	}
 }
 
+func rollbackAtShutdown() {
+	defer log.Flush()
+
+	// Close all open connections. If they're waiting for reads, this will cause
+	// them to error out, which will automatically rollback open transactions.
+	func() {
+		vtgateHandle.mu.Lock()
+		defer vtgateHandle.mu.Unlock()
+		for c := range vtgateHandle.connections {
+			log.Infof("Rolling back transactions associated with connection ID: %v", c.ConnectionID)
+			c.Close()
+		}
+	}()
+
+	// If vtgate is instead busy executing a query, the number of open conns
+	// will be non-zero. Give another second for those queries to finish.
+	for i := 0; i < 100; i++ {
+		if vtgateHandle.numConnections() == 0 {
+			log.Infof("All connections have been rolled back.")
+			return
+		}
+		time.Sleep(10 * time.Millisecond)
+	}
+	log.Errorf("All connections did not go idle. Shutting down anyway.")
+}
+
 func init() {
 	servenv.OnRun(initMySQLProtocol)
 	servenv.OnTermSync(shutdownMysqlProtocolAndDrain)
+	servenv.OnClose(rollbackAtShutdown)
 }
 
 var pluginInitializers []func()