Repository: pingcap/dumpling
Pull request #257 (2021-03-28 08:30:15)

dump: add split big tidb query to several small queries
<!--
Thank you for contributing to Dumpling! Please read the [CONTRIBUTING](https://github.com/pingcap/dumpling/blob/master/CONTRIBUTING.md) document **BEFORE** filing this PR.

If you want to open the **Challenge Program** pull request, please use the following template:
https://raw.githubusercontent.com/pingcap/.github/master/.github/PULL_REQUEST_TEMPLATE/challenge-program.md

You can use it with query parameters: https://github.com/pingcap/dumpling/compare/master...${you branch}?template=challenge-program.md
-->

### What problem does this PR solve? <!--add issue link with summary if exists-->
Although we can make sure TiDB won't OOM through the help of TiDB tablesample. But if we didn't specify `--rows` to open tablesample, TiDB will still OOM.

### What is changed and how it works?
implement solution 2 in https://github.com/pingcap/dumpling/issues/232
Try to split the whole dump table SQL to several small queries to concatenate their results.

### Check List <!--REMOVE the items that are not applicable-->

Tests <!-- At least one of them must be included. -->

 - Unit test
 - Integration test
 - Manual test:
 Test dumping 1.2TB data serially (with clustered index enabled), the data rows lines can match.
```log
[2021/03/28 22:31:00.483 +00:00] [INFO] [collector.go:212] ["backup Success summary: total backup ranges: 3, total success: 3, total failed: 0, total take(backup time): 7h12m34.659704485s, total take(real time): 7h12m34.659726417s, total size(MB): 1178105.36, avg speed(MB/s): 45.39, total rows: 2172298300"]
```
```
mysql> select count(*) from usertable;
+------------+
| count(*)   |
+------------+
| 2172298300 |
+------------+
1 row in set (11 min 59.50 sec)
```

Side effects

 - Increased code complexity

Related changes

 - Need to cherry-pick to the release branch
 - Need to update the documentation
 
### Release note

<!-- bugfixes or new feature need a release note, must in the form of a list, such as

- support -T/--tables-list argument

or if no need to be included in the release note, just add the following line

- No release note
-->
- add split big tidb query to several small queries to avoid TiDB OOM

ti-chi-bot: [REVIEW NOTIFICATION]

This pull request has not been approved.


To complete the [pull request process](https://book.prow.tidb.io/#/en/workflows/pr), please ask the reviewers in the [list](https://prow.tidb.io/tichi/repos/pingcap/dumpling/pulls/257/owners) to review by filling `/cc @reviewer` in the comment.
After your PR has acquired the required number of LGTMs, you can assign this pull request to the committer in the [list](https://prow.tidb.io/tichi/repos/pingcap/dumpling/pulls/257/owners) by filling  `/assign @committer` in the comment to help you merge this pull request.

The full list of commands accepted by this bot can be found [here](https://prow.tidb.io/command-help?repo=pingcap%2Fdumpling).

<details>

Reviewer can indicate their review by writing `/lgtm` in a comment.
Reviewer can cancel approval by writing `/lgtm cancel` in a comment.
</details>

<!--Review Notification Identifier-->
ti-chi-bot: @lichunzhu: PR needs rebase.

<details>

Instructions for interacting with me using PR comments are available [here](https://git.k8s.io/community/contributors/guide/pull-requests.md).  If you have questions or suggestions related to my behavior, please file an issue against the [kubernetes/test-infra](https://github.com/kubernetes/test-infra/issues/new?title=Prow%20issue:) repository.
</details>
lichunzhu: /merge
ti-chi-bot: @lichunzhu: `/merge` in this pull request requires 1 `/lgtm`.

<details>

In response to [this](https://github.com/pingcap/dumpling/pull/257#issuecomment-809253384):

>/merge


Instructions for interacting with me using PR comments are available [here](https://prow.tidb.io/command-help).  If you have questions or suggestions related to my behavior, please file an issue against the [ti-community-infra/tichi](https://github.com/ti-community-infra/tichi/issues/new?title=Prow%20issue:) repository.
</details>
lichunzhu: ### Dump huge table test

Dump 1150.50GB clustered index data without -r argument
dumpling:v5.0.0-rc,mydumper will OOM.
dumpling:70a709a cost 7h12m34s, avg speed: 45.39 MB/s

### Comparison test result

1844.61MB data from a TiDB:v5.0.0-nightly, cluster indexed table. dump the data without -r argument
| binary | cost time | speed |
| ------ | --------- | ------  |
| dumpling:70a709a | 39.00s | 47.29 MB/s
| dumpling:v5.0.0-rc | 34.49s | 53.49 MB/s
| mydumper | 33.28s | 55.42 MB/s

dumpling:70a709a has some speed regression. Mainly because it opens and closes more mysql connections than other binaries.


ti-chi-bot: @lichunzhu: PR needs rebase.

<details>

Instructions for interacting with me using PR comments are available [here](https://git.k8s.io/community/contributors/guide/pull-requests.md).  If you have questions or suggestions related to my behavior, please file an issue against the [kubernetes/test-infra](https://github.com/kubernetes/test-infra/issues/new?title=Prow%20issue:) repository.
</details>
lance6716: /lgtm
ti-chi-bot: [REVIEW NOTIFICATION]

This pull request has been approved by:

- lance6716




To complete the [pull request process](https://book.prow.tidb.io/#/en/workflows/pr), please ask the reviewers in the [list](https://prow.tidb.io/tichi/repos/pingcap/dumpling/pulls/257/owners) to review by filling `/cc @reviewer` in the comment.
After your PR has acquired the required number of LGTMs, you can assign this pull request to the committer in the [list](https://prow.tidb.io/tichi/repos/pingcap/dumpling/pulls/257/owners) by filling  `/assign @committer` in the comment to help you merge this pull request.

The full list of commands accepted by this bot can be found [here](https://prow.tidb.io/command-help?repo=pingcap%2Fdumpling).

<details>

Reviewer can indicate their review by writing `/lgtm` in a comment.
Reviewer can cancel approval by writing `/lgtm cancel` in a comment.
</details>

<!--Review Notification Identifier-->
lichunzhu: /merge
ti-chi-bot: This pull request has been accepted and is ready to merge. <details>Commit hash: 4f22704fdcb20a6ee0bca87ea04d1b3914d24600</details>
lichunzhu: /cherrypick release-5.0


ti-chi-bot: @lichunzhu: once the present PR merges, I will cherry-pick it on top of release-5.0 in a new PR and assign it to you.

<details>

In response to [this](https://github.com/pingcap/dumpling/pull/257#issuecomment-811873014):

>/cherrypick release-5.0
>
>


Instructions for interacting with me using PR comments are available [here](https://git.k8s.io/community/contributors/guide/pull-requests.md).  If you have questions or suggestions related to my behavior, please file an issue against the [kubernetes/test-infra](https://github.com/kubernetes/test-infra/issues/new?title=Prow%20issue:) repository.
</details>

Patch:
diff --git a/v4/export/dump.go b/v4/export/dump.go
index 81f7b19..6eb5670 100755
--- a/v4/export/dump.go
+++ b/v4/export/dump.go
@@ -79,6 +79,7 @@ func (d *Dumper) Dump() (dumpErr error) {
 		conCtrl ConsistencyController
 	)
 	tctx, conf, pool := d.tctx, d.conf, d.dbHandle
+	tctx.L().Info("begin to run Dump", zap.Stringer("conf", conf))
 	m := newGlobalMetadata(tctx, d.extStore, conf.Snapshot)
 	defer func() {
 		if dumpErr == nil {
@@ -295,7 +296,7 @@ func (d *Dumper) dumpDatabases(tctx *tcontext.Context, metaConn *sql.Conn, taskC
 		}
 
 		for _, table := range tables {
-			d.L().Debug("start dumping table...", zap.String("database", dbName),
+			tctx.L().Debug("start dumping table...", zap.String("database", dbName),
 				zap.String("table", table.Name))
 			meta, err := dumpTableMeta(conf, metaConn, dbName, table)
 			if err != nil {
@@ -336,10 +337,71 @@ func (d *Dumper) dumpTableData(tctx *tcontext.Context, conn *sql.Conn, meta Tabl
 	return d.concurrentDumpTable(tctx, conn, meta, taskChan)
 }
 
-func (d *Dumper) sequentialDumpTable(tctx *tcontext.Context, conn *sql.Conn, meta TableMeta, taskChan chan<- Task) error {
+func (d *Dumper) buildConcatTask(tctx *tcontext.Context, conn *sql.Conn, meta TableMeta) (*TaskTableData, error) {
+	tableChan := make(chan Task, 128)
+	errCh := make(chan error, 1)
+	go func() {
+		// adjust rows to suitable rows for this table
+		d.conf.Rows = GetSuitableRows(tctx, conn, meta.DatabaseName(), meta.TableName())
+		err := d.concurrentDumpTable(tctx, conn, meta, tableChan)
+		d.conf.Rows = UnspecifiedSize
+		if err != nil {
+			errCh <- err
+		} else {
+			close(errCh)
+		}
+	}()
+	tableDataArr := make([]*tableData, 0)
+	handleSubTask := func(task Task) {
+		tableTask, ok := task.(*TaskTableData)
+		if !ok {
+			tctx.L().Warn("unexpected task when splitting table chunks", zap.String("task", tableTask.Brief()))
+			return
+		}
+		tableDataInst, ok := tableTask.Data.(*tableData)
+		if !ok {
+			tctx.L().Warn("unexpected task.Data when splitting table chunks", zap.String("task", tableTask.Brief()))
+			return
+		}
+		tableDataArr = append(tableDataArr, tableDataInst)
+	}
+	for {
+		select {
+		case err, ok := <-errCh:
+			if !ok {
+				// make sure all the subtasks in tableChan are handled
+				for len(tableChan) > 0 {
+					task := <-tableChan
+					handleSubTask(task)
+				}
+				if len(tableDataArr) <= 1 {
+					return nil, nil
+				}
+				queries := make([]string, 0, len(tableDataArr))
+				colLen := tableDataArr[0].colLen
+				for _, tableDataInst := range tableDataArr {
+					queries = append(queries, tableDataInst.query)
+					if colLen != tableDataInst.colLen {
+						tctx.L().Warn("colLen varies for same table",
+							zap.Int("oldColLen", colLen),
+							zap.String("oldQuery", queries[0]),
+							zap.Int("newColLen", tableDataInst.colLen),
+							zap.String("newQuery", tableDataInst.query))
+						return nil, nil
+					}
+				}
+				return NewTaskTableData(meta, newMultiQueriesChunk(queries, colLen), 0, 1), nil
+			}
+			return nil, err
+		case task := <-tableChan:
+			handleSubTask(task)
+		}
+	}
+}
+
+func (d *Dumper) dumpWholeTableDirectly(tctx *tcontext.Context, conn *sql.Conn, meta TableMeta, taskChan chan<- Task) error {
 	conf := d.conf
-	db, tbl := meta.DatabaseName(), meta.TableName()
-	tableIR, err := SelectAllFromTable(conf, conn, db, tbl)
+	tableIR, err := SelectAllFromTable(conf, conn, meta)
 	if err != nil {
 		return err
 	}
@@ -348,52 +410,73 @@ func (d *Dumper) sequentialDumpTable(tctx *tcontext.Context, conn *sql.Conn, met
 	if ctxDone {
 		return tctx.Err()
 	}
-
 	return nil
 }
 
+func (d *Dumper) sequentialDumpTable(tctx *tcontext.Context, conn *sql.Conn, meta TableMeta, taskChan chan<- Task) error {
+	conf := d.conf
+	if conf.ServerInfo.ServerType == ServerTypeTiDB {
+		task, err := d.buildConcatTask(tctx, conn, meta)
+		if err != nil {
+			return errors.Trace(err)
+		}
+		if task != nil {
+			ctxDone := d.sendTaskToChan(tctx, task, taskChan)
+			if ctxDone {
+				return tctx.Err()
+			}
+			return nil
+		}
+		tctx.L().Info("didn't build tidb concat sqls, will select all from table now",
+			zap.String("database", meta.DatabaseName()),
+			zap.String("table", meta.TableName()))
+	}
+	return d.dumpWholeTableDirectly(tctx, conn, meta, taskChan)
+}
+
+// concurrentDumpTable tries to split table into several chunks to dump
 func (d *Dumper) concurrentDumpTable(tctx *tcontext.Context, conn *sql.Conn, meta TableMeta, taskChan chan<- Task) error {
 	conf := d.conf
 	db, tbl := meta.DatabaseName(), meta.TableName()
 	if conf.ServerInfo.ServerType == ServerTypeTiDB &&
 		conf.ServerInfo.ServerVersion != nil &&
 		conf.ServerInfo.ServerVersion.Compare(*tableSampleVersion) >= 0 {
-		d.L().Debug("dumping TiDB tables with TABLESAMPLE",
+		tctx.L().Debug("dumping TiDB tables with TABLESAMPLE",
 			zap.String("database", db), zap.String("table", tbl))
 		return d.concurrentDumpTiDBTables(tctx, conn, meta, taskChan)
 	}
 	field, err := pickupPossibleField(db, tbl, conn, conf)
 	if err != nil {
-		return nil
+		return err
 	}
 	if field == "" {
 		// skip split chunk logic if not found proper field
-		d.L().Warn("fallback to sequential dump due to no proper field",
+		tctx.L().Warn("fallback to sequential dump due to no proper field",
 			zap.String("database", db), zap.String("table", tbl))
-		return d.sequentialDumpTable(tctx, conn, meta, taskChan)
+		return d.dumpWholeTableDirectly(tctx, conn, meta, taskChan)
 	}
 
 	min, max, err := d.selectMinAndMaxIntValue(conn, db, tbl, field)
 	if err != nil {
 		return err
 	}
-	d.L().Debug("get int bounding values",
+	tctx.L().Debug("get int bounding values",
 		zap.String("lower", min.String()),
 		zap.String("upper", max.String()))
 
 	count := estimateCount(d.tctx, db, tbl, conn, field, conf)
-	d.L().Info("get estimated rows count",
+	tctx.L().Info("get estimated rows count",
 		zap.String("database", db),
 		zap.String("table", tbl),
 		zap.Uint64("estimateCount", count))
 	if count < conf.Rows {
 		// skip chunk logic if estimates are low
-		d.L().Warn("skip concurrent dump due to estimate count < rows",
+		tctx.L().Warn("skip concurrent dump due to estimate count < rows",
 			zap.Uint64("estimate count", count),
 			zap.Uint64("conf.rows", conf.Rows),
 			zap.String("database", db),
 			zap.String("table", tbl))
-		return d.sequentialDumpTable(tctx, conn, meta, taskChan)
+		return d.dumpWholeTableDirectly(tctx, conn, meta, taskChan)
 	}
 
 	// every chunk would have eventual adjustments
@@ -496,7 +579,7 @@ func (d *Dumper) concurrentDumpTiDBTables(tctx *tcontext.Context, conn *sql.Conn
 		return err
 	}
 	if len(handleVals) == 0 {
-		return nil
+		return d.dumpWholeTableDirectly(tctx, conn, meta, taskChan)
 	}
 	selectField, selectLen, err := buildSelectField(conn, db, tbl, conf.CompleteInsert)
 	if err != nil {
diff --git a/v4/export/ir_impl.go b/v4/export/ir_impl.go
index c9c216f..ee5e01d 100644
--- a/v4/export/ir_impl.go
+++ b/v4/export/ir_impl.go
@@ -51,6 +51,117 @@ func (iter *rowIter) HasNext() bool {
 	return iter.hasNext
 }
 
+// multiQueriesChunkIter implements the SQLRowIter interface.
+// Note: To create a rowIter, please use `newRowIter()` instead of struct literal.
+type multiQueriesChunkIter struct {
+	tctx    *tcontext.Context
+	conn    *sql.Conn
+	rows    *sql.Rows
+	hasNext bool
+	id      int
+	queries []string
+	args    []interface{}
+	err     error
+}
+
+func newMultiQueryChunkIter(tctx *tcontext.Context, conn *sql.Conn, queries []string, argLen int) *multiQueriesChunkIter {
+	r := &multiQueriesChunkIter{
+		tctx:    tctx,
+		conn:    conn,
+		queries: queries,
+		id:      0,
+		args:    make([]interface{}, argLen),
+	}
+	r.nextRows()
+	return r
+}
+
+func (iter *multiQueriesChunkIter) nextRows() {
+	if iter.id >= len(iter.queries) {
+		iter.hasNext = false
+		return
+	}
+	var err error
+	defer func() {
+		if err != nil {
+			iter.hasNext = false
+			iter.err = errors.Trace(err)
+		}
+	}()
+	tctx, conn := iter.tctx, iter.conn
+	// avoid the empty chunk
+	for iter.id < len(iter.queries) {
+		rows := iter.rows
+		if rows != nil {
+			err = rows.Close()
+			if err != nil {
+				return
+			}
+			err = rows.Err()
+			if err != nil {
+				return
+			}
+		}
+		tctx.L().Debug("try to start nextRows", zap.String("query", iter.queries[iter.id]))
+		rows, err = conn.QueryContext(tctx, iter.queries[iter.id])
+		if err != nil {
+			return
+		}
+		if err = rows.Err(); err != nil {
+			return
+		}
+		iter.id++
+		iter.rows = rows
+		iter.hasNext = iter.rows.Next()
+		if iter.hasNext {
+			return
+		}
+	}
+}
+
+func (iter *multiQueriesChunkIter) Close() error {
+	if iter.err != nil {
+		return iter.err
+	}
+	if iter.rows != nil {
+		return iter.rows.Close()
+	}
+	return nil
+}
+
+func (iter *multiQueriesChunkIter) Decode(row RowReceiver) error {
+	if iter.err != nil {
+		return iter.err
+	}
+	if iter.rows == nil {
+		return errors.Errorf("no valid rows found, id: %d", iter.id)
+	}
+	return decodeFromRows(iter.rows, iter.args, row)
+}
+
+func (iter *multiQueriesChunkIter) Error() error {
+	if iter.err != nil {
+		return iter.err
+	}
+	if iter.rows != nil {
+		return errors.Trace(iter.rows.Err())
+	}
+	return nil
+}
+
+func (iter *multiQueriesChunkIter) Next() {
+	if iter.err == nil {
+		iter.hasNext = iter.rows.Next()
+		if !iter.hasNext {
+			iter.nextRows()
+		}
+	}
+}
+
+func (iter *multiQueriesChunkIter) HasNext() bool {
+	return iter.hasNext
+}
+
 type stringIter struct {
 	idx int
 	ss  []string
@@ -131,7 +242,7 @@ func (td *tableData) Rows() SQLRowIter {
 }
 
 func (td *tableData) Close() error {
-	return td.rows.Close()
+	return td.SQLRowIter.Close()
 }
 
 func (td *tableData) RawRows() *sql.Rows {
@@ -215,3 +326,39 @@ func (m *metaData) MetaSQL() string {
 	}
 	return m.metaSQL
 }
+
+type multiQueriesChunk struct {
+	tctx    *tcontext.Context
+	conn    *sql.Conn
+	queries []string
+	colLen  int
+	SQLRowIter
+}
+
+func newMultiQueriesChunk(queries []string, colLength int) *multiQueriesChunk {
+	return &multiQueriesChunk{
+		queries: queries,
+		colLen:  colLength,
+	}
+}
+
+func (td *multiQueriesChunk) Start(tctx *tcontext.Context, conn *sql.Conn) error {
+	td.tctx = tctx
+	td.conn = conn
+	return nil
+}
+
+func (td *multiQueriesChunk) Rows() SQLRowIter {
+	if td.SQLRowIter == nil {
+		td.SQLRowIter = newMultiQueryChunkIter(td.tctx, td.conn, td.queries, td.colLen)
+	}
+	return td.SQLRowIter
+}
+
+func (td *multiQueriesChunk) Close() error {
+	return td.SQLRowIter.Close()
+}
+
+func (td *multiQueriesChunk) RawRows() *sql.Rows {
+	return nil
+}
diff --git a/v4/export/sql.go b/v4/export/sql.go
index 5374830..78c8652 100644
--- a/v4/export/sql.go
+++ b/v4/export/sql.go
@@ -186,7 +186,8 @@ func SelectVersion(db *sql.DB) (string, error) {
 }
 
 // SelectAllFromTable dumps data serialized from a specified table
-func SelectAllFromTable(conf *Config, db *sql.Conn, database, table string) (TableDataIR, error) {
+func SelectAllFromTable(conf *Config, db *sql.Conn, meta TableMeta) (TableDataIR, error) {
+	database, table := meta.DatabaseName(), meta.TableName()
 	selectedField, selectLen, err := buildSelectField(db, database, table, conf.CompleteInsert)
 	if err != nil {
 		return nil, err
@@ -267,6 +268,37 @@ func SelectTiDBRowID(db *sql.Conn, database, table string) (bool, error) {
 	return true, nil
 }
 
+// GetSuitableRows gets suitable rows for each table
+func GetSuitableRows(tctx *tcontext.Context, db *sql.Conn, database, table string) uint64 {
+	const (
+		defaultRows  = 200000
+		maxRows      = 1000000
+		bytesPerFile = 128 * 1024 * 1024 // 128MB per file by default
+	)
+	avgRowLength, err := GetAVGRowLength(tctx, db, database, table)
+	if err != nil || avgRowLength == 0 {
+		tctx.L().Debug("fail to get average row length", zap.Uint64("averageRowLength", avgRowLength), zap.Error(err))
+		return defaultRows
+	}
+	estimateRows := bytesPerFile / avgRowLength
+	if estimateRows > maxRows {
+		return maxRows
+	}
+	return estimateRows
+}
+
+// GetAVGRowLength gets whether this table's average row length
+func GetAVGRowLength(tctx *tcontext.Context, db *sql.Conn, database, table string) (uint64, error) {
+	const query = "select AVG_ROW_LENGTH from INFORMATION_SCHEMA.TABLES where table_schema=? and table_name=?;"
+	var avgRowLength uint64
+	row := db.QueryRowContext(tctx, query, database, table)
+	err := row.Scan(&avgRowLength)
+	if err != nil {
+		return 0, errors.Annotatef(err, "sql: %s", query)
+	}
+	return avgRowLength, nil
+}
+
 // GetColumnTypes gets *sql.ColumnTypes from a specified table
 func GetColumnTypes(db *sql.Conn, fields, database, table string) ([]*sql.ColumnType, error) {
 	query := fmt.Sprintf("SELECT %s FROM `%s`.`%s` LIMIT 1", fields, escapeString(database), escapeString(table))


Test patch:
diff --git a/v4/export/sql_test.go b/v4/export/sql_test.go
index 0402d42..bdecade 100644
--- a/v4/export/sql_test.go
+++ b/v4/export/sql_test.go
@@ -4,6 +4,7 @@ package export
 
 import (
 	"context"
+	"database/sql"
 	"errors"
 
 	tcontext "github.com/pingcap/dumpling/v4/context"
@@ -347,6 +348,63 @@ func (s *testSQLSuite) TestShowCreateView(c *C) {
 	c.Assert(mock.ExpectationsWereMet(), IsNil)
 }
 
+func (s *testSQLSuite) TestGetSuitableRows(c *C) {
+	db, mock, err := sqlmock.New()
+	c.Assert(err, IsNil)
+	defer db.Close()
+	conn, err := db.Conn(context.Background())
+	c.Assert(err, IsNil)
+	const query = "select AVG_ROW_LENGTH from INFORMATION_SCHEMA.TABLES where table_schema=\\? and table_name=\\?;"
+	tctx, cancel := tcontext.Background().WithCancel()
+	defer cancel()
+	database := "foo"
+	table := "bar"
+
+	testCases := []struct {
+		avgRowLength uint64
+		expectedRows uint64
+		returnErr    error
+	}{
+		{
+			32,
+			200000,
+			sql.ErrNoRows,
+		},
+		{
+			0,
+			200000,
+			nil,
+		},
+		{
+			32,
+			1000000,
+			nil,
+		},
+		{
+			1024,
+			131072,
+			nil,
+		},
+		{
+			4096,
+			32768,
+			nil,
+		},
+	}
+	for _, testCase := range testCases {
+		if testCase.returnErr == nil {
+			mock.ExpectQuery(query).WithArgs(database, table).
+				WillReturnRows(sqlmock.NewRows([]string{"AVG_ROW_LENGTH"}).
+					AddRow(testCase.avgRowLength))
+		} else {
+			mock.ExpectQuery(query).WithArgs(database, table).
+				WillReturnError(testCase.returnErr)
+		}
+		rows := GetSuitableRows(tctx, conn, database, table)
+		c.Assert(rows, Equals, testCase.expectedRows)
+	}
+}
+
 func (s *testSQLSuite) TestBuildWhereClauses(c *C) {
 	testCases := []struct {
 		handleColNames       []string