Repository: pingcap/tidb-lightning
Pull request #362 (2020-07-31 06:27:25)

restore: fix missing colum infos when restore from checkpoint
<!--
Thank you for contributing to TiDB-Lightning! Please read the [CONTRIBUTING](https://github.com/pingcap/tidb-lightning/blob/master/CONTRIBUTING.md) document **BEFORE** filing this PR.
-->

### What problem does this PR solve? <!--add issue link with summary if exists-->
Fix missing column infos when restore from checkpoint

### What is changed and how it works?
- Write Chunk ColumnPermutation to checkpoint db (file/mysql)  when insert or update chunk checkpoints.
- When restore from checkpoint, also restore column infos for parser from stored  `ColumnPermutation`

NOTE: this fix can't handle following case:

table schema:
```
create table t(a varchar(32), b int);
```

data.sql
```
insert into t (a, b) values 
("test1", 1),
("test2", 2);

insert into t (a) values 
("test3"),
("test4");
```

If multi insert statements contains different columns number for their order are different, Mable recover from checkpoint will still fail because we don't update columns infos on the fly. 

### Check List <!--REMOVE the items that are not applicable-->

Tests <!-- At least one of them must be included. -->

 - Unit test
 - Integration test
 - Manual test (add detailed scripts or steps below)
 - No code

Side effects

Related changes


glorv: @kennytm PTAL
kennytm: PTAL @3pointer

Patch:
diff --git a/lightning/checkpoints/checkpoints.go b/lightning/checkpoints/checkpoints.go
index d88f2bb..d78ea7c 100644
--- a/lightning/checkpoints/checkpoints.go
+++ b/lightning/checkpoints/checkpoints.go
@@ -177,9 +177,10 @@ func (cp *TableCheckpoint) CountChunks() int {
 }
 
 type chunkCheckpointDiff struct {
-	pos      int64
-	rowID    int64
-	checksum verify.KVChecksum
+	pos               int64
+	rowID             int64
+	checksum          verify.KVChecksum
+	columnPermutation []int
 }
 
 type engineCheckpointDiff struct {
@@ -290,20 +291,22 @@ func (merger *StatusCheckpointMerger) MergeInto(cpd *TableCheckpointDiff) {
 }
 
 type ChunkCheckpointMerger struct {
-	EngineID int32
-	Key      ChunkCheckpointKey
-	Checksum verify.KVChecksum
-	Pos      int64
-	RowID    int64
+	EngineID          int32
+	Key               ChunkCheckpointKey
+	Checksum          verify.KVChecksum
+	Pos               int64
+	RowID             int64
+	ColumnPermutation []int
 }
 
 func (merger *ChunkCheckpointMerger) MergeInto(cpd *TableCheckpointDiff) {
 	cpd.insertEngineCheckpointDiff(merger.EngineID, engineCheckpointDiff{
 		chunks: map[ChunkCheckpointKey]chunkCheckpointDiff{
 			merger.Key: {
-				pos:      merger.Pos,
-				rowID:    merger.RowID,
-				checksum: merger.Checksum,
+				pos:               merger.Pos,
+				rowID:             merger.RowID,
+				checksum:          merger.Checksum,
+				columnPermutation: merger.ColumnPermutation,
 			},
 		},
 	})
@@ -624,7 +627,7 @@ func (cpdb *MySQLCheckpointsDB) InsertEngineCheckpoints(ctx context.Context, tab
 				kvc_bytes, kvc_kvs, kvc_checksum, create_time
 			) VALUES (
 				?, ?,
-				?, ?, '[]', FALSE,
+				?, ?, ?, FALSE,
 				?, ?, ?, ?,
 				0, 0, 0, from_unixtime(?)
 			);
@@ -640,9 +643,13 @@ func (cpdb *MySQLCheckpointsDB) InsertEngineCheckpoints(ctx context.Context, tab
 				return errors.Trace(err)
 			}
 			for _, value := range engine.Chunks {
+				columnPerm, err := json.Marshal(value.ColumnPermutation)
+				if err != nil {
+					return errors.Trace(err)
+				}
 				_, err = chunkStmt.ExecContext(
 					c, tableName, engineID,
-					value.Key.Path, value.Key.Offset,
+					value.Key.Path, value.Key.Offset, columnPerm,
 					value.Chunk.Offset, value.Chunk.EndOffset, value.Chunk.PrevRowIDMax, value.Chunk.RowIDMax,
 					value.Timestamp,
 				)
@@ -663,7 +670,7 @@ func (cpdb *MySQLCheckpointsDB) InsertEngineCheckpoints(ctx context.Context, tab
 
 func (cpdb *MySQLCheckpointsDB) Update(checkpointDiffs map[string]*TableCheckpointDiff) {
 	chunkQuery := fmt.Sprintf(`
-		UPDATE %s.%s SET pos = ?, prev_rowid_max = ?, kvc_bytes = ?, kvc_kvs = ?, kvc_checksum = ?
+		UPDATE %s.%s SET pos = ?, prev_rowid_max = ?, kvc_bytes = ?, kvc_kvs = ?, kvc_checksum = ?, columns = ?
 		WHERE (table_name, engine_id, path, offset) = (?, ?, ?, ?);
 	`, cpdb.schema, CheckpointTableNameChunk)
 	rebaseQuery := fmt.Sprintf(`
@@ -717,10 +724,14 @@ func (cpdb *MySQLCheckpointsDB) Update(checkpointDiffs map[string]*TableCheckpoi
 					}
 				}
 				for key, diff := range engineDiff.chunks {
+					columnPerm, err := json.Marshal(diff.columnPermutation)
+					if err != nil {
+						return errors.Trace(err)
+					}
 					if _, e := chunkStmt.ExecContext(
 						c,
 						diff.pos, diff.rowID, diff.checksum.SumSize(), diff.checksum.SumKVS(), diff.checksum.Sum(),
-						tableName, engineID, key.Path, key.Offset,
+						columnPerm, tableName, engineID, key.Path, key.Offset,
 					); e != nil {
 						return errors.Trace(e)
 					}
@@ -901,6 +912,10 @@ func (cpdb *FileCheckpointsDB) InsertEngineCheckpoints(_ context.Context, tableN
 			chunk.PrevRowidMax = value.Chunk.PrevRowIDMax
 			chunk.RowidMax = value.Chunk.RowIDMax
 			chunk.Timestamp = value.Timestamp
+			if len(value.ColumnPermutation) > 0 {
+				chunk.ColumnPermutation = intSlice2Int32Slice(value.ColumnPermutation)
+			}
+
 		}
 		tableModel.Engines[engineID] = engineModel
 	}
@@ -933,6 +948,7 @@ func (cpdb *FileCheckpointsDB) Update(checkpointDiffs map[string]*TableCheckpoin
 				chunkModel.KvcBytes = diff.checksum.SumSize()
 				chunkModel.KvcKvs = diff.checksum.SumKVS()
 				chunkModel.KvcChecksum = diff.checksum.Sum()
+				chunkModel.ColumnPermutation = intSlice2Int32Slice(diff.columnPermutation)
 			}
 		}
 	}
@@ -1293,3 +1309,11 @@ func (cpdb *FileCheckpointsDB) DumpEngines(context.Context, io.Writer) error {
 func (cpdb *FileCheckpointsDB) DumpChunks(context.Context, io.Writer) error {
 	return errors.Errorf("dumping file checkpoint into CSV not unsupported, you may copy %s instead", cpdb.path)
 }
+
+func intSlice2Int32Slice(s []int) []int32 {
+	res := make([]int32, len(s))
+	for _, i := range s {
+		res = append(res, int32(i))
+	}
+	return res
+}
diff --git a/lightning/mydump/parser.go b/lightning/mydump/parser.go
index 87d6fb5..ba4aeb3 100644
--- a/lightning/mydump/parser.go
+++ b/lightning/mydump/parser.go
@@ -112,6 +112,8 @@ type Parser interface {
 	// Columns returns the _lower-case_ column names corresponding to values in
 	// the LastRow.
 	Columns() []string
+	// SetColumns set restored column names to parser
+	SetColumns([]string)
 
 	SetLogger(log.Logger)
 }
@@ -154,6 +156,10 @@ func (parser *blockParser) Columns() []string {
 	return parser.columns
 }
 
+func (parser *blockParser) SetColumns(columns []string) {
+	parser.columns = columns
+}
+
 func (parser *blockParser) logSyntaxError() {
 	content := parser.buf
 	if len(content) > 256 {
diff --git a/lightning/restore/restore.go b/lightning/restore/restore.go
index cb08579..b6e2537 100644
--- a/lightning/restore/restore.go
+++ b/lightning/restore/restore.go
@@ -72,6 +72,10 @@ var DeliverPauser = common.NewPauser()
 func init() {
 	cfg := tidbcfg.GetGlobalConfig()
 	cfg.Log.SlowThreshold = 3000
+	// used in integration tests
+	failpoint.Inject("SetMinDeliverBytes", func(v failpoint.Value) {
+		minDeliverBytes = uint64(v.(int))
+	})
 }
 
 type saveCp struct {
@@ -956,7 +960,7 @@ func (t *TableRestore) restoreEngine(
 		// 	3. load kvs data (into kv deliver server)
 		// 	4. flush kvs data (into tikv node)
 
-		cr, err := newChunkRestore(chunkIndex, rc.cfg, chunk, rc.ioWorkers)
+		cr, err := newChunkRestore(chunkIndex, rc.cfg, chunk, rc.ioWorkers, t.tableInfo)
 		if err != nil {
 			return nil, nil, errors.Trace(err)
 		}
@@ -1247,6 +1251,7 @@ func newChunkRestore(
 	cfg *config.Config,
 	chunk *ChunkCheckpoint,
 	ioWorkers *worker.Pool,
+	tableInfo *TidbTableInfo,
 ) (*chunkRestore, error) {
 	blockBufSize := cfg.Mydumper.ReadBlockSize
 
@@ -1264,6 +1269,10 @@ func newChunkRestore(
 		parser = mydump.NewChunkParser(cfg.TiDB.SQLMode, reader, blockBufSize, ioWorkers)
 	}
 
+	if len(chunk.ColumnPermutation) > 0 {
+		parser.SetColumns(getColumnNames(tableInfo.Core, chunk.ColumnPermutation))
+	}
+
 	parser.SetPos(chunk.Chunk.Offset, chunk.Chunk.PrevRowIDMax)
 
 	return &chunkRestore{
@@ -1406,6 +1415,17 @@ func (t *TableRestore) initializeColumns(columns []string, ccp *ChunkCheckpoint)
 	ccp.ColumnPermutation = colPerm
 }
 
+func getColumnNames(tableInfo *model.TableInfo, permutation []int) []string {
+	names := make([]string, 0, len(permutation))
+	for _, idx := range permutation {
+		// skip columns with index -1
+		if idx >= 0 {
+			names = append(names, tableInfo.Columns[idx].Name.O)
+		}
+	}
+	return names
+}
+
 func (tr *TableRestore) importKV(ctx context.Context, closedEngine *kv.ClosedEngine) error {
 	task := closedEngine.Logger().Begin(zap.InfoLevel, "import and cleanup engine")
 
@@ -1538,9 +1558,9 @@ func increaseGCLifeTime(ctx context.Context, db *sql.DB) (err error) {
 
 ////////////////////////////////////////////////////////////////
 
-const (
-	maxKVQueueSize  = 128   // Cache at most this number of rows before blocking the encode loop
-	minDeliverBytes = 65536 // 64 KB. batch at least this amount of bytes to reduce number of messages
+var (
+	maxKVQueueSize         = 128   // Cache at most this number of rows before blocking the encode loop
+	minDeliverBytes uint64 = 65536 // 64 KB. batch at least this amount of bytes to reduce number of messages
 )
 
 type deliveredKVs struct {
@@ -1635,6 +1655,10 @@ func (cr *chunkRestore) deliverLoop(
 			// No need to save checkpoint if nothing was delivered.
 			saveCheckpoint(rc, t, engineID, cr.chunk)
 		}
+		failpoint.Inject("FailAfterWriteRows", func() {
+			time.Sleep(time.Second)
+			panic("forcing failure due to FailAfterWriteRows")
+		})
 		// TODO: for local backend, we may save checkpoint more frequently, e.g. after writen
 		//10GB kv pairs to data engine, we can do a flush for both data & index engine, then we
 		// can safely update current checkpoint.
@@ -1663,11 +1687,12 @@ func saveCheckpoint(rc *RestoreController, t *TableRestore, engineID int32, chun
 	rc.saveCpCh <- saveCp{
 		tableName: t.tableName,
 		merger: &ChunkCheckpointMerger{
-			EngineID: engineID,
-			Key:      chunk.Key,
-			Checksum: chunk.Checksum,
-			Pos:      chunk.Chunk.Offset,
-			RowID:    chunk.Chunk.PrevRowIDMax,
+			EngineID:          engineID,
+			Key:               chunk.Key,
+			Checksum:          chunk.Checksum,
+			Pos:               chunk.Chunk.Offset,
+			RowID:             chunk.Chunk.PrevRowIDMax,
+			ColumnPermutation: chunk.ColumnPermutation,
 		},
 	}
 }


Test patch:
diff --git a/lightning/checkpoints/checkpoints_sql_test.go b/lightning/checkpoints/checkpoints_sql_test.go
index f981cdd..0078329 100644
--- a/lightning/checkpoints/checkpoints_sql_test.go
+++ b/lightning/checkpoints/checkpoints_sql_test.go
@@ -111,7 +111,7 @@ func (s *cpSQLSuite) TestNormalOperations(c *C) {
 		ExpectPrepare("REPLACE INTO `mock-schema`\\.chunk_v\\d+ .+")
 	insertChunkStmt.
 		ExpectExec().
-		WithArgs("`db1`.`t2`", 0, "/tmp/path/1.sql", 0, 12, 102400, 1, 5000, 1234567890).
+		WithArgs("`db1`.`t2`", 0, "/tmp/path/1.sql", 0, []byte("null"), 12, 102400, 1, 5000, 1234567890).
 		WillReturnResult(sqlmock.NewResult(10, 1))
 	s.mock.ExpectCommit()
 
@@ -173,7 +173,7 @@ func (s *cpSQLSuite) TestNormalOperations(c *C) {
 		ExpectPrepare("UPDATE `mock-schema`\\.chunk_v\\d+ SET pos = .+").
 		ExpectExec().
 		WithArgs(
-			55904, 681, 4491, 586, 486070148917,
+			55904, 681, 4491, 586, 486070148917, []byte("null"),
 			"`db1`.`t2`", 0, "/tmp/path/1.sql", 0,
 		).
 		WillReturnResult(sqlmock.NewResult(11, 1))
diff --git a/lightning/restore/restore_test.go b/lightning/restore/restore_test.go
index e76600d..73af368 100644
--- a/lightning/restore/restore_test.go
+++ b/lightning/restore/restore_test.go
@@ -581,7 +581,7 @@ func (s *chunkRestoreSuite) SetUpTest(c *C) {
 	}
 
 	var err error
-	s.cr, err = newChunkRestore(1, s.cfg, &chunk, w)
+	s.cr, err = newChunkRestore(1, s.cfg, &chunk, w, nil)
 	c.Assert(err, IsNil)
 }
 
diff --git a/tests/checkpoint_columns/config.toml b/tests/checkpoint_columns/config.toml
new file mode 100644
index 0000000..c65861b
--- /dev/null
+++ b/tests/checkpoint_columns/config.toml
@@ -0,0 +1,16 @@
+[lightning]
+table-concurrency = 1
+region-concurrency = 1
+
+[checkpoint]
+enable = true
+schema = "tidb_lightning_checkpoint_test"
+driver = "mysql"
+keep-after-success = true
+
+[mydumper]
+read-block-size = 1
+
+[tikv-importer]
+max-kv-pairs = 0
+
diff --git a/tests/checkpoint_columns/run.sh b/tests/checkpoint_columns/run.sh
new file mode 100755
index 0000000..3ced606
--- /dev/null
+++ b/tests/checkpoint_columns/run.sh
@@ -0,0 +1,48 @@
+#!/bin/sh
+#
+# Copyright 2019 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.
+
+set -euE
+
+# Populate the mydumper source
+DBPATH="$TEST_DIR/cp.mydump"
+
+mkdir -p $DBPATH
+echo 'CREATE DATABASE cp_tsr;' > "$DBPATH/cp_tsr-schema-create.sql"
+echo "CREATE TABLE tbl(i TINYINT PRIMARY KEY, j INT);" > "$DBPATH/cp_tsr.tbl-schema.sql"
+# the column orders in data file is different from table schema order.
+echo "INSERT INTO tbl (j, i) VALUES (3, 1),(4, 2);" > "$DBPATH/cp_tsr.tbl.sql"
+
+# Set minDeliverBytes to a small enough number to only write only 1 row each time
+# Set the failpoint to kill the lightning instance as soon as one row is written
+export GO_FAILPOINTS="github.com/pingcap/tidb-lightning/lightning/restore/FailAfterWriteRows=return;github.com/pingcap/tidb-lightning/lightning/restore/SetMinDeliverBytes=return(1)"
+
+# Start importing the tables.
+run_sql 'DROP DATABASE IF EXISTS cp_tsr'
+run_sql 'DROP DATABASE IF EXISTS tidb_lightning_checkpoint_test'
+
+set +e
+run_lightning -d "$DBPATH" --backend tidb --enable-checkpoint=1 2> /dev/null
+set -e
+run_sql 'SELECT count(*) FROM `cp_tsr`.tbl'
+check_contains "count(*): 1"
+
+# restart lightning from checkpoint, the second line should be written successfully
+export GO_FAILPOINTS=
+set +e
+run_lightning -d "$DBPATH" --backend tidb --enable-checkpoint=1 2> /dev/null
+set -e
+
+run_sql 'SELECT j FROM `cp_tsr`.tbl WHERE i = 2;'
+check_contains "j: 4"