check that the context deadline hasn't expired before running a query

If the context deadline expires before the query has a chance to
execute, the query killer will terminate the connection immediately
after sending the query to mysql.

Instead, add a check at the start of dbConn.execOnce and fail
the query immediately if the context is already past the deadline.

Signed-off-by: Michael Demmer <mdemmer@slack-corp.com>

Author: demmer

go/mysql/fakesqldb/server.go

@@ -70,6 +70,8 @@ type DB struct {
 	name string
 	// isConnFail trigger a panic in the connection handler.
 	isConnFail bool
+	// connDelay causes a sleep in the connection handler
+	connDelay time.Duration
 	// shouldClose, if true, tells ComQuery() to close the connection when
 	// processing the next query. This will trigger a MySQL client error with
 	// errno 2013 ("server lost").
@@ -288,6 +290,10 @@ func (db *DB) NewConnection(c *mysql.Conn) {
 		panic(fmt.Errorf("simulating a connection failure"))
 	}
 
+	if db.connDelay != 0 {
+		time.Sleep(db.connDelay)
+	}
+
 	if conn, ok := db.connections[c.ConnectionID]; ok {
 		db.t.Fatalf("BUG: connection with id: %v is already active. existing conn: %v new conn: %v", c.ConnectionID, conn, c)
 	}
@@ -517,6 +523,13 @@ func (db *DB) DisableConnFail() {
 	db.isConnFail = false
 }
 
+// SetConnDelay delays connections to this fake DB for the given duration
+func (db *DB) SetConnDelay(d time.Duration) {
+	db.mu.Lock()
+	defer db.mu.Unlock()
+	db.connDelay = d
+}
+
 // EnableShouldClose closes the connection when processing the next query.
 func (db *DB) EnableShouldClose() {
 	db.mu.Lock()

go/vt/vttablet/tabletserver/connpool/dbconn.go

@@ -132,6 +132,16 @@ func (dbc *DBConn) execOnce(ctx context.Context, query string, maxrows int, want
 	dbc.current.Set(query)
 	defer dbc.current.Set("")
 
+	// Check if the context is already past its deadline before
+	// trying to execute the query.
+	if ctx.Done() != nil {
+		select {
+		case <-ctx.Done():
+			return nil, fmt.Errorf("%v before execution started", ctx.Err())
+		default:
+		}
+	}
+
 	done, wg := dbc.setDeadline(ctx)
 	if done != nil {
 		defer func() {

go/vt/vttablet/tabletserver/connpool/dbconn_test.go

@@ -117,6 +117,66 @@ func TestDBConnExec(t *testing.T) {
 	compareTimingCounts(t, "Exec", 1, startCounts, tabletenv.MySQLStats.Counts())
 }
 
+func TestDBConnDeadline(t *testing.T) {
+	db := fakesqldb.New(t)
+	defer db.Close()
+	startCounts := tabletenv.MySQLStats.Counts()
+	sql := "select * from test_table limit 1000"
+	expectedResult := &sqltypes.Result{
+		Fields: []*querypb.Field{
+			{Type: sqltypes.VarChar},
+		},
+		RowsAffected: 1,
+		Rows: [][]sqltypes.Value{
+			{sqltypes.NewVarChar("123")},
+		},
+	}
+	db.AddQuery(sql, expectedResult)
+
+	connPool := newPool()
+	connPool.Open(db.ConnParams(), db.ConnParams(), db.ConnParams())
+	defer connPool.Close()
+
+	db.SetConnDelay(100 * time.Millisecond)
+	ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(50*time.Millisecond))
+	defer cancel()
+
+	dbConn, err := NewDBConn(connPool, db.ConnParams())
+	defer dbConn.Close()
+	if err != nil {
+		t.Fatalf("should not get an error, err: %v", err)
+	}
+
+	_, err = dbConn.Exec(ctx, sql, 1, false)
+	want := "context deadline exceeded before execution started"
+	if err == nil || !strings.Contains(err.Error(), want) {
+		t.Errorf("Exec: %v, want %s", err, want)
+	}
+
+	compareTimingCounts(t, "Connect", 1, startCounts, tabletenv.MySQLStats.Counts())
+	compareTimingCounts(t, "ConnectError", 0, startCounts, tabletenv.MySQLStats.Counts())
+	compareTimingCounts(t, "Exec", 0, startCounts, tabletenv.MySQLStats.Counts())
+
+	startCounts = tabletenv.MySQLStats.Counts()
+
+	ctx, cancel = context.WithDeadline(context.Background(), time.Now().Add(10*time.Second))
+	defer cancel()
+
+	result, err := dbConn.Exec(ctx, sql, 1, false)
+	if err != nil {
+		t.Fatalf("should not get an error, err: %v", err)
+	}
+	expectedResult.Fields = nil
+	if !reflect.DeepEqual(expectedResult, result) {
+		t.Errorf("Exec: %v, want %v", expectedResult, result)
+	}
+
+	compareTimingCounts(t, "Connect", 0, startCounts, tabletenv.MySQLStats.Counts())
+	compareTimingCounts(t, "ConnectError", 0, startCounts, tabletenv.MySQLStats.Counts())
+	compareTimingCounts(t, "Exec", 1, startCounts, tabletenv.MySQLStats.Counts())
+
+}
+
 func TestDBConnKill(t *testing.T) {
 	db := fakesqldb.New(t)
 	defer db.Close()