checker(dm): check the number of connections before starting

dm/checker/checker.go

@@ -88,6 +88,8 @@ type Checker struct {
 	warnCnt int64
 
 	onlineDDL onlineddl.OnlinePlugin
+
+	stCfgs []*config.SubTaskConfig
 }
 
 // NewChecker returns a checker.
@@ -97,6 +99,7 @@ func NewChecker(cfgs []*config.SubTaskConfig, checkingItems map[string]string, e
 		checkingItems: checkingItems,
 		errCnt:        errCnt,
 		warnCnt:       warnCnt,
+		stCfgs:        cfgs,
 	}
 
 	for _, cfg := range cfgs {
@@ -144,6 +147,25 @@ func (c *Checker) Init(ctx context.Context) (err error) {
 	if egErr := eg.Wait(); egErr != nil {
 		return egErr
 	}
+	// check connections
+	if _, ok := c.checkingItems[config.ConnAmountChecking]; ok {
+		if len(c.stCfgs) > 0 {
+			switch c.stCfgs[0].Mode {
+			case config.ModeAll:
+				// TODO: check the connections for syncer
+				// TODO: check for incremental mode
+				c.checkList = append(c.checkList, checker.NewLoaderConnAmountChecker(c.instances[0].targetDB, c.stCfgs))
+				for i, inst := range c.instances {
+					c.checkList = append(c.checkList, checker.NewDumperConnAmountChecker(inst.sourceDB, c.stCfgs[i].MydumperConfig.Threads))
+				}
+			case config.ModeFull:
+				c.checkList = append(c.checkList, checker.NewLoaderConnAmountChecker(c.instances[0].targetDB, c.stCfgs))
+				for i, inst := range c.instances {
+					c.checkList = append(c.checkList, checker.NewDumperConnAmountChecker(inst.sourceDB, c.stCfgs[i].MydumperConfig.Threads))
+				}
+			}
+		}
+	}
 	// targetTableID => source => [tables]
 	sharding := make(map[string]map[string][]*filter.Table)
 	shardingCounter := make(map[string]int)

dm/config/checking_item.go

@@ -35,6 +35,7 @@ const (
 	ShardAutoIncrementIDChecking = "auto_increment_ID"
 	OnlineDDLChecking            = "online_ddl"
 	BinlogDBChecking             = "binlog_db"
+	ConnAmountChecking           = "conn_amount"
 )
 
 // AllCheckingItems contains all checking items.
@@ -52,6 +53,7 @@ var AllCheckingItems = map[string]string{
 	ShardAutoIncrementIDChecking: "conflict auto increment ID of shard tables checking item",
 	OnlineDDLChecking:            "online ddl checking item",
 	BinlogDBChecking:             "binlog db checking item",
+	ConnAmountChecking:           "connection number checking item",
 }
 
 // MaxSourceIDLength is the max length for dm-worker source id.

dm/dumpling/dumpling.go

@@ -164,7 +164,6 @@ func (m *Dumpling) Process(ctx context.Context, pr chan pb.ProcessResult) {
 
 	newCtx, cancel := context.WithCancel(ctx)
 	var dumpling *export.Dumper
-
 	if dumpling, err = export.NewDumper(newCtx, m.dumpConfig); err == nil {
 		m.mu.Lock()
 		m.core = dumpling

dm/loader/loader.go

@@ -614,7 +614,12 @@ func (l *Loader) Process(ctx context.Context, pr chan pb.ProcessResult) {
 			errs = append(errs, err)
 		}
 	}()
-
+	failpoint.Inject("longLoadProcess", func(val failpoint.Value) {
+		if sec, ok := val.(int); ok {
+			l.logger.Info("long loader unit", zap.Int("second", sec))
+			time.Sleep(time.Duration(sec) * time.Second)
+		}
+	})
 	err = l.Restore(newCtx)
 	close(l.runFatalChan) // Restore returned, all potential fatal sent to l.runFatalChan
 	cancel()              // cancel the goroutines created in `Restore`.

dm/pkg/checker/conn_checker.go

@@ -0,0 +1,196 @@
+// Copyright 2022 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package checker
+
+import (
+	"context"
+	"database/sql"
+
+	"github.com/pingcap/tiflow/dm/config"
+	"github.com/pingcap/tiflow/dm/pkg/conn"
+	tcontext "github.com/pingcap/tiflow/dm/pkg/context"
+	"github.com/pingcap/tiflow/dm/pkg/log"
+	"go.uber.org/zap"
+)
+
+type connAmountChecker struct {
+	toCheckDB *conn.BaseDB
+	stCfgs    []*config.SubTaskConfig
+
+	getConfigConn func(stCfgs []*config.SubTaskConfig) int
+	workerName    string
+	unlimitedConn bool
+}
+
+func newConnAmountChecker(toCheckDB *conn.BaseDB, stCfgs []*config.SubTaskConfig, fn func(stCfgs []*config.SubTaskConfig) int, workerName string) connAmountChecker {
+	return connAmountChecker{
+		toCheckDB:     toCheckDB,
+		stCfgs:        stCfgs,
+		getConfigConn: fn,
+		workerName:    workerName,
+	}
+}
+
+func (c *connAmountChecker) check(ctx context.Context, checkerName string) *Result {
+	result := &Result{
+		Name:  checkerName,
+		Desc:  "check if user-specified number of connection exceeds database's maximum",
+		State: StateFailure,
+	}
+	baseConn, err := c.toCheckDB.GetBaseConn(ctx)
+	//nolint:errcheck
+	defer c.toCheckDB.CloseBaseConn(baseConn)
+	if err != nil {
+		markCheckError(result, err)
+		return result
+	}
+	var rows, processRows *sql.Rows
+	rows, err = baseConn.QuerySQL(tcontext.NewContext(ctx, log.L()), "SHOW GLOBAL VARIABLES LIKE 'max_connections'")
+	if err != nil {
+		markCheckError(result, err)
+		return result
+	}
+	defer rows.Close()
+	var (
+		maxConn  int
+		variable string
+	)
+	for rows.Next() {
+		err = rows.Scan(&variable, &maxConn)
+		if err != nil {
+			markCheckError(result, err)
+			return result
+		}
+	}
+	if maxConn == 0 {
+		c.unlimitedConn = true
+	}
+	processRows, err = baseConn.QuerySQL(tcontext.NewContext(ctx, log.L()), "SHOW PROCESSLIST")
+	if err != nil {
+		markCheckError(result, err)
+		return result
+	}
+	defer processRows.Close()
+	usedConn := 0
+	for processRows.Next() {
+		usedConn++
+	}
+	usedConn -= 1 // exclude the connection used for show processlist
+	log.L().Debug("connection checker", zap.Int("max_connections", maxConn), zap.Int("used_connections", usedConn))
+	neededConn := c.getConfigConn(c.stCfgs)
+	switch {
+	case c.unlimitedConn:
+		// zero max_connections means unlimited connections
+		// we shouldn't report a warning.
+		c.unlimitedConn = true
+		result.State = StateSuccess
+	case neededConn > maxConn:
+		// nonzero max_connections and needed connections exceed max_connections
+		// FYI: https://github.com/pingcap/tidb/pull/35453
+		// currently, TiDB's max_connections is set to 0 representing unlimited connections,
+		// while for MySQL, 0 is not a legal value (never retrieve from it).
+		result.Errors = append(
+			result.Errors,
+			NewError(
+				"database's max_connections: %d is less than the number %s needs: %d",
+				maxConn,
+				c.workerName,
+				neededConn,
+			),
+		)
+		result.Instruction = "set larger max_connections or reduce the pool size of dm"
+		result.State = StateFailure
+	default:
+		// nonzero max_connections and needed connections are less than or equal to max_connections
+		result.State = StateSuccess
+		if maxConn-usedConn < neededConn {
+			result.State = StateWarning
+			result.Errors = append(
+				result.Errors,
+				NewError(
+					"database's max_connections: %d, used_connections: %d, available_connections: %d is less than %s needs: %d",
+					maxConn,
+					usedConn,
+					maxConn-usedConn,
+					c.workerName,
+					neededConn,
+				),
+			)
+		}
+	}
+	return result
+}
+
+type LoaderConnAmountChecker struct {
+	connAmountChecker
+}
+
+func NewLoaderConnAmountChecker(targetDB *conn.BaseDB, stCfgs []*config.SubTaskConfig) RealChecker {
+	return &LoaderConnAmountChecker{
+		connAmountChecker: newConnAmountChecker(targetDB, stCfgs, func(stCfgs []*config.SubTaskConfig) int {
+			loaderConn := 0
+			for _, stCfg := range stCfgs {
+				// loader's worker and checkpoint (always keeps one db connection)
+				loaderConn += stCfg.LoaderConfig.PoolSize + 1
+			}
+			return loaderConn
+		}, "loader"),
+	}
+}
+
+func (l *LoaderConnAmountChecker) Name() string {
+	return "loader_conn_amount_checker"
+}
+
+func (l *LoaderConnAmountChecker) Check(ctx context.Context) *Result {
+	result := l.check(ctx, l.Name())
+	if !l.unlimitedConn && result.State == StateFailure {
+		// if the max_connections is set as a specific number
+		// and we failed because of the number connecions needed is smaller than max_connections
+		for _, stCfg := range l.stCfgs {
+			if stCfg.NeedUseLightning() {
+				// if we're using lightning, this error should be omitted
+				// because lightning doesn't need to keep connections while restoring.
+				result.Errors = append(
+					result.Errors,
+					NewWarn("task precheck cannot accurately check the number of connection needed for Lightning, please set a sufficiently large connections for TiDB"),
+				)
+				result.State = StateWarning
+				break
+			}
+		}
+	}
+	return result
+}
+
+func NewDumperConnAmountChecker(sourceDB *conn.BaseDB, dumperThreads int) RealChecker {
+	return &DumperConnAmountChecker{
+		connAmountChecker: newConnAmountChecker(sourceDB, nil, func(_ []*config.SubTaskConfig) int {
+			// one for generating SQL, another for consistency control
+			return dumperThreads + 2
+		}, "dumper"),
+	}
+}
+
+type DumperConnAmountChecker struct {
+	connAmountChecker
+}
+
+func (d *DumperConnAmountChecker) Check(ctx context.Context) *Result {
+	return d.check(ctx, d.Name())
+}
+
+func (d *DumperConnAmountChecker) Name() string {
+	return "dumper_conn_amount_checker"
+}

dm/pkg/checker/conn_checker_test.go

@@ -0,0 +1,68 @@
+// Copyright 2022 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package checker
+
+import (
+	"testing"
+
+	"github.com/DATA-DOG/go-sqlmock"
+	"github.com/pingcap/tidb/dumpling/context"
+	"github.com/pingcap/tiflow/dm/config"
+	"github.com/pingcap/tiflow/dm/pkg/conn"
+	"github.com/stretchr/testify/require"
+)
+
+func TestConnAmountChecker(t *testing.T) {
+	var err error
+	db, dbMock, err := sqlmock.New()
+	require.NoError(t, err)
+	stCfgs := []*config.SubTaskConfig{
+		{
+			SyncerConfig: config.SyncerConfig{
+				WorkerCount: 10,
+			},
+			LoaderConfig: config.LoaderConfig{
+				PoolSize: 16,
+			},
+		},
+	}
+	baseDB := conn.NewBaseDB(db, func() {})
+	// test loader: fail
+	dbMock.ExpectQuery("SHOW GLOBAL VARIABLES LIKE 'max_connections'").WillReturnRows(sqlmock.NewRows([]string{"Variable_name", "Value"}).
+		AddRow("max_connections", 16))
+	dbMock.ExpectQuery("SHOW PROCESSLIST").WillReturnRows(sqlmock.NewRows(
+		[]string{"Id", "User", "Host", "db", "Command", "Time", "State", "Info"}).
+		AddRow(1, "root", "localhost", "test", "Query", 0, "init", ""),
+	)
+	loaderChecker := NewLoaderConnAmountChecker(baseDB, stCfgs)
+	result := loaderChecker.Check(context.Background())
+	require.Equal(t, 1, len(result.Errors))
+	require.Equal(t, StateFailure, result.State)
+	require.Regexp(t, "(.|\n)*is less than the number loader(.|\n)*", result.Errors[0].ShortErr)
+
+	// test loader: success
+	db, dbMock, err = sqlmock.New()
+	require.NoError(t, err)
+	baseDB = conn.NewBaseDB(db, func() {})
+	dbMock.ExpectQuery("SHOW GLOBAL VARIABLES LIKE 'max_connections'").WillReturnRows(sqlmock.NewRows([]string{"Variable_name", "Value"}).
+		AddRow("max_connections", 17))
+	dbMock.ExpectQuery("SHOW PROCESSLIST").WillReturnRows(sqlmock.NewRows(
+		[]string{"Id", "User", "Host", "db", "Command", "Time", "State", "Info"}).
+		AddRow(1, "root", "localhost", "test", "Query", 0, "init", ""),
+	)
+	loaderChecker = NewLoaderConnAmountChecker(baseDB, stCfgs)
+	result = loaderChecker.Check(context.Background())
+	require.Equal(t, 0, len(result.Errors))
+	require.Equal(t, StateSuccess, result.State)
+}

dm/tests/dmctl_command/conf/dm-task3.yaml

@@ -0,0 +1,37 @@
+---
+name: test
+task-mode: full
+meta-schema: "dm_meta"
+enable-heartbeat: false
+shard-mode: "pessimistic"
+target-database:
+  host: "127.0.0.1"
+  port: 4000
+  user: "root"
+  password: ""
+
+mysql-instances:
+  - source-id: "mysql-replica-01"
+    block-allow-list:  "instance"
+    mydumper-config-name: "global"
+    loader-config-name: "global"
+
+block-allow-list:
+  instance:
+    do-dbs: ["dmctl_conn"]
+    do-tables:
+    -  db-name: "dmctl_conn"
+       tbl-name: "*"
+
+mydumpers:
+  global:
+    threads: 2
+    chunk-filesize: 64
+    skip-tz-utc: true
+    extra-args: ""
+
+loaders:
+  global:
+    pool-size: 5
+    import-mode: "loader"
+    dir: "./dumped_data"