executor: fix load data losing connection when batch_dml_size is set (#22724)

executor/builder.go

@@ -737,6 +737,8 @@ func (b *executorBuilder) buildLoadData(v *plannercore.LoadData) Executor {
 		Table:        tbl,
 		Columns:      v.Columns,
 		GenExprs:     v.GenCols.Exprs,
+		isLoadData:   true,
+		txnInUse:     sync.Mutex{},
 	}
 	err := insertVal.initInsertColumns()
 	if err != nil {

executor/insert_common.go

@@ -18,6 +18,7 @@ import (
 	"context"
 	"fmt"
 	"math"
+	"sync"
 	"time"
 
 	"github.com/opentracing/opentracing-go"
@@ -81,6 +82,12 @@ type InsertValues struct {
 	memTracker     *memory.Tracker
 
 	stats *InsertRuntimeStat
+
+	// LoadData use two goroutines. One for generate batch data,
+	// The other one for commit task, which will invalid txn.
+	// We use mutex to protect routine from using invalid txn.
+	isLoadData bool
+	txnInUse   sync.Mutex
 }
 
 type defaultVal struct {
@@ -857,10 +864,6 @@ func (e *InsertValues) adjustAutoRandomDatum(ctx context.Context, d types.Datum,
 	// Change NULL to auto id.
 	// Change value 0 to auto id, if NoAutoValueOnZero SQL mode is not set.
 	if d.IsNull() || e.ctx.GetSessionVars().SQLMode&mysql.ModeNoAutoValueOnZero == 0 {
-		_, err := e.ctx.Txn(true)
-		if err != nil {
-			return types.Datum{}, errors.Trace(err)
-		}
 		recordID, err = e.allocAutoRandomID(&c.FieldType)
 		if err != nil {
 			return types.Datum{}, err
@@ -894,6 +897,14 @@ func (e *InsertValues) allocAutoRandomID(fieldType *types.FieldType) (int64, err
 	if tables.OverflowShardBits(autoRandomID, tableInfo.AutoRandomBits, layout.TypeBitsLength, layout.HasSignBit) {
 		return 0, autoid.ErrAutoRandReadFailed
 	}
+	if e.isLoadData {
+		e.txnInUse.Lock()
+		defer e.txnInUse.Unlock()
+	}
+	_, err = e.ctx.Txn(true)
+	if err != nil {
+		return 0, err
+	}
 	shard := tables.CalcShard(tableInfo.AutoRandomBits, e.ctx.GetSessionVars().TxnCtx.StartTS, layout.TypeBitsLength, layout.HasSignBit)
 	autoRandomID |= shard
 	return autoRandomID, nil

executor/load_data.go

@@ -199,6 +199,8 @@ func (e *LoadDataInfo) CommitOneTask(ctx context.Context, task CommitTask) error
 		logutil.Logger(ctx).Error("commit error commit", zap.Error(err))
 		return err
 	}
+	e.txnInUse.Lock()
+	defer e.txnInUse.Unlock()
 	// Make sure that there are no retries when committing.
 	if err = e.Ctx.RefreshTxnCtx(ctx); err != nil {
 		logutil.Logger(ctx).Error("commit error refresh", zap.Error(err))

server/server_test.go

@@ -903,6 +903,74 @@ func (cli *testServerClient) runTestLoadData(c *C, server *Server) {
 	})
 }
 
+func (cli *testServerClient) runTestLoadDataAutoRandom(c *C) {
+	path := "/tmp/load_data_txn_error.csv"
+
+	fp, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0600)
+	c.Assert(err, IsNil)
+	c.Assert(fp, NotNil)
+
+	defer func() {
+		_ = os.Remove(path)
+	}()
+
+	cksum1 := 0
+	cksum2 := 0
+	for i := 0; i < 50000; i++ {
+		n1 := rand.Intn(1000)
+		n2 := rand.Intn(1000)
+		str1 := strconv.Itoa(n1)
+		str2 := strconv.Itoa(n2)
+		row := str1 + "\t" + str2
+		_, err := fp.WriteString(row)
+		c.Assert(err, IsNil)
+		_, err = fp.WriteString("\n")
+		c.Assert(err, IsNil)
+
+		if i == 0 {
+			cksum1 = n1
+			cksum2 = n2
+		} else {
+			cksum1 = cksum1 ^ n1
+			cksum2 = cksum2 ^ n2
+		}
+	}
+
+	err = fp.Close()
+	c.Assert(err, IsNil)
+
+	cli.runTestsOnNewDB(c, func(config *mysql.Config) {
+		config.AllowAllFiles = true
+		config.Params = map[string]string{"sql_mode": "''"}
+	}, "load_data_batch_dml", func(dbt *DBTest) {
+		// Set batch size, and check if load data got a invalid txn error.
+		dbt.mustExec("set @@session.tidb_dml_batch_size = 128")
+		dbt.mustExec("drop table if exists t")
+		dbt.mustExec("create table t(c1 bigint auto_random primary key, c2 bigint, c3 bigint)")
+		dbt.mustExec(fmt.Sprintf("load data local infile %q into table t (c2, c3)", path))
+
+		var (
+			rowCnt    int
+			colCkSum1 int
+			colCkSum2 int
+		)
+		rows := dbt.mustQuery("select count(*) from t")
+		dbt.Check(rows.Next(), IsTrue, Commentf("unexpected data"))
+		err = rows.Scan(&rowCnt)
+		dbt.Check(err, IsNil)
+		dbt.Check(rowCnt, DeepEquals, 50000)
+		dbt.Check(rows.Next(), IsFalse, Commentf("unexpected data"))
+
+		rows = dbt.mustQuery("select bit_xor(c2), bit_xor(c3) from t")
+		dbt.Check(rows.Next(), IsTrue, Commentf("unexpected data"))
+		err = rows.Scan(&colCkSum1, &colCkSum2)
+		dbt.Check(err, IsNil)
+		dbt.Check(colCkSum1, DeepEquals, cksum1)
+		dbt.Check(colCkSum2, DeepEquals, cksum2)
+		dbt.Check(rows.Next(), IsFalse, Commentf("unexpected data"))
+	})
+}
+
 func (cli *testServerClient) runTestConcurrentUpdate(c *C) {
 	dbName := "Concurrent"
 	cli.runTestsOnNewDB(c, func(config *mysql.Config) {

server/tidb_test.go

@@ -148,6 +148,12 @@ func (ts *tidbTestSerialSuite) TestLoadData(c *C) {
 	ts.runTestLoadDataForSlowLog(c, ts.server)
 }
 
+// Fix issue#22540. Change tidb_dml_batch_size,
+// then check if load data into table with auto random column works properly.
+func (ts *tidbTestSerialSuite) TestLoadDataAutoRandom(c *C) {
+	ts.runTestLoadDataAutoRandom(c)
+}
+
 func (ts *tidbTestSerialSuite) TestExplainFor(c *C) {
 	ts.runTestExplainForConn(c)
 }