planner: Add tidb_max_estimated_cost to prevent trying to run too expensive queries

errors.toml

@@ -2756,6 +2756,11 @@ error = '''
 '%s' is unsupported on cache tables.
 '''
 
+["planner:8266"]
+error = '''
+Optimizer cost exceeds tidb_max_estimated_cost: %g > %g
+'''
+
 ["privilege:1045"]
 error = '''
 Access denied for user '%-.48s'@'%-.255s' (using password: %s)

pkg/errno/errcode.go

@@ -1150,6 +1150,7 @@ const (
 	ErrGlobalIndexNotExplicitlySet = 8264
 
 	ErrWarnGlobalIndexNeedManuallyAnalyze = 8265
+	ErrMaxEstimatedCostExceeded           = 8266
 
 	// Resource group errors.
 	ErrResourceGroupExists                    = 8248

pkg/errno/errname.go

@@ -1193,4 +1193,5 @@ var MySQLErrName = map[uint16]*mysql.ErrMessage{
 
 	ErrWarnGlobalIndexNeedManuallyAnalyze: mysql.Message("Auto analyze is not effective for index '%-.192s', need analyze manually", nil),
 	ErrTimeStampInDSTTransition:           mysql.Message("Timestamp is not valid, since it is in Daylight Saving Time transition '%s' for time zone '%s'", nil),
+	ErrMaxEstimatedCostExceeded:           mysql.Message("Optimizer cost exceeds tidb_max_estimated_cost: %g > %g", nil),
 }

pkg/planner/core/optimizer.go

@@ -50,6 +50,7 @@ import (
 	"github.com/pingcap/tidb/pkg/privilege"
 	"github.com/pingcap/tidb/pkg/sessionctx"
 	"github.com/pingcap/tidb/pkg/sessionctx/vardef"
+	"github.com/pingcap/tidb/pkg/sessionctx/variable"
 	"github.com/pingcap/tidb/pkg/types"
 	"github.com/pingcap/tidb/pkg/util"
 	"github.com/pingcap/tidb/pkg/util/dbterror/plannererrors"
@@ -293,6 +294,19 @@ func doOptimize(ctx context.Context, sctx base.PlanContext, flag uint64, logic b
 	return VolcanoOptimize(ctx, sctx, flag, logic)
 }
 
+func checkMaxEstimatedCost(sessVars *variable.SessionVars, cost float64) error {
+	if sessVars.MaxEstimatedCost != 0.0 &&
+		!sessVars.InRestrictedSQL && // Allow internal queries!
+		cost > sessVars.MaxEstimatedCost {
+		// Still allow EXPLAIN, unless EXPLAIN ANALYZE
+		if !sessVars.StmtCtx.InExplainStmt ||
+			sessVars.StmtCtx.InExplainAnalyzeStmt {
+			return plannererrors.ErrMaxEstimatedCostExceeded.FastGenByArgs(cost, sessVars.MaxEstimatedCost)
+		}
+	}
+	return nil
+}
+
 // CascadesOptimize includes: normalization, cascadesOptimize, and physicalOptimize.
 func CascadesOptimize(ctx context.Context, sctx base.PlanContext, flag uint64, logic base.LogicalPlan) (base.LogicalPlan, base.PhysicalPlan, float64, error) {
 	sessVars := sctx.GetSessionVars()
@@ -322,6 +336,10 @@ func CascadesOptimize(ctx context.Context, sctx base.PlanContext, flag uint64, l
 	if err != nil {
 		return nil, nil, 0, err
 	}
+	err = checkMaxEstimatedCost(sessVars, cost)
+	if err != nil {
+		return nil, nil, cost, err
+	}
 
 	finalPlan := postOptimize(ctx, sctx, physical)
 	if sessVars.StmtCtx.EnableOptimizerCETrace {
@@ -1078,7 +1096,8 @@ func isLogicalRuleDisabled(r base.LogicalOptRule) bool {
 }
 
 func physicalOptimize(logic base.LogicalPlan) (plan base.PhysicalPlan, cost float64, err error) {
-	if logic.SCtx().GetSessionVars().StmtCtx.EnableOptimizerDebugTrace {
+	sessVars := logic.SCtx().GetSessionVars()
+	if sessVars.StmtCtx.EnableOptimizerDebugTrace {
 		debugtrace.EnterContextCommon(logic.SCtx())
 		defer debugtrace.LeaveContextCommon(logic.SCtx())
 	}
@@ -1093,26 +1112,30 @@ func physicalOptimize(logic base.LogicalPlan) (plan base.PhysicalPlan, cost floa
 		ExpectedCnt: math.MaxFloat64,
 	}
 
-	logic.SCtx().GetSessionVars().StmtCtx.TaskMapBakTS = 0
+	sessVars.StmtCtx.TaskMapBakTS = 0
 	t, err := physicalop.FindBestTask(logic, prop)
 	if err != nil {
 		return nil, 0, err
 	}
 	if t.Invalid() {
 		errMsg := "Can't find a proper physical plan for this query"
-		if config.GetGlobalConfig().DisaggregatedTiFlash && !logic.SCtx().GetSessionVars().IsMPPAllowed() {
+		if config.GetGlobalConfig().DisaggregatedTiFlash && !sessVars.IsMPPAllowed() {
 			errMsg += ": cop and batchCop are not allowed in disaggregated tiflash mode, you should turn on tidb_allow_mpp switch"
 		}
 		return nil, 0, plannererrors.ErrInternal.GenWithStackByArgs(errMsg)
 	}
 
 	// collect the warnings from task.
-	logic.SCtx().GetSessionVars().StmtCtx.AppendWarnings(t.(*physicalop.RootTask).Warnings.GetWarnings())
+	sessVars.StmtCtx.AppendWarnings(t.(*physicalop.RootTask).Warnings.GetWarnings())
 
 	if err = t.Plan().ResolveIndices(); err != nil {
 		return nil, 0, err
 	}
 	cost, err = getPlanCost(t.Plan(), property.RootTaskType, costusage.NewDefaultPlanCostOption())
+	err = checkMaxEstimatedCost(sessVars, cost)
+	if err != nil {
+		return nil, cost, err
+	}
 	return t.Plan(), cost, err
 }
 

pkg/planner/core/plan_cost_ver1_test.go

@@ -123,3 +123,78 @@ func TestScanOnSmallTable(t *testing.T) {
 	}
 	require.True(t, useTiKVScan, "should use tikv scan, but got:\n%s", resStr)
 }
+
+func TestMaxEstimatedCost(t *testing.T) {
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+
+	tk.MustExec(`set @@tidb_cost_model_version=2`)
+	tk.MustExec("use test")
+	tk.MustExec(`create table t (a int auto_increment primary key, b int, filler varchar(255))`)
+	tk.MustExec(`insert into t (b) values (1),(2),(3),(4)`)
+	tk.MustExec(`insert into t (b) select t.b + t1.b + t2.b + t3.b + t4.b + t5.b from t, t t1, t t2, t t3, t t4, t t5`)
+	tk.MustExec(`analyze table t`)
+
+	// Using Max Estimated Cost: 60e6, which should theoretically be in the order of 1 minute
+	res := tk.MustQuery(`explain format=verbose select /*+ SET_VAR(tidb_max_estimated_cost="60e6") */ t.b + t1.b + t2.b + t3.b + t4.b + t5.b from t, t t1, t t2, t t3, t t4, t t5 where t.a = 3 and t1.b = 3 and t2.a = 3`)
+	estCost, err := strconv.ParseFloat(res.Rows()[0][2].(string), 64)
+	require.NoError(t, err)
+	require.Greater(t, estCost, 60e6)
+	tk.MustContainErrMsg(`select /*+ SET_VAR(tidb_max_estimated_cost="60e6") */ t.b + t1.b + t2.b + t3.b + t4.b + t5.b from t, t t1, t t2, t t3, t t4, t t5 where t.a = 3 and t1.b = 3 and t2.a = 3`,
+		"[planner:8266]Optimizer cost exceeds tidb_max_estimated_cost: ")
+	tk.MustContainErrMsg(`explain analyze format=verbose select /*+ SET_VAR(tidb_max_estimated_cost="60e6") */ t.b + t1.b + t2.b + t3.b + t4.b + t5.b from t, t t1, t t2, t t3, t t4, t t5 where t.a = 3 and t1.b = 3 and t2.a = 3`,
+		"[planner:8266]Optimizer cost exceeds tidb_max_estimated_cost: ")
+
+	tk.MustExec(`set @@tidb_cost_model_version=1`)
+	res = tk.MustQuery(`explain format=verbose select /*+ SET_VAR(tidb_max_estimated_cost="60e6") */ t.b + t1.b + t2.b + t3.b + t4.b + t5.b from t, t t1, t t2, t t3, t t4, t t5 where t.a = 3 and t1.b = 3 and t2.a = 3`)
+	estCost, err = strconv.ParseFloat(res.Rows()[0][2].(string), 64)
+	require.NoError(t, err)
+	require.Greater(t, estCost, 60e6)
+	tk.MustContainErrMsg(`select /*+ SET_VAR(tidb_max_estimated_cost="60e6") */ t.b + t1.b + t2.b + t3.b + t4.b + t5.b from t, t t1, t t2, t t3, t t4, t t5 where t.a = 3 and t1.b = 3 and t2.a = 3`,
+		"[planner:8266]Optimizer cost exceeds tidb_max_estimated_cost: ")
+	tk.MustContainErrMsg(`explain analyze format=verbose select /*+ SET_VAR(tidb_max_estimated_cost="60e6") */ t.b + t1.b + t2.b + t3.b + t4.b + t5.b from t, t t1, t t2, t t3, t t4, t t5 where t.a = 3 and t1.b = 3 and t2.a = 3`,
+		"[planner:8266]Optimizer cost exceeds tidb_max_estimated_cost: ")
+	tk.MustExec(`set @@tidb_cost_model_version=2`)
+	tk.MustExec(`alter table t add index idx_b (b)`)
+	tk.MustExec(`analyze table t`)
+
+	res = tk.MustQuery(`explain format=verbose select /*+ SET_VAR(tidb_max_estimated_cost="60e6") */ t.b + t1.b + t2.b + t3.b + t4.b + t5.b from t, t t1, t t2, t t3, t t4, t t5 where t.a = 3 and t1.b = 3 and t2.a = 3`)
+	estCost, err = strconv.ParseFloat(res.Rows()[0][2].(string), 64)
+	require.NoError(t, err)
+	require.Greater(t, estCost, 60e6)
+	tk.MustContainErrMsg(`select /*+ SET_VAR(tidb_max_estimated_cost="60e6") */ t.b + t1.b + t2.b + t3.b + t4.b + t5.b from t, t t1, t t2, t t3, t t4, t t5 where t.a = 3 and t1.b = 3 and t2.a = 3`,
+		"[planner:8266]Optimizer cost exceeds tidb_max_estimated_cost: ")
+	tk.MustContainErrMsg(`explain analyze format=verbose select /*+ SET_VAR(tidb_max_estimated_cost="60e6") */ t.b + t1.b + t2.b + t3.b + t4.b + t5.b from t, t t1, t t2, t t3, t t4, t t5 where t.a = 3 and t1.b = 3 and t2.a = 3`,
+		"[planner:8266]Optimizer cost exceeds tidb_max_estimated_cost: ")
+
+	tk.MustExec(`set @@tidb_cost_model_version=1`)
+	res = tk.MustQuery(`explain format=verbose select /*+ SET_VAR(tidb_max_estimated_cost="60e6") */ t.b + t1.b + t2.b + t3.b + t4.b + t5.b from t, t t1, t t2, t t3, t t4, t t5 where t.a = 3 and t1.b = 3 and t2.a = 3`)
+	estCost, err = strconv.ParseFloat(res.Rows()[0][2].(string), 64)
+	require.NoError(t, err)
+	require.Greater(t, estCost, 60e6)
+	tk.MustContainErrMsg(`select /*+ SET_VAR(tidb_max_estimated_cost="60e6") */ t.b + t1.b + t2.b + t3.b + t4.b + t5.b from t, t t1, t t2, t t3, t t4, t t5 where t.a = 3 and t1.b = 3 and t2.a = 3`,
+		"[planner:8266]Optimizer cost exceeds tidb_max_estimated_cost: ")
+	tk.MustContainErrMsg(`explain analyze select /*+ SET_VAR(tidb_max_estimated_cost="60e6") */ t.b + t1.b + t2.b + t3.b + t4.b + t5.b from t, t t1, t t2, t t3, t t4, t t5 where t.a = 3 and t1.b = 3 and t2.a = 3`,
+		"[planner:8266]Optimizer cost exceeds tidb_max_estimated_cost: ")
+
+	tk.MustExec(`set tidb_max_estimated_cost=60e6`)
+	tk.MustExec(`set @@tidb_cost_model_version=2`)
+	res = tk.MustQuery(`explain format=verbose select t.b + t1.b + t2.b + t3.b + t4.b + t5.b from t, t t1, t t2, t t3, t t4, t t5 where t.a = 3 and t1.b = 3 and t2.a = 3`)
+	estCost, err = strconv.ParseFloat(res.Rows()[0][2].(string), 64)
+	require.NoError(t, err)
+	require.Greater(t, estCost, 60e6)
+	tk.MustContainErrMsg(`select t.b + t1.b + t2.b + t3.b + t4.b + t5.b from t, t t1, t t2, t t3, t t4, t t5 where t.a = 3 and t1.b = 3 and t2.a = 3`,
+		"[planner:8266]Optimizer cost exceeds tidb_max_estimated_cost: ")
+	tk.MustContainErrMsg(`explain analyze format=verbose select t.b + t1.b + t2.b + t3.b + t4.b + t5.b from t, t t1, t t2, t t3, t t4, t t5 where t.a = 3 and t1.b = 3 and t2.a = 3`,
+		"[planner:8266]Optimizer cost exceeds tidb_max_estimated_cost: ")
+
+	tk.MustExec(`set @@tidb_cost_model_version=1`)
+	res = tk.MustQuery(`explain format=verbose select t.b + t1.b + t2.b + t3.b + t4.b + t5.b from t, t t1, t t2, t t3, t t4, t t5 where t.a = 3 and t1.b = 3 and t2.a = 3`)
+	estCost, err = strconv.ParseFloat(res.Rows()[0][2].(string), 64)
+	require.NoError(t, err)
+	require.Greater(t, estCost, 60e6)
+	tk.MustContainErrMsg(`select t.b + t1.b + t2.b + t3.b + t4.b + t5.b from t, t t1, t t2, t t3, t t4, t t5 where t.a = 3 and t1.b = 3 and t2.a = 3`,
+		"[planner:8266]Optimizer cost exceeds tidb_max_estimated_cost: ")
+	tk.MustContainErrMsg(`explain analyze select t.b + t1.b + t2.b + t3.b + t4.b + t5.b from t, t t1, t t2, t t3, t t4, t t5 where t.a = 3 and t1.b = 3 and t2.a = 3`,
+		"[planner:8266]Optimizer cost exceeds tidb_max_estimated_cost: ")
+}

pkg/sessionctx/vardef/tidb_vars.go

@@ -1079,6 +1079,10 @@ const (
 
 	// TiDBAccelerateUserCreationUpdate decides whether tidb will load & update the whole user's data in-memory.
 	TiDBAccelerateUserCreationUpdate = "tidb_accelerate_user_creation_update"
+
+	// TiDBMaxEstimatedCost set a maximum cost of plans from the optimizer to be executed. Higher costs estimates
+	// would simply return an error, stating that the estimated costs is too high.
+	TiDBMaxEstimatedCost = "tidb_max_estimated_cost"
 )
 
 // TiDB vars that have only global scope
@@ -1745,6 +1749,7 @@ const (
 	DefTiDBLoadBindingTimeout                         = 200
 	DefTiDBEnableBindingUsage                         = true
 	DefTiDBAdvancerCheckPointLagLimit                 = 48 * time.Hour
+	DefTiDBMaxEstimatedCost                           = 0.0
 	DefTiDBMemArbitratorSoftLimitText                 = memory.ArbitratorSoftLimitModDisableName
 	DefTiDBMemArbitratorModeText                      = memory.ArbitratorModeDisableName
 	DefTiDBMemArbitratorQueryReservedText             = "0"

pkg/sessionctx/variable/session.go

@@ -1088,6 +1088,10 @@ type SessionVars struct {
 	IndexJoinCostFactor        float64
 	SelectivityFactor          float64
 
+	// MaxEstimatedCost is the maximum estimated cost for a DML statement.
+	// If the value is 0.0, then not enabled.
+	MaxEstimatedCost float64
+
 	// enableForceInlineCTE is used to enable/disable force inline CTE.
 	enableForceInlineCTE bool
 

pkg/sessionctx/variable/sysvar.go

@@ -3800,6 +3800,20 @@ var defaultSysVars = []*SysVar{
 			return vardef.AdvancerCheckPointLagLimit.Load().String(), nil
 		},
 	},
+	{
+		Scope: vardef.ScopeGlobal | vardef.ScopeSession,
+		Name:  vardef.TiDBMaxEstimatedCost,
+		Value: strconv.FormatFloat(vardef.DefTiDBMaxEstimatedCost, 'f', -1, 64),
+		Type:  vardef.TypeFloat,
+		SetSession: func(vars *SessionVars, s string) error {
+			vars.MaxEstimatedCost = tidbOptFloat64(s, vardef.DefTiDBMaxEstimatedCost)
+			return nil
+		},
+		MinValue: 0,
+		MaxValue: math.MaxUint64,
+		// because the special character in custom syntax cannot be correctly handled in set_var hint
+		IsHintUpdatableVerified: true,
+	},
 	{Scope: vardef.ScopeGlobal | vardef.ScopeSession, Name: vardef.TiDBSlowLogRules, Value: "", Type: vardef.TypeStr,
 		SetSession: func(s *SessionVars, val string) error {
 			slowLogRules, err := ParseSessionSlowLogRules(val)

pkg/sessionctx/variable/sysvar_test.go

@@ -1949,3 +1949,49 @@ func TestTiDBOptSelectivityFactor(t *testing.T) {
 	warn := vars.StmtCtx.GetWarnings()[0].Err
 	require.Equal(t, "[variable:1292]Truncated incorrect tidb_opt_selectivity_factor value: '1.1'", warn.Error())
 }
+
+func TestTiDBMaxEstimatedCostVar(t *testing.T) {
+	ctx := context.Background()
+	vars := NewSessionVars(nil)
+	mock := NewMockGlobalAccessor4Tests()
+	mock.SessionVars = vars
+	vars.GlobalVarsAccessor = mock
+	val, err := vars.GetSessionOrGlobalSystemVar(context.Background(), vardef.TiDBMaxEstimatedCost)
+	require.NoError(t, err)
+	require.Equal(t, "0", val)
+
+	// set valid SESSION value
+	require.NoError(t, vars.SetSystemVar(vardef.TiDBMaxEstimatedCost, "9.7"))
+	val, err = vars.GetSessionOrGlobalSystemVar(context.Background(), vardef.TiDBMaxEstimatedCost)
+	require.NoError(t, err)
+	require.Equal(t, "9.7", val)
+	require.Len(t, vars.StmtCtx.GetWarnings(), 0)
+
+	// set valid GLOBAL value
+	err = mock.SetGlobalSysVar(ctx, vardef.TiDBMaxEstimatedCost, "111.1")
+	require.NoError(t, err)
+	val, err = mock.GetGlobalSysVar(vardef.TiDBMaxEstimatedCost)
+	require.NoError(t, err)
+	require.Equal(t, "111.1", val)
+	require.Len(t, vars.StmtCtx.GetWarnings(), 0)
+
+	// set invalid value
+	err = mock.SetGlobalSysVar(ctx, vardef.TiDBMaxEstimatedCost, "-0.1")
+	require.NoError(t, err)
+	require.Len(t, vars.StmtCtx.GetWarnings(), 1)
+	warn := vars.StmtCtx.GetWarnings()[0].Err
+	require.Equal(t, "[variable:1292]Truncated incorrect tidb_max_estimated_cost value: '-0.1'", warn.Error())
+	val, err = mock.GetGlobalSysVar(vardef.TiDBMaxEstimatedCost)
+	require.NoError(t, err)
+	// TODO: in case of error/warning/truncation, should we really set it to DEFAULT / disable it? Not leave it untouched?
+	require.Equal(t, "0", val)
+	err = vars.SetSystemVar(vardef.TiDBMaxEstimatedCost, "10e64")
+	require.NoError(t, err)
+	// Nothing is clearing the warnings
+	require.Len(t, vars.StmtCtx.GetWarnings(), 2)
+	warn = vars.StmtCtx.GetWarnings()[1].Err
+	require.Equal(t, "[variable:1292]Truncated incorrect tidb_max_estimated_cost value: '10e64'", warn.Error())
+	val, err = mock.GetGlobalSysVar(vardef.TiDBMaxEstimatedCost)
+	require.NoError(t, err)
+	require.Equal(t, "0", val)
+}

pkg/util/dbterror/plannererrors/planner_terror.go

@@ -121,4 +121,6 @@ var (
 	ErrRowIsReferenced2 = dbterror.ClassOptimizer.NewStd(mysql.ErrRowIsReferenced2)
 	ErrNoReferencedRow2 = dbterror.ClassOptimizer.NewStd(mysql.ErrNoReferencedRow2)
 	ErrSpDoesNotExist   = dbterror.ClassOptimizer.NewStd(mysql.ErrSpDoesNotExist)
+
+	ErrMaxEstimatedCostExceeded = dbterror.ClassOptimizer.NewStd(mysql.ErrMaxEstimatedCostExceeded)
 )