Repository: pingcap/dumpling
Pull request #264 (2021-04-01 12:38:16)

[release-5.0] dump: add split big tidb query to several small queries
close https://github.com/pingcap/dumpling/issues/263

This is a manual cherry-pick of #257

/assign lichunzhu

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/264/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/264/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: @lance6716 PTAL
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/264/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/264/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: e0873ed66a71eb7e0ac9eb044adb7e0f65e118e9</details>

Patch:
diff --git a/v4/export/dump.go b/v4/export/dump.go
index 5b11915..5e77ede 100755
--- a/v4/export/dump.go
+++ b/v4/export/dump.go
@@ -290,7 +290,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 {
@@ -331,10 +331,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
 	}
@@ -343,52 +404,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
@@ -491,7 +573,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 3c6ea35..c9c7c11 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