add more params

Signed-off-by: husharp <[email protected]>

go.mod

@@ -327,6 +327,7 @@ replace (
 	// fix potential security issue(CVE-2020-26160) introduced by indirect dependency.
 	github.com/dgrijalva/jwt-go => github.com/form3tech-oss/jwt-go v3.2.6-0.20210809144907-32ab6a8243d7+incompatible
 	github.com/go-ldap/ldap/v3 => github.com/YangKeao/ldap/v3 v3.4.5-0.20230421065457-369a3bab1117
+	github.com/pingcap/kvproto => github.com/HuSharp/kvproto v0.0.0-20240902044512-f2444fdd4baf
 	github.com/pingcap/tidb/pkg/parser => ./pkg/parser
 
 	// TODO: `sourcegraph.com/sourcegraph/appdash` has been archived, and the original host has been removed.

pkg/ddl/resource_group.go

@@ -220,6 +220,13 @@ func SetDirectResourceGroupSettings(groupInfo *model.ResourceGroupInfo, opt *ast
 		if len(opt.RunawayOptionList) == 0 {
 			resourceGroupSettings.Runaway = nil
 		}
+		if resourceGroupSettings.Runaway == nil {
+			resourceGroupSettings.Runaway = &model.ResourceGroupRunawaySettings{}
+		} else { // reset the runaway rule settings
+			resourceGroupSettings.Runaway.ExecElapsedTimeMs = 0
+			resourceGroupSettings.Runaway.ProcessedKeys = 0
+			resourceGroupSettings.Runaway.RequestUnit = 0
+		}
 		for _, opt := range opt.RunawayOptionList {
 			if err := SetDirectResourceGroupRunawayOption(resourceGroupSettings, opt); err != nil {
 				return err
@@ -257,18 +264,23 @@ func SetDirectResourceGroupRUSecondOption(resourceGroupSettings *model.ResourceG
 
 // SetDirectResourceGroupRunawayOption tries to set runaway part of the ResourceGroupSettings.
 func SetDirectResourceGroupRunawayOption(resourceGroupSettings *model.ResourceGroupSettings, opt *ast.ResourceGroupRunawayOption) error {
-	if resourceGroupSettings.Runaway == nil {
-		resourceGroupSettings.Runaway = &model.ResourceGroupRunawaySettings{}
-	}
 	settings := resourceGroupSettings.Runaway
 	switch opt.Tp {
 	case ast.RunawayRule:
-		// because execute time won't be too long, we use `time` pkg which does not support to parse unit 'd'.
-		dur, err := time.ParseDuration(opt.RuleOption.ExecElapsed)
-		if err != nil {
-			return err
+		switch opt.RuleOption.Tp {
+		case ast.RunawayRuleExecElapsed:
+			// because execute time won't be too long, we use `time` pkg which does not support to parse unit 'd'.
+			dur, err := time.ParseDuration(opt.RuleOption.ExecElapsed)
+			if err != nil {
+				return err
+			}
+			settings.ExecElapsedTimeMs = uint64(dur.Milliseconds())
+		case ast.RunawayRuleProcessedKeys:
+			settings.ProcessedKeys = opt.RuleOption.ProcessedKeys
+		case ast.RunawayRuleRequestUnit:
+			settings.RequestUnit = opt.RuleOption.RequestUnit
 		}
-		settings.ExecElapsedTimeMs = uint64(dur.Milliseconds())
+
 	case ast.RunawayAction:
 		settings.Action = opt.ActionOption.Type
 	case ast.RunawayWatch:

pkg/ddl/resourcegroup/errors.go

@@ -31,8 +31,8 @@ var (
 	ErrUnknownResourceGroupMode = errors.New("unknown resource group mode")
 	// ErrDroppingInternalResourceGroup is from group.go
 	ErrDroppingInternalResourceGroup = errors.New("can't drop reserved resource group")
-	// ErrInvalidResourceGroupRunawayExecElapsedTime is from group.go.
-	ErrInvalidResourceGroupRunawayExecElapsedTime = errors.New("invalid exec elapsed time")
+	// ErrResourceGroupRunawayRuleIsEmpty is from group.go.
+	ErrResourceGroupRunawayRuleIsEmpty = errors.New("please set at least one field(exec_elapsed_time_ms, processed_keys, request_unit)")
 	// ErrUnknownResourceGroupRunawayAction is from group.go.
 	ErrUnknownResourceGroupRunawayAction = errors.New("unknown resource group runaway action")
 )

pkg/ddl/resourcegroup/group.go

@@ -37,13 +37,15 @@ func NewGroupFromOptions(groupName string, options *model.ResourceGroupSettings)
 
 	group.Priority = uint32(options.Priority)
 	if options.Runaway != nil {
+		if options.Runaway.ExecElapsedTimeMs == 0 && options.Runaway.ProcessedKeys == 0 && options.Runaway.RequestUnit == 0 {
+			return nil, ErrResourceGroupRunawayRuleIsEmpty
+		}
 		runaway := &rmpb.RunawaySettings{
 			Rule: &rmpb.RunawayRule{},
 		}
-		if options.Runaway.ExecElapsedTimeMs == 0 {
-			return nil, ErrInvalidResourceGroupRunawayExecElapsedTime
-		}
 		runaway.Rule.ExecElapsedTimeMs = options.Runaway.ExecElapsedTimeMs
+		runaway.Rule.ProcessedKeys = options.Runaway.ProcessedKeys
+		runaway.Rule.RequestUnit = options.Runaway.RequestUnit
 		if options.Runaway.Action == model.RunawayActionNone {
 			return nil, ErrUnknownResourceGroupRunawayAction
 		}

pkg/ddl/tests/resourcegroup/resource_group_test.go

@@ -211,7 +211,7 @@ func TestResourceGroupBasic(t *testing.T) {
 	tk.MustContainErrMsg("create resource group x ru_per_sec=1000 burstable QUERY_LIMIT=(EXEC_ELAPSED='15s' action kill action cooldown)", "Dupliated runaway options specified")
 	tk.MustContainErrMsg("create resource group x ru_per_sec=1000 QUERY_LIMIT=(EXEC_ELAPSED='15s') burstable priority=Low, QUERY_LIMIT=(EXEC_ELAPSED='15s')", "Dupliated options specified")
 	tk.MustContainErrMsg("create resource group x ru_per_sec=1000 QUERY_LIMIT=(EXEC_ELAPSED='15s') QUERY_LIMIT=(EXEC_ELAPSED='15s')", "Dupliated options specified")
-	tk.MustContainErrMsg("create resource group x ru_per_sec=1000 QUERY_LIMIT=(action kill)", "invalid exec elapsed time")
+	tk.MustContainErrMsg("create resource group x ru_per_sec=1000 QUERY_LIMIT=(action kill)", "please set at least one field(exec_elapsed_time_ms, process_keys, request_unit)")
 	tk.MustGetErrCode("create resource group x ru_per_sec=1000 QUERY_LIMIT=(EXEC_ELAPSED='15s' action kil)", mysql.ErrParse)
 	tk.MustContainErrMsg("create resource group x ru_per_sec=1000 QUERY_LIMIT=(EXEC_ELAPSED='15s')", "unknown resource group runaway action")
 	tk.MustGetErrCode("create resource group x ru_per_sec=1000 EXEC_ELAPSED='15s' action kill", mysql.ErrParse)
@@ -223,6 +223,13 @@ func TestResourceGroupBasic(t *testing.T) {
 	// Check information schema table information_schema.resource_groups
 	tk.MustExec("create resource group x RU_PER_SEC=1000 PRIORITY=LOW")
 	tk.MustQuery("select * from information_schema.resource_groups where name = 'x'").Check(testkit.Rows("x 1000 LOW NO <nil> <nil>"))
+	tk.MustExec("alter resource group x RU_PER_SEC=2000 BURSTABLE QUERY_LIMIT=(EXEC_ELAPSED='15s' PROCESSED_KEYS=100 action kill)")
+	tk.MustQuery("select * from information_schema.resource_groups where name = 'x'").Check(testkit.Rows("x 2000 LOW YES EXEC_ELAPSED='15s', PROCESSED_KEYS=100, ACTION=KILL <nil>"))
+	tk.MustExec("alter resource group x RU_PER_SEC=2000 BURSTABLE QUERY_LIMIT=(EXEC_ELAPSED='15s' PROCESSED_KEYS=100 REQUEST_UNIT=100 action kill)")
+	tk.MustQuery("select * from information_schema.resource_groups where name = 'x'").Check(testkit.Rows("x 2000 LOW YES EXEC_ELAPSED='15s', PROCESSED_KEYS=100, REQUEST_UNIT=100, ACTION=KILL <nil>"))
+	tk.MustExec("alter resource group x RU_PER_SEC=2000 BURSTABLE QUERY_LIMIT=(PROCESSED_KEYS=200 REQUEST_UNIT=300 action kill)")
+	tk.MustQuery("select * from information_schema.resource_groups where name = 'x'").Check(testkit.Rows("x 2000 LOW YES PROCESSED_KEYS=200, REQUEST_UNIT=300, ACTION=KILL <nil>"))
+	tk.MustQuery("show create resource group x").Check(testkit.Rows("x CREATE RESOURCE GROUP `x` RU_PER_SEC=2000, PRIORITY=LOW, BURSTABLE, QUERY_LIMIT=(PROCESSED_KEYS=200 REQUEST_UNIT=300 ACTION=KILL)"))
 	tk.MustExec("alter resource group x RU_PER_SEC=2000 BURSTABLE QUERY_LIMIT=(EXEC_ELAPSED='15s' action kill)")
 	tk.MustQuery("select * from information_schema.resource_groups where name = 'x'").Check(testkit.Rows("x 2000 LOW YES EXEC_ELAPSED='15s', ACTION=KILL <nil>"))
 	tk.MustQuery("show create resource group x").Check(testkit.Rows("x CREATE RESOURCE GROUP `x` RU_PER_SEC=2000, PRIORITY=LOW, BURSTABLE, QUERY_LIMIT=(EXEC_ELAPSED=\"15s\" ACTION=KILL)"))

pkg/domain/resourcegroup/BUILD.bazel

@@ -15,6 +15,7 @@ go_library(
         "@com_github_prometheus_client_golang//prometheus",
         "@com_github_tikv_client_go_v2//tikv",
         "@com_github_tikv_client_go_v2//tikvrpc",
+        "@com_github_tikv_client_go_v2//util",
         "@com_github_tikv_pd_client//resource_group/controller",
         "@org_uber_go_zap//:zap",
     ],