feat: authoritative query timeout for vttablet from vtgate

go/test/endtoend/vtgate/queries/timeout/main_test.go

@@ -63,8 +63,9 @@ func TestMain(m *testing.M) {
 
 		clusterInstance.VtTabletExtraArgs = append(clusterInstance.VtTabletExtraArgs,
 			"--queryserver-config-max-result-size", "1000000",
-			"--queryserver-config-query-timeout", "200s",
-			"--queryserver-config-query-pool-timeout", "200s")
+			"--queryserver-config-query-timeout", "2s",
+			"--queryserver-config-transaction-timeout", "3s",
+			"--queryserver-config-query-pool-timeout", "2s")
 		// Start Unsharded keyspace
 		ukeyspace := &cluster.Keyspace{
 			Name:      uks,

go/test/endtoend/vtgate/queries/timeout/timeout_test.go

@@ -17,11 +17,13 @@ limitations under the License.
 package misc
 
 import (
+	"context"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
+	"vitess.io/vitess/go/mysql"
 	"vitess.io/vitess/go/test/endtoend/cluster"
 	"vitess.io/vitess/go/test/endtoend/utils"
 )
@@ -66,6 +68,9 @@ func TestQueryTimeoutWithDual(t *testing.T) {
 	assert.Error(t, err)
 	_, err = utils.ExecAllowError(t, mcmp.VtConn, "select /*vt+ QUERY_TIMEOUT_MS=15 */ sleep(0.001) from dual")
 	assert.NoError(t, err)
+	// infinite query timeout overriding all defaults
+	_, err = utils.ExecAllowError(t, mcmp.VtConn, "select /*vt+ QUERY_TIMEOUT_MS=0 */ sleep(5) from dual")
+	assert.NoError(t, err)
 }
 
 func TestQueryTimeoutWithTables(t *testing.T) {
@@ -124,3 +129,53 @@ func TestQueryTimeoutWithShardTargeting(t *testing.T) {
 		})
 	}
 }
+
+func TestQueryTimeoutWithoutVTGateDefault(t *testing.T) {
+	// disable query timeout
+	clusterInstance.VtGateExtraArgs = append(clusterInstance.VtGateExtraArgs,
+		"--query-timeout", "0")
+	require.NoError(t,
+		clusterInstance.RestartVtgate())
+
+	// update vtgate params
+	vtParams = clusterInstance.GetVTParams(keyspaceName)
+
+	mcmp, closer := start(t)
+	defer closer()
+
+	// tablet query timeout of 2s
+	_, err := utils.ExecAllowError(t, mcmp.VtConn, "select sleep(5) from dual")
+	assert.Error(t, err)
+
+	// infinite timeout using query hint
+	utils.Exec(t, mcmp.VtConn, "select /*vt+ QUERY_TIMEOUT_MS=0 */ sleep(5) from dual")
+
+	// checking again without query hint, tablet query timeout of 2s should be applied
+	_, err = utils.ExecAllowError(t, mcmp.VtConn, "select sleep(5) from dual")
+	assert.Error(t, err)
+
+	// set timeout of 20ms
+	utils.Exec(t, mcmp.VtConn, "set query_timeout=20")
+
+	// query timeout of 20ms should be applied
+	_, err = utils.ExecAllowError(t, mcmp.VtConn, "select sleep(1) from dual")
+	assert.Error(t, err)
+
+	// infinite timeout using query hint will override session timeout.
+	utils.Exec(t, mcmp.VtConn, "select /*vt+ QUERY_TIMEOUT_MS=0 */ sleep(5) from dual")
+
+	// open second session
+	conn2, err := mysql.Connect(context.Background(), &vtParams)
+	require.NoError(t, err)
+	defer conn2.Close()
+
+	// tablet query timeout of 2s should be applied, as session timeout is not set on this connection.
+	utils.Exec(t, conn2, "select sleep(1) from dual")
+	_, err = utils.ExecAllowError(t, conn2, "select sleep(5) from dual")
+	assert.Error(t, err)
+
+	// reset session on first connection, tablet query timeout of 2s should be applied.
+	utils.Exec(t, mcmp.VtConn, "set query_timeout=0")
+	_, err = utils.ExecAllowError(t, mcmp.VtConn, "select sleep(5) from dual")
+	assert.Error(t, err)
+}

go/vt/sqlparser/comments.go

@@ -554,27 +554,6 @@ func AllowScatterDirective(stmt Statement) bool {
 	return checkDirective(stmt, DirectiveAllowScatter)
 }
 
