update

pingcap · Nov 14, 2024 · 48335a6 · 48335a6
1 parent a276364
commit 48335a6
Show file tree

Hide file tree

Showing 6 changed files with 53 additions and 0 deletions.
diff --git a/ddl/ingest/config.go b/ddl/ingest/config.go
@@ -22,6 +22,7 @@ import (
 	"github.com/pingcap/tidb/br/pkg/lightning/checkpoints"
 	lightning "github.com/pingcap/tidb/br/pkg/lightning/config"
 	tidb "github.com/pingcap/tidb/config"
+	"github.com/pingcap/tidb/sessionctx/variable"
 	"github.com/pingcap/tidb/util/logutil"
 	"github.com/pingcap/tidb/util/size"
 	"go.uber.org/zap"
@@ -42,6 +43,7 @@ func genConfig(memRoot MemRoot, jobID int64, unique bool) (*Config, error) {
 	cfg.TikvImporter.Backend = lightning.BackendLocal
 	// Each backend will build a single dir in lightning dir.
 	cfg.TikvImporter.SortedKVDir = filepath.Join(LitSortPath, EncodeBackendTag(jobID))
+	cfg.TikvImporter.StoreWriteBWLimit = lightning.ByteSize(variable.DDLReorgMaxWriteSpeed.Load())
 	if ImporterRangeConcurrencyForTest != nil {
 		cfg.TikvImporter.RangeConcurrency = int(ImporterRangeConcurrencyForTest.Load())
 	}

diff --git a/executor/BUILD.bazel b/executor/BUILD.bazel
@@ -446,6 +446,7 @@ go_test(
         "//util/tableutil",
         "//util/timeutil",
         "//util/topsql/state",
+        "@com_github_docker_go_units//:go-units",
         "@com_github_golang_protobuf//proto",
         "@com_github_gorilla_mux//:mux",
         "@com_github_pingcap_errors//:errors",

diff --git a/executor/ddl_test.go b/executor/ddl_test.go
@@ -23,6 +23,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/docker/go-units"
 	"github.com/pingcap/failpoint"
 	"github.com/pingcap/tidb/ddl/schematracker"
 	ddltestutil "github.com/pingcap/tidb/ddl/testutil"
@@ -1318,6 +1319,31 @@ func TestSetDDLErrorCountLimit(t *testing.T) {
 	res.Check(testkit.Rows("100"))
 }
 
+func TestSetDDLReorgMaxWriteSpeed(t *testing.T) {
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+	require.Equal(t, int64(variable.DefTiDBDDLReorgMaxWriteSpeed), variable.DDLReorgMaxWriteSpeed.Load())
+
+	// valid values
+	for _, val := range []int64{1, 0, 100, 1024 * 1024, 2147483647, units.PiB} {
+		tk.MustExec(fmt.Sprintf("set @@global.tidb_ddl_reorg_max_write_speed = %d", val))
+		require.Equal(t, val, variable.DDLReorgMaxWriteSpeed.Load())
+		tk.MustQuery("select @@global.tidb_ddl_reorg_max_write_speed").Check(testkit.Rows(strconv.FormatInt(val, 10)))
+	}
+	for _, val := range []string{"1", "0", "100", "2KB", "3MiB", "4 gb", "2147483647", "1125899906842624" /* 1PiB */} {
+		tk.MustExec(fmt.Sprintf("set @@global.tidb_ddl_reorg_max_write_speed = '%s'", val))
+		expected, err := units.RAMInBytes(val)
+		require.NoError(t, err)
+		require.Equal(t, expected, variable.DDLReorgMaxWriteSpeed.Load())
+		tk.MustQuery("select @@global.tidb_ddl_reorg_max_write_speed").Check(testkit.Rows(strconv.FormatInt(expected, 10)))
+	}
+
+	// invalid values
+	tk.MustExecToErr("set @@global.tidb_ddl_reorg_max_write_speed = -1")
+	tk.MustExecToErr("set @@global.tidb_ddl_reorg_max_write_speed = invalid_val")
+	tk.MustExecToErr("set @@global.tidb_ddl_reorg_max_write_speed = %d", units.PiB+1)
+}
+
 func TestLoadDDLDistributeVars(t *testing.T) {
 	store := testkit.CreateMockStore(t)
 	tk := testkit.NewTestKit(t, store)

diff --git a/sessionctx/variable/BUILD.bazel b/sessionctx/variable/BUILD.bazel
@@ -62,6 +62,7 @@ go_library(
         "//util/tls",
         "//util/topsql/state",
         "//util/versioninfo",
+        "@com_github_docker_go_units//:go-units",
         "@com_github_pingcap_errors//:errors",
         "@com_github_tikv_client_go_v2//config",
         "@com_github_tikv_client_go_v2//kv",

diff --git a/sessionctx/variable/sysvar.go b/sessionctx/variable/sysvar.go
@@ -25,6 +25,7 @@ import (
 	"sync/atomic"
 	"time"
 
+	"github.com/docker/go-units"
 	"github.com/pingcap/errors"
 	"github.com/pingcap/tidb/config"
 	"github.com/pingcap/tidb/keyspace"
@@ -688,6 +689,23 @@ var defaultSysVars = []*SysVar{
 		SetDDLReorgBatchSize(int32(tidbOptPositiveInt32(val, DefTiDBDDLReorgBatchSize)))
 		return nil
 	}},
+	{Scope: ScopeGlobal, Name: TiDBDDLReorgMaxWriteSpeed, Value: strconv.Itoa(DefTiDBDDLReorgMaxWriteSpeed), Type: TypeStr,
+		SetGlobal: func(_ context.Context, s *SessionVars, val string) error {
+			i64, err := units.RAMInBytes(val)
+			if err != nil {
+				return errors.Trace(err)
+			}
+			if i64 < 0 || i64 > units.PiB {
+				// Here we limit the max value to 1 PiB instead of math.MaxInt64, since:
+				// 1. it is large enough
+				// 2. units.RAMInBytes would first cast the size to a float, and may lose precision when the size is too large
+				return fmt.Errorf("invalid value for '%d', it should be within [%d, %d]", i64, 0, units.PiB)
+			}
+			DDLReorgMaxWriteSpeed.Store(i64)
+			return nil
+		}, GetGlobal: func(_ context.Context, sv *SessionVars) (string, error) {
+			return strconv.FormatInt(DDLReorgMaxWriteSpeed.Load(), 10), nil
+		}},
 	{Scope: ScopeGlobal, Name: TiDBDDLErrorCountLimit, Value: strconv.Itoa(DefTiDBDDLErrorCountLimit), Type: TypeUnsigned, MinValue: 0, MaxValue: math.MaxInt64, SetGlobal: func(_ context.Context, s *SessionVars, val string) error {
 		SetDDLErrorCountLimit(TidbOptInt64(val, DefTiDBDDLErrorCountLimit))
 		return nil

diff --git a/sessionctx/variable/tidb_vars.go b/sessionctx/variable/tidb_vars.go
@@ -485,6 +485,9 @@ const (
 	// It can be: PRIORITY_LOW, PRIORITY_NORMAL, PRIORITY_HIGH
 	TiDBDDLReorgPriority = "tidb_ddl_reorg_priority"
 
+	// TiDBDDLReorgMaxWriteSpeed defines the max write limitation for the lightning local backend
+	TiDBDDLReorgMaxWriteSpeed = "tidb_ddl_reorg_max_write_speed"
+
 	// TiDBEnableAutoIncrementInGenerated disables the mysql compatibility check on using auto-incremented columns in
 	// expression indexes and generated columns described here https://dev.mysql.com/doc/refman/5.7/en/create-table-generated-columns.html for details.
 	TiDBEnableAutoIncrementInGenerated = "tidb_enable_auto_increment_in_generated"
@@ -1148,6 +1151,7 @@ const (
 	DefTiDBDDLReorgBatchSize                       = 256
 	DefTiDBDDLFlashbackConcurrency                 = 64
 	DefTiDBDDLErrorCountLimit                      = 512
+	DefTiDBDDLReorgMaxWriteSpeed                   = 0
 	DefTiDBMaxDeltaSchemaCount                     = 1024
 	DefTiDBPlacementMode                           = PlacementModeStrict
 	DefTiDBEnableAutoIncrementInGenerated          = false
@@ -1367,6 +1371,7 @@ var (
 	ddlFlashbackConcurrency int32 = DefTiDBDDLFlashbackConcurrency
 	ddlErrorCountLimit      int64 = DefTiDBDDLErrorCountLimit
 	ddlReorgRowFormat       int64 = DefTiDBRowFormatV2
+	DDLReorgMaxWriteSpeed         = atomic.NewInt64(DefTiDBDDLReorgMaxWriteSpeed)
 	maxDeltaSchemaCount     int64 = DefTiDBMaxDeltaSchemaCount
 	// DDLSlowOprThreshold is the threshold for ddl slow operations, uint is millisecond.
 	DDLSlowOprThreshold                  = config.GetGlobalConfig().Instance.DDLSlowOprThreshold