Repository: pingcap/dumpling
Pull request #258 (2021-03-29 09:04:59)

[release-5.0] dump: fix the problem that dumpling gets stuck after the database connection is lost
This is an automated cherry-pick of #256

/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/258/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/258/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: /lgtm
ti-chi-bot: [REVIEW NOTIFICATION]

This pull request has been approved by:

- lichunzhu




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/258/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/258/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: eea1bf420e56f9b000bffd26c4d0c7490d218f5a</details>
lichunzhu: /merge
zhouqiang-cl: /merge
ti-chi-bot: @zhouqiang-cl: `/merge` is only allowed for the committers in [list](https://prow.tidb.io/tichi/repos/pingcap/dumpling/pulls/258/owners).

<details>

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

>/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: /merge

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 d0e26ed..cbce362 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
 	}
@@ -211,11 +212,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 {
@@ -273,7 +274,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 {
@@ -282,7 +283,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),
@@ -294,11 +298,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
 				}
@@ -309,18 +319,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)
@@ -328,12 +338,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 &&
@@ -341,7 +354,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 {
@@ -351,7 +364,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)
@@ -374,7 +387,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
@@ -410,9 +423,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++
@@ -420,8 +433,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
@@ -468,7 +481,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()
 
@@ -489,9 +502,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
@@ -618,12 +631,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)
+}