review 1 fix

Author: Artem Klyukvin

connection.go

@@ -50,13 +50,18 @@ type mysqlConn struct {
 	closed   atomicBool  // set when conn is closed, before closech is closed
 
 	// for killing query after timeout
-	id  int
-	d   MySQLDriver
-	dsn string
+	id int
+	d  MySQLDriver
 }
 
 func (mc *mysqlConn) kill() error {
-	conn, err := mc.d.Open(mc.dsn)
+	t := 50 * time.Millisecond
+	killCfg := *mc.cfg
+	killCfg.Timeout = t
+	killCfg.ReadTimeout = t
+	killCfg.WriteTimeout = t
+
+	conn, err := mc.d.Open(killCfg.FormatDSN())
 	if err != nil {
 		return err
 	}
@@ -461,8 +466,10 @@ func (mc *mysqlConn) getSystemVar(name string) ([]byte, error) {
 // finish is called when the query has canceled.
 func (mc *mysqlConn) cancel(err error) {
 	mc.canceled.Set(err)
-	// do not put kill to cleanup to prevent cyclic kills
-	mc.kill()
+	if mc.cfg.KillQueryOnTimeout {
+		// do not put kill to cleanup to prevent cyclic kills
+		mc.kill()
+	}
 	mc.cleanup()
 }
 

connection_go18.go

@@ -160,10 +160,6 @@ func (mc *mysqlConn) watchCancel(ctx context.Context) error {
 	select {
 	default:
 	case <-ctx.Done():
-		killErr := mc.kill()
-		if killErr != nil {
-			errLog.Print("failed to kill query: ", killErr)
-		}
 		return ctx.Err()
 	}
 	if mc.watcher == nil {

driver.go

@@ -151,8 +151,6 @@ func (d MySQLDriver) Open(dsn string) (driver.Conn, error) {
 	}
 
 	mc.d = d
-	mc.dsn = dsn
-
 	return mc, nil
 }
 

driver_go18_test.go

@@ -126,17 +126,19 @@ func getQueryProcess(db *sql.DB, dbName, query string) (*dbProcess, error) {
 	return longProcess, err
 }
 
-func killQuery(db *sql.DB, dbName, query string, cancel context.CancelFunc) error {
+var expectedKilledErr = fmt.Errorf("process expected to be killed")
+
+func killQuery(db *sql.DB, dbName, query string, timeout time.Duration, cancel context.CancelFunc) error {
 	process, err := getQueryProcess(db, dbName, query)
 	if err != nil {
 		return fmt.Errorf("failed to get mysql process: %v", err)
 	}
 	cancel()
 
-	end := time.Now().Add(killTimeout)
+	end := time.Now().Add(timeout)
 	for time.Now().Before(end) {
 		if checkProcessExists(dbName, process.ID, db) {
-			err = fmt.Errorf("process %d expected to be killed", process.ID)
+			err = expectedKilledErr
 			time.Sleep(pollTimeout)
 		} else {
 			err = nil
@@ -173,7 +175,7 @@ func testCancel(dbt *DBTest, ctx context.Context, cancel context.CancelFunc, que
 	}()
 
 	// it is safe to not use timeouts here since they are inside the killQuery function
-	err = killQuery(dbt.db, dbname, query, cancel)
+	err = killQuery(dbt.db, dbname, query, killTimeout, cancel)
 	if err != nil {
 		dbt.Error(err)
 		return
@@ -195,6 +197,62 @@ func testCancel(dbt *DBTest, ctx context.Context, cancel context.CancelFunc, que
 	tx.Commit()
 }
 
+func testCancelNoKill(dbt *DBTest, ctx context.Context, cancel context.CancelFunc, query string, queryFunc func() error) {
+	tx, err := dbt.db.BeginTx(context.Background(), nil)
+	if err != nil {
+		dbt.Fatal(err)
+		return
+	}
+
+	_, err = tx.Exec("LOCK TABLES test WRITE")
+	if err != nil {
+		tx.Rollback()
+		dbt.Fatal(err)
+	}
+
+	errChan := make(chan error)
+	go func() {
+		// This query will be canceled.
+		err = queryFunc()
+		if err != nil && err != context.Canceled {
+			errLog.Print(err)
+		}
+		if err != context.Canceled && ctx.Err() != context.Canceled {
+			errChan <- fmt.Errorf("expected context.Canceled, got %v", err)
+			return
+		}
+		errChan <- nil
+	}()
+
+	// it is safe to not use timeouts here since they are inside the killQuery function
+	err = killQuery(dbt.db, dbname, query, 500*time.Millisecond, cancel)
+	if err != expectedKilledErr {
+		if err == nil {
+			dbt.Errorf("query kill expected to fail")
+		} else {
+			dbt.Errorf(fmt.Sprintf("unexpected error %s", err))
+		}
+	}
+
+	_, err = tx.Exec("UNLOCK TABLES")
+	if err != nil {
+		tx.Rollback()
+		dbt.Fatal(err)
+	}
+	tx.Commit()
+
+	<-errChan
+}
+
+func getKillDSN() string {
+	cfg, err := ParseDSN(dsn)
+	if err != nil {
+		panic(err)
+	}
+	cfg.KillQueryOnTimeout = true
+	return cfg.FormatDSN()
+}
+
 func TestMultiResultSet(t *testing.T) {
 	type result struct {
 		values  [][]int
@@ -385,12 +443,62 @@ func TestPingContext(t *testing.T) {
 	})
 }
 
-func TestContextCancelExec(t *testing.T) {
+func TestContextCancelNoKill(t *testing.T) {
 	runTests(t, dsn, func(dbt *DBTest) {
 		dbt.mustExec("CREATE TABLE test (v INTEGER)")
 		ctx, cancel := context.WithCancel(context.Background())
 		exec := "INSERT INTO test VALUES(1)"
 
+		testCancelNoKill(dbt, ctx, cancel, exec, func() error {
+			_, err := dbt.db.ExecContext(ctx, exec)
+			return err
+		})
+
+		// Check how many times the query is executed.
+		var v int
+		var err error
+		for i := 0; i != 3; i++ {
+			err = nil
+			if err := dbt.db.QueryRow("SELECT COUNT(*) FROM test").Scan(&v); err != nil {
+				dbt.Fatalf("%s", err.Error())
+				return
+			}
+			if v != 1 {
+				err = fmt.Errorf("expected val to be 1, got %d", v)
+			}
+
+			if err != nil {
+				time.Sleep(100 * time.Millisecond) // wait while insert is executed after table lock released
+			}
+		}
+		if err != nil {
+			dbt.Error(err)
+			return
+		}
+
+		// Context is already canceled, so error should come before execution.
+		if _, err := dbt.db.ExecContext(ctx, "INSERT INTO test VALUES (1)"); err == nil {
+			dbt.Error("expected error")
+		} else if err.Error() != "context canceled" {
+			dbt.Fatalf("unexpected error: %s", err)
+		}
+
+		// The second insert query will fail, so the table has no changes.
+		if err := dbt.db.QueryRow("SELECT COUNT(*) FROM test").Scan(&v); err != nil {
+			dbt.Fatalf("%s", err.Error())
+		}
+		if v != 1 {
+			dbt.Errorf("expected val to be 1, got %d", v)
+		}
+	})
+}
+
+func TestContextCancelExec(t *testing.T) {
+	runTests(t, getKillDSN(), func(dbt *DBTest) {
+		dbt.mustExec("CREATE TABLE test (v INTEGER)")
+		ctx, cancel := context.WithCancel(context.Background())
+		exec := "INSERT INTO test VALUES(1)"
+
 		testCancel(dbt, ctx, cancel, exec, func() error {
 			_, err := dbt.db.ExecContext(ctx, exec)
 			return err
@@ -423,7 +531,7 @@ func TestContextCancelExec(t *testing.T) {
 }
 
 func TestContextCancelQuery(t *testing.T) {
-	runTests(t, dsn, func(dbt *DBTest) {
+	runTests(t, getKillDSN(), func(dbt *DBTest) {
 		dbt.mustExec("CREATE TABLE test (v INTEGER)")
 		ctx, cancel := context.WithCancel(context.Background())
 		query := "SELECT 1 FROM test"
@@ -501,7 +609,7 @@ func TestContextCancelPrepare(t *testing.T) {
 }
 
 func TestContextCancelStmtExec(t *testing.T) {
-	runTests(t, dsn, func(dbt *DBTest) {
+	runTests(t, getKillDSN(), func(dbt *DBTest) {
 		dbt.mustExec("CREATE TABLE test (v INTEGER)")
 		ctx, cancel := context.WithCancel(context.Background())
 		exec := "INSERT INTO test VALUES(1)"
@@ -528,7 +636,7 @@ func TestContextCancelStmtExec(t *testing.T) {
 }
 
 func TestContextCancelStmtQuery(t *testing.T) {
-	runTests(t, dsn, func(dbt *DBTest) {
+	runTests(t, getKillDSN(), func(dbt *DBTest) {
 		dbt.mustExec("CREATE TABLE test (v INTEGER)")
 		ctx, cancel := context.WithCancel(context.Background())
 		query := "SELECT 1 FROM test"
@@ -555,7 +663,7 @@ func TestContextCancelStmtQuery(t *testing.T) {
 }
 
 func TestContextCancelBegin(t *testing.T) {
-	runTests(t, dsn, func(dbt *DBTest) {
+	runTests(t, getKillDSN(), func(dbt *DBTest) {
 		dbt.mustExec("CREATE TABLE test (v INTEGER)")
 		ctx, cancel := context.WithCancel(context.Background())
 		query := "SELECT 1 FROM test"

dsn.go

@@ -57,6 +57,7 @@ type Config struct {
 	MultiStatements         bool // Allow multiple statements in one query
 	ParseTime               bool // Parse time values to time.Time
 	RejectReadOnly          bool // Reject read-only connections
+	KillQueryOnTimeout      bool // kill query on the server side if context timed out
 }
 
 // NewConfig creates a new Config and sets default values.
@@ -254,6 +255,15 @@ func (cfg *Config) FormatDSN() string {
 		}
 	}
 
+	if cfg.KillQueryOnTimeout {
+		if hasParam {
+			buf.WriteString("&killQueryOnTimeout=true")
+		} else {
+			hasParam = true
+			buf.WriteString("?killQueryOnTimeout=true")
+		}
+	}
+
 	if cfg.Timeout > 0 {
 		if hasParam {
 			buf.WriteString("&timeout=")
@@ -512,6 +522,14 @@ func parseDSNParams(cfg *Config, params string) (err error) {
 				return errors.New("invalid bool value: " + value)
 			}
 
+		// Kill queries on context timeout
+		case "killQueryOnTimeout":
+			var isBool bool
+			cfg.KillQueryOnTimeout, isBool = readBool(value)
+			if !isBool {
+				return errors.New("invalid bool value: " + value)
+			}
+
 		// Strict mode
 		case "strict":
 			panic("strict mode has been removed. See https://github.com/go-sql-driver/mysql/wiki/strict-mode")

packets.go

@@ -179,10 +179,9 @@ func (mc *mysqlConn) readInitPacket() ([]byte, error) {
 	}
 
 	// server version [null terminated string]
+	// connection id [4 bytes]
 	idPos := 1 + bytes.IndexByte(data[1:], 0x00) + 1
 	mc.id = int(binary.LittleEndian.Uint32(data[idPos : idPos+4]))
-
-	// connection id [4 bytes]
 	pos := idPos + 4
 
 	// first part of the password cipher [8 bytes]