Repository: pingcap/dumpling
Pull request #256 (2021-03-27 09:36:53)

dump: fix the problem that dumpling gets stuck after the database connection is lost
<!--
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-->
fix https://github.com/pingcap/dumpling/issues/243

### What is changed and how it works?
Use `writerCtx` to avoid `taskChan` get blocked after the writer quitted. 

### Check List <!--REMOVE the items that are not applicable-->

Tests <!-- At least one of them must be included. -->

 - Unit test
 - Integration test

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
-->
- fix the problem that dumpling gets stuck after the database connection is lost

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/256/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/256/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-->
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/256/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/256/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: 5479520f0da97cfd0678f3ad94f14a574869b2e6</details>

Patch:
diff --git a/v4/export/config.go b/v4/export/config.go
index 2548dbe..00d8680 100644
--- a/v4/export/config.go
+++ b/v4/export/config.go
@@ -560,7 +560,7 @@ const (
 
 var (
 	gcSafePointVersion = semver.New("4.0.0")
-	tableSampleVersion = semver.New("5.0.0")
+	tableSampleVersion = semver.New("5.0.0-nightly")
 )
 
 // ServerInfo is the combination of ServerType and ServerInfo
diff --git a/v4/export/dump.go b/v4/export/dump.go
index eec79fd..e0b5f55 100755
--- a/v4/export/dump.go
+++ b/v4/export/dump.go
@@ -164,7 +164,8 @@ func (d *Dumper) Dump() (dumpErr error) {
 	taskChan := make(chan Task, defaultDumpThreads)
 	AddGauge(taskChannelCapacity, conf.Labels, defaultDumpThreads)
 	wg, writingCtx := errgroup.WithContext(tctx)
-	writers, tearDownWriters, err := d.startWriters(tctx.WithContext(writingCtx), wg, taskChan, rebuildConn)
+	writerCtx := tctx.WithContext(writingCtx)
+	writers, tearDownWriters, err := d.startWriters(writerCtx, wg, taskChan, rebuildConn)
 	if err != nil {
 		return err
 	}
@@ -217,11 +218,11 @@ func (d *Dumper) Dump() (dumpErr error) {
 	}
 
 	if conf.SQL == "" {
-		if err = d.dumpDatabases(metaConn, taskChan); err != nil {
+		if err = d.dumpDatabases(writerCtx, metaConn, taskChan); err != nil && !errors.ErrorEqual(err, context.Canceled) {
 			return err
 		}
 	} else {
-		d.dumpSQL(taskChan)
+		d.dumpSQL(writerCtx, taskChan)
 	}
 	close(taskChan)
 	if err := wg.Wait(); err != nil {
@@ -279,7 +280,7 @@ func (d *Dumper) startWriters(tctx *tcontext.Context, wg *errgroup.Group, taskCh
 	return writers, tearDown, nil
 }
 
-func (d *Dumper) dumpDatabases(metaConn *sql.Conn, taskChan chan<- Task) error {
+func (d *Dumper) dumpDatabases(tctx *tcontext.Context, metaConn *sql.Conn, taskChan chan<- Task) error {
 	conf := d.conf
 	allTables := conf.Tables
 	for dbName, tables := range allTables {
@@ -288,7 +289,10 @@ func (d *Dumper) dumpDatabases(metaConn *sql.Conn, taskChan chan<- Task) error {
 			return err
 		}
 		task := NewTaskDatabaseMeta(dbName, createDatabaseSQL)
-		d.sendTaskToChan(task, taskChan)
+		ctxDone := d.sendTaskToChan(tctx, task, taskChan)
+		if ctxDone {
+			return tctx.Err()
+		}
 
 		for _, table := range tables {
 			d.L().Debug("start dumping table...", zap.String("database", dbName),
@@ -300,11 +304,17 @@ func (d *Dumper) dumpDatabases(metaConn *sql.Conn, taskChan chan<- Task) error {
 
 			if table.Type == TableTypeView {
 				task := NewTaskViewMeta(dbName, table.Name, meta.ShowCreateTable(), meta.ShowCreateView())
-				d.sendTaskToChan(task, taskChan)
+				ctxDone = d.sendTaskToChan(tctx, task, taskChan)
+				if ctxDone {
+					return tctx.Err()
+				}
 			} else {
 				task := NewTaskTableMeta(dbName, table.Name, meta.ShowCreateTable())
-				d.sendTaskToChan(task, taskChan)
-				err = d.dumpTableData(metaConn, meta, taskChan)
+				ctxDone = d.sendTaskToChan(tctx, task, taskChan)
+				if ctxDone {
+					return tctx.Err()
+				}
+				err = d.dumpTableData(tctx, metaConn, meta, taskChan)
 				if err != nil {
 					return err
 				}
@@ -315,18 +325,18 @@ func (d *Dumper) dumpDatabases(metaConn *sql.Conn, taskChan chan<- Task) error {
 	return nil
 }
 
-func (d *Dumper) dumpTableData(conn *sql.Conn, meta TableMeta, taskChan chan<- Task) error {
+func (d *Dumper) dumpTableData(tctx *tcontext.Context, conn *sql.Conn, meta TableMeta, taskChan chan<- Task) error {
 	conf := d.conf
 	if conf.NoData {
 		return nil
 	}
 	if conf.Rows == UnspecifiedSize {
-		return d.sequentialDumpTable(conn, meta, taskChan)
+		return d.sequentialDumpTable(tctx, conn, meta, taskChan)
 	}
-	return d.concurrentDumpTable(conn, meta, taskChan)
+	return d.concurrentDumpTable(tctx, conn, meta, taskChan)
 }
 
-func (d *Dumper) sequentialDumpTable(conn *sql.Conn, meta TableMeta, taskChan chan<- Task) error {
+func (d *Dumper) sequentialDumpTable(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)
@@ -334,12 +344,15 @@ func (d *Dumper) sequentialDumpTable(conn *sql.Conn, meta TableMeta, taskChan ch
 		return err
 	}
 	task := NewTaskTableData(meta, tableIR, 0, 1)
-	d.sendTaskToChan(task, taskChan)
+	ctxDone := d.sendTaskToChan(tctx, task, taskChan)
+	if ctxDone {
+		return tctx.Err()
+	}
 
 	return nil
 }
 
-func (d *Dumper) concurrentDumpTable(conn *sql.Conn, meta TableMeta, taskChan chan<- Task) error {
+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 &&
@@ -347,7 +360,7 @@ func (d *Dumper) concurrentDumpTable(conn *sql.Conn, meta TableMeta, taskChan ch
 		conf.ServerInfo.ServerVersion.Compare(*tableSampleVersion) >= 0 {
 		d.L().Debug("dumping TiDB tables with TABLESAMPLE",
 			zap.String("database", db), zap.String("table", tbl))
-		return d.concurrentDumpTiDBTables(conn, meta, taskChan)
+		return d.concurrentDumpTiDBTables(tctx, conn, meta, taskChan)
 	}
 	field, err := pickupPossibleField(db, tbl, conn, conf)
 	if err != nil {
@@ -357,7 +370,7 @@ func (d *Dumper) concurrentDumpTable(conn *sql.Conn, meta TableMeta, taskChan ch
 		// skip split chunk logic if not found proper field
 		d.L().Warn("fallback to sequential dump due to no proper field",
 			zap.String("database", db), zap.String("table", tbl))
-		return d.sequentialDumpTable(conn, meta, taskChan)
+		return d.sequentialDumpTable(tctx, conn, meta, taskChan)
 	}
 
 	min, max, err := d.selectMinAndMaxIntValue(conn, db, tbl, field)
@@ -380,7 +393,7 @@ func (d *Dumper) concurrentDumpTable(conn *sql.Conn, meta TableMeta, taskChan ch
 			zap.Uint64("conf.rows", conf.Rows),
 			zap.String("database", db),
 			zap.String("table", tbl))
-		return d.sequentialDumpTable(conn, meta, taskChan)
+		return d.sequentialDumpTable(tctx, conn, meta, taskChan)
 	}
 
 	// every chunk would have eventual adjustments
@@ -416,9 +429,9 @@ func (d *Dumper) concurrentDumpTable(conn *sql.Conn, meta TableMeta, taskChan ch
 			nullValueCondition = ""
 		}
 		task := NewTaskTableData(meta, newTableData(query, selectLen, false), chunkIndex, int(totalChunks))
-		ctxDone := d.sendTaskToChan(task, taskChan)
+		ctxDone := d.sendTaskToChan(tctx, task, taskChan)
 		if ctxDone {
-			break
+			return tctx.Err()
 		}
 		cutoff = nextCutOff
 		chunkIndex++
@@ -426,8 +439,8 @@ func (d *Dumper) concurrentDumpTable(conn *sql.Conn, meta TableMeta, taskChan ch
 	return nil
 }
 
-func (d *Dumper) sendTaskToChan(task Task, taskChan chan<- Task) (ctxDone bool) {
-	tctx, conf := d.tctx, d.conf
+func (d *Dumper) sendTaskToChan(tctx *tcontext.Context, task Task, taskChan chan<- Task) (ctxDone bool) {
+	conf := d.conf
 	select {
 	case <-tctx.Done():
 		return true
@@ -474,7 +487,7 @@ func (d *Dumper) selectMinAndMaxIntValue(conn *sql.Conn, db, tbl, field string)
 	return min, max, nil
 }
 
-func (d *Dumper) concurrentDumpTiDBTables(conn *sql.Conn, meta TableMeta, taskChan chan<- Task) error {
+func (d *Dumper) concurrentDumpTiDBTables(tctx *tcontext.Context, conn *sql.Conn, meta TableMeta, taskChan chan<- Task) error {
 	conf := d.conf
 	db, tbl := meta.DatabaseName(), meta.TableName()
 
@@ -495,9 +508,9 @@ func (d *Dumper) concurrentDumpTiDBTables(conn *sql.Conn, meta TableMeta, taskCh
 	for i, w := range where {
 		query := buildSelectQuery(db, tbl, selectField, buildWhereCondition(conf, w), orderByClause)
 		task := NewTaskTableData(meta, newTableData(query, selectLen, false), i, len(where))
-		ctxDone := d.sendTaskToChan(task, taskChan)
+		ctxDone := d.sendTaskToChan(tctx, task, taskChan)
 		if ctxDone {
-			break
+			return tctx.Err()
 		}
 	}
 	return nil
@@ -624,12 +637,12 @@ func dumpTableMeta(conf *Config, conn *sql.Conn, db string, table *TableInfo) (T
 	return meta, nil
 }
 
-func (d *Dumper) dumpSQL(taskChan chan<- Task) {
+func (d *Dumper) dumpSQL(tctx *tcontext.Context, taskChan chan<- Task) {
 	conf := d.conf
 	meta := &tableMeta{}
 	data := newTableData(conf.SQL, 0, true)
 	task := NewTaskTableData(meta, data, 0, 1)
-	d.sendTaskToChan(task, taskChan)
+	d.sendTaskToChan(tctx, task, taskChan)
 }
 
 func canRebuildConn(consistency string, trxConsistencyOnly bool) bool {


Test patch:
diff --git a/v4/export/dump_test.go b/v4/export/dump_test.go
new file mode 100644
index 0000000..b9d8ec6
--- /dev/null
+++ b/v4/export/dump_test.go
@@ -0,0 +1,55 @@
+// Copyright 2021 PingCAP, Inc. Licensed under Apache-2.0.
+
+package export
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	tcontext "github.com/pingcap/dumpling/v4/context"
+
+	"github.com/DATA-DOG/go-sqlmock"
+	. "github.com/pingcap/check"
+	"github.com/pingcap/errors"
+	"golang.org/x/sync/errgroup"
+)
+
+func (s *testSQLSuite) TestDumpBlock(c *C) {
+	db, mock, err := sqlmock.New()
+	c.Assert(err, IsNil)
+	defer db.Close()
+
+	database := "test"
+	mock.ExpectQuery(fmt.Sprintf("SHOW CREATE DATABASE `%s`", escapeString(database))).
+		WillReturnRows(sqlmock.NewRows([]string{"Database", "Create Database"}).
+			AddRow("test", "CREATE DATABASE `test` /*!40100 DEFAULT CHARACTER SET utf8mb4 */"))
+
+	tctx, cancel := tcontext.Background().WithCancel()
+	defer cancel()
+	conn, err := db.Conn(tctx)
+	c.Assert(err, IsNil)
+
+	d := &Dumper{
+		tctx:      tctx,
+		conf:      DefaultConfig(),
+		cancelCtx: cancel,
+	}
+	wg, writingCtx := errgroup.WithContext(tctx)
+	writerErr := errors.New("writer error")
+
+	wg.Go(func() error {
+		return errors.Trace(writerErr)
+	})
+	wg.Go(func() error {
+		time.Sleep(time.Second)
+		return context.Canceled
+	})
+	writerCtx := tctx.WithContext(writingCtx)
+	// simulate taskChan is full
+	taskChan := make(chan Task, 1)
+	taskChan <- &TaskDatabaseMeta{}
+	d.conf.Tables = DatabaseTables{}.AppendTable("test", nil)
+	c.Assert(errors.ErrorEqual(d.dumpDatabases(writerCtx, conn, taskChan), context.Canceled), IsTrue)
+	c.Assert(errors.ErrorEqual(wg.Wait(), writerErr), IsTrue)
+}