WIP again

Signed-off-by: Tim Vaillancourt <[email protected]>

go/vt/vttablet/tabletserver/query_executor.go

@@ -65,8 +65,9 @@ type QueryExecutor struct {
 }
 
 const (
-	streamRowsSize         = 256
-	maxQueryBufferDuration = 15 * time.Second
+	streamRowsSize           = 256
+	maxQueryBufferDuration   = 15 * time.Second
+	queryTimeoutMysqlMaxWait = time.Second
 )
 
 var (
@@ -918,6 +919,7 @@ func (qre *QueryExecutor) generateFinalSQL(parsedQuery *sqlparser.ParsedQuery, b
 	if err != nil {
 		return "", "", vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "%s", err)
 	}
+	query = addMySQLOptimizerHints(qre.tsv.config, query)
 	if qre.tsv.config.AnnotateQueries {
 		username := callerid.GetPrincipal(callerid.EffectiveCallerIDFromContext(qre.ctx))
 		if username == "" {
@@ -934,10 +936,6 @@ func (qre *QueryExecutor) generateFinalSQL(parsedQuery *sqlparser.ParsedQuery, b
 		buf.WriteString(qre.marginComments.Leading)
 		qre.marginComments.Leading = buf.String()
 	}
-	if qre.tsv.config.QueryTimeoutMethod == "mysql" {
-		timeout := qre.tsv.config.Oltp.QueryTimeout
-		query = addMySQLMaxExecutionTimeComment(query, qre.tsv.config.QueryTimeout)
-	}
 
 	if qre.marginComments.Leading == "" && qre.marginComments.Trailing == "" {
 		return query, query, nil
@@ -951,13 +949,27 @@ func (qre *QueryExecutor) generateFinalSQL(parsedQuery *sqlparser.ParsedQuery, b
 	return buf.String(), query, nil
 }
 
-func addMySQLMaxExecutionTimeComment(query string, queryTimeoutMs int64) string {
-	fields := strings.SplitN(query, " ", 2)
-	return strings.Join([]string{
-		fields[0],
-		fmt.Sprintf("/*+ MAX_EXECUTION_TIME(%d) */", queryTimeoutMs),
-		fields[1],
-	}, " ")
+func addMySQLOptimizerHints(config *tabletenv.TabletConfig, query string) string {
+	hints := make([]string, 0)
+
+	if config.Oltp.QueryTimeoutMethod == "mysql" && config.Oltp.QueryTimeoutSeconds > 0 {
+		// The MAX_EXECUTION_TIME(N) hint sets a statement execution timeout of N milliseconds.
+		// https://dev.mysql.com/doc/refman/8.0/en/optimizer-hints.html#optimizer-hints-execution-time
+		hints = append(hints,
+			fmt.Sprintf("MAX_EXECUTION_TIME(%d)", config.Oltp.QueryTimeoutSeconds.Get().Milliseconds()),
+		)
+	}
+
+	if len(hints) > 0 {
+		fields := strings.SplitN(query, " ", 2)
+		return strings.Join([]string{
+			fields[0],
+			"/*+ " + strings.Join(hints, " ") + " */",
+			fields[1],
+		}, " ")
+	}
+
+	return query
 }
 
 func rewriteOUTParamError(err error) error {

go/vt/vttablet/tabletserver/query_executor_test.go

@@ -1786,3 +1786,16 @@ func TestQueryExecSchemaReloadCount(t *testing.T) {
 		})
 	}
 }
+
+func TestAddMySQLOptimizerHints(t *testing.T) {
+	config := tabletenv.NewDefaultConfig()
+	{
+		config.Oltp.QueryTimeoutMethod = "vttablet"
+		t.Logf("sql: %v", addMySQLOptimizerHints(config, "select * from something"))
+	}
+	{
+		config.Oltp.QueryTimeoutMethod = "mysql"
+		config.Oltp.QueryTimeoutSeconds = tabletenv.Seconds(1)
+		t.Logf("sql: %v", addMySQLOptimizerHints(config, "select * from something"))
+	}
+}

go/vt/vttablet/tabletserver/tabletenv/config.go

@@ -37,12 +37,13 @@ import (
 
 // These constants represent values for various config parameters.
 const (
-	Enable       = "enable"
-	Disable      = "disable"
-	Dryrun       = "dryRun"
-	NotOnPrimary = "notOnPrimary"
-	Polling      = "polling"
-	Heartbeat    = "heartbeat"
+	Enable                    = "enable"
+	Disable                   = "disable"
+	Dryrun                    = "dryRun"
+	NotOnPrimary              = "notOnPrimary"
+	Polling                   = "polling"
+	Heartbeat                 = "heartbeat"
+	DefaultQueryTimeoutMethod = "vttablet"
 )
 
 var (
@@ -183,6 +184,8 @@ func registerTabletEnvFlags(fs *pflag.FlagSet) {
 	fs.Int64Var(&currentConfig.RowStreamer.MaxMySQLReplLagSecs, "vreplication_copy_phase_max_mysql_replication_lag", 43200, "The maximum MySQL replication lag (in seconds) that can exist on a vstreamer (source) before starting another round of copying rows. This helps to limit the impact on the source tablet.")
 
 	fs.BoolVar(&currentConfig.EnableViews, "queryserver-enable-views", false, "Enable views support in vttablet.")
+
+	fs.StringVar(&currentConfig.Oltp.QueryTimeoutMethod, "query-timeout-method", defaultConfig.Oltp.QueryTimeoutMethod, "The method to be used to kill MySQL queries, options: 'vttablet' and 'mysql'. 'vttablet' issues a MySQL KILL operation whereas 'mysql' pushes the kill to MySQL.")
 }
 
 var (
@@ -346,6 +349,7 @@ type OlapConfig struct {
 // OltpConfig contains the config for oltp settings.
 type OltpConfig struct {
 	QueryTimeoutSeconds Seconds `json:"queryTimeoutSeconds,omitempty"`
+	QueryTimeoutMethod  string  `json:"queryTimeoutMethod,omitempty"`
 	TxTimeoutSeconds    Seconds `json:"txTimeoutSeconds,omitempty"`
 	MaxRows             int     `json:"maxRows,omitempty"`
 	WarnRows            int     `json:"warnRows,omitempty"`
@@ -518,6 +522,7 @@ var defaultConfig = TabletConfig{
 	},
 	Oltp: OltpConfig{
 		QueryTimeoutSeconds: 30,
+		QueryTimeoutMethod:  DefaultQueryTimeoutMethod,
 		TxTimeoutSeconds:    30,
 		MaxRows:             10000,
 	},

go/vt/vttablet/tabletserver/tabletserver.go

@@ -159,7 +159,12 @@ func NewTabletServer(name string, config *tabletenv.TabletConfig, topoServer *to
 		topoServer:             topoServer,
 		alias:                  proto.Clone(alias).(*topodatapb.TabletAlias),
 	}
-	tsv.QueryTimeout.Store(config.Oltp.QueryTimeoutSeconds.Get().Nanoseconds())
+
+	queryTimeoutNanos := config.Oltp.QueryTimeoutSeconds.Get().Nanoseconds()
+	if config.Oltp.QueryTimeoutMethod == "mysql" {
+		queryTimeoutNanos = queryTimeoutNanos + queryTimeoutMysqlMaxWait.Nanoseconds()
+	}
+	tsv.QueryTimeout.Store(queryTimeoutNanos)
 
 	tsOnce.Do(func() { srvTopoServer = srvtopo.NewResilientServer(topoServer, "TabletSrvTopo") })