-// ForeignKeyChecksState returns the state of foreign_key_checks variable if it is part of a SET_VAR optimizer hint in the comments.
-func ForeignKeyChecksState(stmt Statement) *bool {
-	cmt, ok := stmt.(Commented)
-	if ok {
-		fkChecksVal := cmt.GetParsedComments().GetMySQLSetVarValue(sysvars.ForeignKeyChecks)
-		// If the value of the `foreign_key_checks` optimizer hint is something that doesn't make sense,
-		// then MySQL just ignores it and treats it like the case, where it is unspecified. We are choosing
-		// to have the same behaviour here. If the value doesn't match any of the acceptable values, we return nil,
-		// that signifies that no value was specified.
-		switch strings.ToLower(fkChecksVal) {
-		case "on", "1":
-			fkState := true
-			return &fkState
-		case "off", "0":
-			fkState := false
-			return &fkState
-		}
-	}
-	return nil
-}
-
 func checkDirective(stmt Statement, key string) bool {
 	cmt, ok := stmt.(Commented)
 	if ok {
@@ -583,42 +562,43 @@ func checkDirective(stmt Statement, key string) bool {
 	return false
 }
 
-// GetPriorityFromStatement gets the priority from the provided Statement, using DirectivePriority
-func GetPriorityFromStatement(statement Statement) (string, error) {
-	commentedStatement, ok := statement.(Commented)
-	// This would mean that the statement lacks comments, so we can't obtain the workload from it. Hence default to
-	// empty priority
+type QueryHints struct {
+	IgnoreMaxMemoryRows bool
+	Consolidator        querypb.ExecuteOptions_Consolidator
+	Workload            string
+	ForeignKeyChecks    *bool
+	Priority            string
+	Timeout             *int
+}
+
+func BuildQueryHints(stmt Statement) (qh QueryHints, err error) {
+	qh = QueryHints{}
+
+	comment, ok := stmt.(Commented)
 	if !ok {
-		return "", nil
+		return qh, nil
 	}
 
-	directives := commentedStatement.GetParsedComments().Directives()
-	priority, ok := directives.GetString(DirectivePriority, "")
-	if !ok || priority == "" {
-		return "", nil
-	}
+	directives := comment.GetParsedComments().Directives()
 
-	intPriority, err := strconv.Atoi(priority)
-	if err != nil || intPriority < 0 || intPriority > MaxPriorityValue {
-		return "", ErrInvalidPriority
+	qh.Priority, err = getPriority(directives)
+	if err != nil {
+		return qh, err
 	}
+	qh.IgnoreMaxMemoryRows = directives.IsSet(DirectiveIgnoreMaxMemoryRows)
+	qh.Consolidator = getConsolidator(stmt, directives)
+	qh.Workload = getWorkload(directives)
+	qh.ForeignKeyChecks = getForeignKeyChecksState(comment)
+	qh.Timeout = getQueryTimeout(directives)
 
-	return priority, nil
+	return qh, nil
 }
 
-// Consolidator returns the consolidator option.
-func Consolidator(stmt Statement) querypb.ExecuteOptions_Consolidator {
-	var comments *ParsedComments
-	switch stmt := stmt.(type) {
-	case *Select:
-		comments = stmt.Comments
-	default:
-		return querypb.ExecuteOptions_CONSOLIDATOR_UNSPECIFIED
-	}
-	if comments == nil {
+// getConsolidator returns the consolidator option.
+func getConsolidator(stmt Statement, directives *CommentDirectives) querypb.ExecuteOptions_Consolidator {
+	if _, isSelect := stmt.(SelectStatement); !isSelect {
 		return querypb.ExecuteOptions_CONSOLIDATOR_UNSPECIFIED
 	}
-	directives := comments.Directives()
 	strv, isSet := directives.GetString(DirectiveConsolidator, "")
 	if !isSet {
 		return querypb.ExecuteOptions_CONSOLIDATOR_UNSPECIFIED
@@ -629,18 +609,56 @@ func Consolidator(stmt Statement) querypb.ExecuteOptions_Consolidator {
 	return querypb.ExecuteOptions_CONSOLIDATOR_UNSPECIFIED
 }
 
-// GetWorkloadNameFromStatement gets the workload name from the provided Statement, using workloadLabel as the name of
+// getWorkload gets the workload name from the provided Statement, using workloadLabel as the name of
 // the query directive that specifies it.
-func GetWorkloadNameFromStatement(statement Statement) string {
-	commentedStatement, ok := statement.(Commented)
-	// This would mean that the statement lacks comments, so we can't obtain the workload from it. Hence default to
-	// empty workload name
-	if !ok {
-		return ""
+func getWorkload(directives *CommentDirectives) string {
+	workloadName, _ := directives.GetString(DirectiveWorkloadName, "")
+	return workloadName
+}
+
+// getForeignKeyChecksState returns the state of foreign_key_checks variable if it is part of a SET_VAR optimizer hint in the comments.
+func getForeignKeyChecksState(cmt Commented) *bool {
+	fkChecksVal := cmt.GetParsedComments().GetMySQLSetVarValue(sysvars.ForeignKeyChecks)
+	// If the value of the `foreign_key_checks` optimizer hint is something that doesn't make sense,
+	// then MySQL just ignores it and treats it like the case, where it is unspecified. We are choosing
+	// to have the same behaviour here. If the value doesn't match any of the acceptable values, we return nil,
+	// that signifies that no value was specified.
+	switch strings.ToLower(fkChecksVal) {
+	case "on", "1":
+		fkState := true
+		return &fkState
+	case "off", "0":
+		fkState := false
+		return &fkState
 	}
+	return nil
+}
 
-	directives := commentedStatement.GetParsedComments().Directives()
-	workloadName, _ := directives.GetString(DirectiveWorkloadName, "")
+// getPriority gets the priority from the provided Statement, using DirectivePriority
+func getPriority(directives *CommentDirectives) (string, error) {
+	priority, ok := directives.GetString(DirectivePriority, "")
+	if !ok || priority == "" {
+		return "", nil
+	}
 
-	return workloadName
+	intPriority, err := strconv.Atoi(priority)
+	if err != nil || intPriority < 0 || intPriority > MaxPriorityValue {
+		return "", ErrInvalidPriority
+	}
+
+	return priority, nil
+}
+
+// getQueryTimeout gets the query timeout from the provided Statement, using DirectiveQueryTimeout
+func getQueryTimeout(directives *CommentDirectives) *int {
+	timeoutString, ok := directives.GetString(DirectiveQueryTimeout, "")
+	if !ok || timeoutString == "" {
+		return nil
+	}
+
+	timeout, err := strconv.Atoi(timeoutString)
+	if err != nil || timeout < 0 {
+		return nil
+	}
+	return &timeout
 }