Add engine primitive to handle query timeout

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

@@ -94,10 +94,41 @@ func TestQueryTimeoutWithTables(t *testing.T) {
 	utils.Exec(t, mcmp.VtConn, "select /*vt+ QUERY_TIMEOUT_MS=500 */ sleep(0.1) from t1 where id1 = 1")
 	_, err = utils.ExecAllowError(t, mcmp.VtConn, "select /*vt+ QUERY_TIMEOUT_MS=20 */ sleep(0.1) from t1 where id1 = 1")
 	require.Error(t, err)
-	assert.Contains(t, err.Error(), "context deadline exceeded")
+	if utils.BinaryIsAtLeastAtVersion(21, "vtgate") {
+		assert.ErrorContains(t, err, "Query execution was interrupted, maximum statement execution time exceeded")
+	} else {
+		assert.Contains(t, err.Error(), "context deadline exceeded")
+	}
 	assert.Contains(t, err.Error(), "(errno 1317) (sqlstate 70100)")
 }
 
+// TestOverallQueryTimeout tests that the query timeout is applied to the overall execution of a query
+// and not just individual routes.
+func TestOverallQueryTimeout(t *testing.T) {
+	utils.SkipIfBinaryIsBelowVersion(t, 21, "vtgate")
+	mcmp, closer := start(t)
+	defer closer()
+
+	mcmp.Exec("insert into t1(id1, id2) values (2,2),(3,3)")
+	// After inserting the rows above, if we run the following query, we will end up doing join on vtgate
+	// that issues one select query on the left side and 2 on the right side. The queries on the right side
+	// take 2 and 3 seconds each to run. If we have an overall timeout for 4 seconds, then it should fail.
+
+	_, err := utils.ExecAllowError(t, mcmp.VtConn, "select /*vt+ QUERY_TIMEOUT_MS=4000 */ sleep(u2.id2), u1.id2 from t1 u1 join t1 u2 where u1.id2 = u2.id1")
+	assert.Error(t, err)
+	assert.ErrorContains(t, err, "Query execution was interrupted, maximum statement execution time exceeded")
+
+	// Let's also check that setting the session variable also works.
+	utils.Exec(t, mcmp.VtConn, "set query_timeout=4000")
+	_, err = utils.ExecAllowError(t, mcmp.VtConn, "select sleep(u2.id2), u1.id2 from t1 u1 join t1 u2 where u1.id2 = u2.id1")
+	assert.Error(t, err)
+	assert.ErrorContains(t, err, "Query execution was interrupted, maximum statement execution time exceeded")
+
+	// Increasing the timeout should pass the query.
+	utils.Exec(t, mcmp.VtConn, "set query_timeout=10000")
+	_ = utils.Exec(t, mcmp.VtConn, "select sleep(u2.id2), u1.id2 from t1 u1 join t1 u2 where u1.id2 = u2.id1")
+}
+
 // TestQueryTimeoutWithShardTargeting tests the query timeout with shard targeting.
 func TestQueryTimeoutWithShardTargeting(t *testing.T) {
 	mcmp, closer := start(t)

go/test/vschemawrapper/vschema_wrapper.go

@@ -47,6 +47,7 @@ type VSchemaWrapper struct {
 	Dest                  key.Destination
 	SysVarEnabled         bool
 	ForeignKeyChecksState *bool
+	SessionQueryTimeout   int
 	Version               plancontext.PlannerVersion
 	EnableViews           bool
 	TestBuilder           func(query string, vschema plancontext.VSchema, keyspace string) (*engine.Plan, error)
@@ -132,6 +133,13 @@ func (vw *VSchemaWrapper) Environment() *vtenv.Environment {
 	return vw.Env
 }
 
+func (vw *VSchemaWrapper) GetQueryTimeout(queryTimeoutFromComments int) int {
+	if queryTimeoutFromComments != 0 {
+		return queryTimeoutFromComments
+	}
+	return vw.SessionQueryTimeout
+}
+
 func (vw *VSchemaWrapper) PlannerWarning(_ string) {
 }
 

go/vt/vterrors/code.go

@@ -119,6 +119,8 @@ var (
 	VT14004 = errorWithoutState("VT14004", vtrpcpb.Code_UNAVAILABLE, "cannot find keyspace for: %s", "The specified keyspace could not be found.")
 	VT14005 = errorWithoutState("VT14005", vtrpcpb.Code_UNAVAILABLE, "cannot lookup sidecar database for keyspace: %s", "Failed to read sidecar database identifier.")
 
+	VT15001 = errorWithoutState("VT15001", vtrpcpb.Code_DEADLINE_EXCEEDED, "Query execution was interrupted, maximum statement execution time exceeded", "Query execution was interrupted, maximum statement execution time exceeded")
+
 	// Errors is a list of errors that must match all the variables
 	// defined above to enable auto-documentation of error codes.
 	Errors = []func(args ...any) *VitessError{

go/vt/vtgate/engine/fake_primitive_test.go

@@ -22,6 +22,7 @@ import (
 	"reflect"
 	"strings"
 	"testing"
+	"time"
 
 	"golang.org/x/sync/errgroup"
 
@@ -41,6 +42,9 @@ type fakePrimitive struct {
 
 	log []string
 
+	// sleepTime is the time for which the fake primitive sleeps before returning the results.
+	sleepTime time.Duration
+
 	allResultsInOneCall bool
 
 	async bool
@@ -71,6 +75,9 @@ func (f *fakePrimitive) GetTableName() string {
 
 func (f *fakePrimitive) TryExecute(ctx context.Context, vcursor VCursor, bindVars map[string]*querypb.BindVariable, wantfields bool) (*sqltypes.Result, error) {
 	f.log = append(f.log, fmt.Sprintf("Execute %v %v", printBindVars(bindVars), wantfields))
+	if f.sleepTime != 0 {
+		time.Sleep(f.sleepTime)
+	}
 	if f.results == nil {
 		return nil, f.sendErr
 	}
@@ -85,6 +92,9 @@ func (f *fakePrimitive) TryExecute(ctx context.Context, vcursor VCursor, bindVar
 
 func (f *fakePrimitive) TryStreamExecute(ctx context.Context, vcursor VCursor, bindVars map[string]*querypb.BindVariable, wantfields bool, callback func(*sqltypes.Result) error) error {
 	f.log = append(f.log, fmt.Sprintf("StreamExecute %v %v", printBindVars(bindVars), wantfields))
+	if f.sleepTime != 0 {
+		time.Sleep(f.sleepTime)
+	}
 	if f.results == nil {
 		return f.sendErr
 	}

go/vt/vtgate/engine/fake_vcursor_test.go

@@ -54,7 +54,8 @@ var _ SessionActions = (*noopVCursor)(nil)
 
 // noopVCursor is used to build other vcursors.
 type noopVCursor struct {
-	inTx bool
+	inTx         bool
+	queryTimeout int
 }
 
 // MySQLVersion implements VCursor.
@@ -298,7 +299,10 @@ func (t *noopVCursor) SetQueryTimeout(maxExecutionTime int64) {
 }
 
 func (t *noopVCursor) GetQueryTimeout(queryTimeoutFromComments int) int {
-	return queryTimeoutFromComments
+	if queryTimeoutFromComments != 0 {
+		return queryTimeoutFromComments
+	}
+	return t.queryTimeout
 }
 
 func (t *noopVCursor) SetSkipQueryPlanCache(context.Context, bool) error {

go/vt/vtgate/engine/timeout_handler.go

@@ -0,0 +1,104 @@
+package engine
+
+import (
+	"context"
+	"time"
+
+	"vitess.io/vitess/go/sqltypes"
+	querypb "vitess.io/vitess/go/vt/proto/query"
+	"vitess.io/vitess/go/vt/vterrors"
+)
+
+// TimeoutHandler is a primitive that adds a timeout to the execution of a query.
+type TimeoutHandler struct {
+	Timeout int
+	Input   Primitive
+}
+
+var _ Primitive = (*TimeoutHandler)(nil)
+
+// NewTimeoutHandler creates a new timeout handler.
+func NewTimeoutHandler(input Primitive, timeout int) *TimeoutHandler {
+	return &TimeoutHandler{
+		Timeout: timeout,
+		Input:   input,
+	}
+}
+
+// RouteType is part of the Primitive interface
+func (t *TimeoutHandler) RouteType() string {
+	return t.Input.RouteType()
+}
+
+// GetKeyspaceName is part of the Primitive interface
+func (t *TimeoutHandler) GetKeyspaceName() string {
+	return t.Input.GetKeyspaceName()
+}
+
+// GetTableName is part of the Primitive interface
+func (t *TimeoutHandler) GetTableName() string {
+	return t.Input.GetTableName()
+}
+
+// GetFields is part of the Primitive interface
+func (t *TimeoutHandler) GetFields(ctx context.Context, vcursor VCursor, bindVars map[string]*querypb.BindVariable) (*sqltypes.Result, error) {
+	return t.Input.GetFields(ctx, vcursor, bindVars)
+}
+
+// NeedsTransaction is part of the Primitive interface
+func (t *TimeoutHandler) NeedsTransaction() bool {
+	return t.Input.NeedsTransaction()
+}
+
+// TryExecute is part of the Primitive interface
+func (t *TimeoutHandler) TryExecute(ctx context.Context, vcursor VCursor, bindVars map[string]*querypb.BindVariable, wantfields bool) (res *sqltypes.Result, err error) {
+	ctx, cancel := context.WithTimeout(ctx, time.Duration(t.Timeout)*time.Millisecond)
+	defer cancel()
+
+	complete := make(chan any)
+	go func() {
+		res, err = t.Input.TryExecute(ctx, vcursor, bindVars, wantfields)
+		close(complete)
+	}()
+
+	select {
+	case <-ctx.Done():
+		return nil, vterrors.VT15001()
+	case <-complete:
+		return res, err
+	}
+}
+
+// TryStreamExecute is part of the Primitive interface
+func (t *TimeoutHandler) TryStreamExecute(ctx context.Context, vcursor VCursor, bindVars map[string]*querypb.BindVariable, wantfields bool, callback func(*sqltypes.Result) error) (err error) {
+	ctx, cancel := context.WithTimeout(ctx, time.Duration(t.Timeout)*time.Millisecond)
+	defer cancel()
+
+	complete := make(chan any)
+	go func() {
+		err = t.Input.TryStreamExecute(ctx, vcursor, bindVars, wantfields, callback)
+		close(complete)
+	}()
+
+	select {
+	case <-ctx.Done():
+		return vterrors.VT15001()
+	case <-complete:
+		return err
+	}
+}
+
+// Inputs is part of the Primitive interface
+func (t *TimeoutHandler) Inputs() ([]Primitive, []map[string]any) {
+	return []Primitive{t.Input}, nil
+}
+
+// description is part of the Primitive interface
+func (t *TimeoutHandler) description() PrimitiveDescription {
+	return PrimitiveDescription{
+		OperatorType: "TimeoutHandler",
+		Other: map[string]any{
+			"Timeout": t.Timeout,
+		},
+	}
+}

go/vt/vtgate/engine/timeout_handler_test.go

@@ -0,0 +1,59 @@
+package engine
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+
+	"vitess.io/vitess/go/sqltypes"
+)
+
+// TestTimeoutHandler tests timeout handler primitive.
+func TestTimeoutHandler(t *testing.T) {
+	tests := []struct {
+		name      string
+		sleepTime time.Duration
+		timeout   int
+		wantErr   string
+	}{
+		{
+			name:      "Timeout without failure",
+			sleepTime: 100 * time.Millisecond,
+			timeout:   1000,
+			wantErr:   "",
+		}, {
+			name:      "Timeout with failure",
+			sleepTime: 2 * time.Second,
+			timeout:   100,
+			wantErr:   "VT15001: Query execution was interrupted, maximum statement execution time exceeded",
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			_, err := createTimeoutHandlerForTesting(tt.timeout, tt.sleepTime).TryExecute(context.Background(), &noopVCursor{}, nil, false)
+			if tt.wantErr != "" {
+				require.EqualError(t, err, tt.wantErr)
+			} else {
+				require.NoError(t, err)
+			}
+			err = createTimeoutHandlerForTesting(tt.timeout, tt.sleepTime).TryStreamExecute(context.Background(), &noopVCursor{}, nil, false, func(result *sqltypes.Result) error {
+				return nil
+			})
+			if tt.wantErr != "" {
+				require.EqualError(t, err, tt.wantErr)
+			} else {
+				require.NoError(t, err)
+			}
+		})
+	}
+}
+
+// createTimeoutHandlerForTesting creates a TimeoutHandler for testing that has a fakePrimitive as an input.
+func createTimeoutHandlerForTesting(timeout int, sleepTime time.Duration) *TimeoutHandler {
+	return NewTimeoutHandler(&fakePrimitive{
+		results:   nil,
+		sleepTime: sleepTime,
+	}, timeout)
+}

go/vt/vtgate/planbuilder/builder.go

@@ -153,13 +153,7 @@ func buildRoutePlan(stmt sqlparser.Statement, reservedVars *sqlparser.ReservedVa
 
 func createInstructionFor(ctx context.Context, query string, stmt sqlparser.Statement, reservedVars *sqlparser.ReservedVars, vschema plancontext.VSchema, enableOnlineDDL, enableDirectDDL bool) (*planResult, error) {
 	switch stmt := stmt.(type) {
-	case *sqlparser.Select, *sqlparser.Insert, *sqlparser.Update, *sqlparser.Delete:
-		configuredPlanner, err := getConfiguredPlanner(vschema, stmt, query)
-		if err != nil {
-			return nil, err
-		}
-		return buildRoutePlan(stmt, reservedVars, vschema, configuredPlanner)
-	case *sqlparser.Union:
+	case *sqlparser.Select, *sqlparser.Insert, *sqlparser.Update, *sqlparser.Delete, *sqlparser.Union:
 		configuredPlanner, err := getConfiguredPlanner(vschema, stmt, query)
 		if err != nil {
 			return nil, err

go/vt/vtgate/planbuilder/plan_test.go

@@ -260,6 +260,16 @@ func (s *planTestSuite) addPKsProvided(vschema *vindexes.VSchema, ks string, tbl
 	}
 }
 
+func (s *planTestSuite) TestQueryTimeout() {
+	vschemaWrapper := &vschemawrapper.VSchemaWrapper{
+		V:                   loadSchema(s.T(), "vschemas/schema.json", true),
+		Env:                 vtenv.NewTestEnv(),
+		SessionQueryTimeout: 200,
+	}
+
+	s.testFile("query_timeout_cases.json", vschemaWrapper, false)
+}
+
 func (s *planTestSuite) TestSystemTables57() {
 	// first we move everything to use 5.7 logic
 	env, err := vtenv.New(vtenv.Options{
@@ -683,7 +693,7 @@ func (s *planTestSuite) testFile(filename string, vschema *vschemawrapper.VSchem
 					if tcase.Skip {
 						t.Skip(message)
 					} else {
-						t.Errorf(message)
+						t.Error(message)
 					}
 				} else if tcase.Skip {
 					t.Errorf("query is correct even though it is skipped:\n %s", tcase.Query)

go/vt/vtgate/planbuilder/plancontext/planning_context_test.go

@@ -280,6 +280,10 @@ func (v *vschema) Environment() *vtenv.Environment {
 	return vtenv.NewTestEnv()
 }
 
+func (v *vschema) GetQueryTimeout(queryTimeoutFromComments int) int {
+	return queryTimeoutFromComments
+}
+
 func (v *vschema) ErrorIfShardedF(keyspace *vindexes.Keyspace, warn, errFmt string, params ...any) error {
 	// TODO implement me
 	panic("implement me")

go/vt/vtgate/planbuilder/plancontext/vschema.go

@@ -43,6 +43,7 @@ type VSchema interface {
 	SetPlannerVersion(pv PlannerVersion)
 	ConnCollation() collations.ID
 	Environment() *vtenv.Environment
+	GetQueryTimeout(queryTimeoutFromComments int) int
 
 	// ErrorIfShardedF will return an error if the keyspace is sharded,
 	// and produce a warning if the vtgate if configured to do so

go/vt/vtgate/planbuilder/select.go

@@ -229,9 +229,23 @@ func newBuildSelectPlan(
 		return nil, nil, err
 	}
 
+	plan = handleTimeout(plan, queryTimeout(selStmt.GetParsedComments().Directives()), vschema)
+
 	return plan, operators.TablesUsed(op), nil
 }
 
+// handleTimeout checks if there is a timeout that needs to be added to the query. If there is one
+// then we wrap the plan in a TimeoutHandler primitive. If there is no timeout, we return the plan as is.
+// Because we are accessing the query timeout stored in the session state, we have to also add this value to the plan key
+// so that we don't end up using the same plan when the session variable changes.
+func handleTimeout(plan engine.Primitive, queryTimeoutComment int, vschema plancontext.VSchema) engine.Primitive {
+	finalQueryTimeout := vschema.GetQueryTimeout(queryTimeoutComment)
+	if finalQueryTimeout == 0 {
+		return plan
+	}
+	return engine.NewTimeoutHandler(plan, finalQueryTimeout)
+}
+
 func createSelectOperator(ctx *plancontext.PlanningContext, selStmt sqlparser.SelectStatement, reservedVars *sqlparser.ReservedVars) (operators.Operator, error) {
 	err := queryRewrite(ctx, selStmt)
 	if err != nil {