planner: support using "KiB/MiB/GiB" to set isntance_plan_cache_targe…

…t/max_mem_size (#55316) ref #54057
pingcap · Aug 9, 2024 · 2a731e6 · 2a731e6
1 parent 829bd5b
commit 2a731e6
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 8 deletions.
diff --git a/pkg/executor/set_test.go b/pkg/executor/set_test.go
@@ -776,6 +776,12 @@ func TestSetVar(t *testing.T) {
 	tk.MustQuery("select @@global.tidb_instance_plan_cache_target_mem_size").Check(testkit.Rows("114857600"))
 	tk.MustExec("set global tidb_instance_plan_cache_max_mem_size = 135829120")
 	tk.MustQuery("select @@global.tidb_instance_plan_cache_max_mem_size").Check(testkit.Rows("135829120"))
+	tk.MustExec("set global tidb_instance_plan_cache_max_mem_size = 1GiB")
+	tk.MustQuery("select @@global.tidb_instance_plan_cache_max_mem_size").Check(testkit.Rows("1073741824"))
+	tk.MustExec("set global tidb_instance_plan_cache_target_mem_size = 999MiB")
+	tk.MustQuery("select @@global.tidb_instance_plan_cache_target_mem_size").Check(testkit.Rows("1047527424"))
+	tk.MustExec("set global tidb_instance_plan_cache_target_mem_size = 998MiB")
+	tk.MustQuery("select @@global.tidb_instance_plan_cache_target_mem_size").Check(testkit.Rows("1046478848"))
 
 	// test variables for cost model ver2
 	tk.MustQuery("select @@tidb_cost_model_version").Check(testkit.Rows(fmt.Sprintf("%v", variable.DefTiDBCostModelVer)))

diff --git a/pkg/sessionctx/variable/sysvar.go b/pkg/sessionctx/variable/sysvar.go
@@ -1344,28 +1344,34 @@ var defaultSysVars = []*SysVar{
 			EnableInstancePlanCache.Store(TiDBOptOn(val))
 			return nil
 		}},
-	{Scope: ScopeGlobal, Name: TiDBInstancePlanCacheTargetMemSize, Value: strconv.Itoa(int(DefTiDBInstancePlanCacheTargetMemSize)), Type: TypeInt, MinValue: 1, MaxValue: math.MaxInt32,
+	{Scope: ScopeGlobal, Name: TiDBInstancePlanCacheTargetMemSize, Value: strconv.Itoa(int(DefTiDBInstancePlanCacheTargetMemSize)), Type: TypeStr,
 		GetGlobal: func(_ context.Context, s *SessionVars) (string, error) {
 			return strconv.FormatInt(InstancePlanCacheTargetMemSize.Load(), 10), nil
 		},
 		SetGlobal: func(_ context.Context, s *SessionVars, val string) error {
-			v := TidbOptInt64(val, int64(DefTiDBInstancePlanCacheTargetMemSize))
-			if v > InstancePlanCacheMaxMemSize.Load() {
+			v, str := parseByteSize(val)
+			if str == "" {
+				v = uint64(TidbOptInt64(val, int64(DefTiDBInstancePlanCacheTargetMemSize)))
+			}
+			if v > uint64(InstancePlanCacheMaxMemSize.Load()) {
 				return errors.New("tidb_instance_plan_cache_target_mem_size must be less than tidb_instance_plan_cache_max_mem_size")
 			}
-			InstancePlanCacheTargetMemSize.Store(v)
+			InstancePlanCacheTargetMemSize.Store(int64(v))
 			return nil
 		}},
-	{Scope: ScopeGlobal, Name: TiDBInstancePlanCacheMaxMemSize, Value: strconv.Itoa(int(DefTiDBInstancePlanCacheMaxMemSize)), Type: TypeInt, MinValue: 1, MaxValue: math.MaxInt32,
+	{Scope: ScopeGlobal, Name: TiDBInstancePlanCacheMaxMemSize, Value: strconv.Itoa(int(DefTiDBInstancePlanCacheMaxMemSize)), Type: TypeStr,
 		GetGlobal: func(_ context.Context, s *SessionVars) (string, error) {
 			return strconv.FormatInt(InstancePlanCacheMaxMemSize.Load(), 10), nil
 		},
 		SetGlobal: func(_ context.Context, s *SessionVars, val string) error {
-			v := TidbOptInt64(val, int64(DefTiDBInstancePlanCacheMaxMemSize))
-			if v < InstancePlanCacheTargetMemSize.Load() {
+			v, str := parseByteSize(val)
+			if str == "" {
+				v = uint64(TidbOptInt64(val, int64(DefTiDBInstancePlanCacheTargetMemSize)))
+			}
+			if v < uint64(InstancePlanCacheTargetMemSize.Load()) {
 				return errors.New("tidb_instance_plan_cache_max_mem_size must be greater than tidb_instance_plan_cache_target_mem_size")
 			}
-			InstancePlanCacheMaxMemSize.Store(v)
+			InstancePlanCacheMaxMemSize.Store(int64(v))
 			return nil
 		}},
 	{Scope: ScopeGlobal, Name: TiDBMemOOMAction, Value: DefTiDBMemOOMAction, PossibleValues: []string{"CANCEL", "LOG"}, Type: TypeEnum,

diff --git a/pkg/sessionctx/variable/varsutil.go b/pkg/sessionctx/variable/varsutil.go
@@ -409,15 +409,27 @@ func parseByteSize(s string) (byteSize uint64, normalizedStr string) {
 	if n, err := fmt.Sscanf(s, "%dKB%s", &byteSize, &endString); n == 1 && err == io.EOF {
 		return byteSize << 10, fmt.Sprintf("%dKB", byteSize)
 	}
+	if n, err := fmt.Sscanf(s, "%dKiB%s", &byteSize, &endString); n == 1 && err == io.EOF {
+		return byteSize << 10, fmt.Sprintf("%dKiB", byteSize)
+	}
 	if n, err := fmt.Sscanf(s, "%dMB%s", &byteSize, &endString); n == 1 && err == io.EOF {
 		return byteSize << 20, fmt.Sprintf("%dMB", byteSize)
 	}
+	if n, err := fmt.Sscanf(s, "%dMiB%s", &byteSize, &endString); n == 1 && err == io.EOF {
+		return byteSize << 20, fmt.Sprintf("%dMiB", byteSize)
+	}
 	if n, err := fmt.Sscanf(s, "%dGB%s", &byteSize, &endString); n == 1 && err == io.EOF {
 		return byteSize << 30, fmt.Sprintf("%dGB", byteSize)
 	}
+	if n, err := fmt.Sscanf(s, "%dGiB%s", &byteSize, &endString); n == 1 && err == io.EOF {
+		return byteSize << 30, fmt.Sprintf("%dGiB", byteSize)
+	}
 	if n, err := fmt.Sscanf(s, "%dTB%s", &byteSize, &endString); n == 1 && err == io.EOF {
 		return byteSize << 40, fmt.Sprintf("%dTB", byteSize)
 	}
+	if n, err := fmt.Sscanf(s, "%dTiB%s", &byteSize, &endString); n == 1 && err == io.EOF {
+		return byteSize << 40, fmt.Sprintf("%dTiB", byteSize)
+	}
 	return 0, ""
 }