optimize the performance of lightning

cmd/tidb-lightning/main.go

@@ -17,6 +17,7 @@ import (
 	"fmt"
 	"os"
 	"os/signal"
+	"runtime/debug"
 	"syscall"
 
 	"github.com/pingcap/tidb-lightning/lightning"
@@ -44,6 +45,21 @@ func main() {
 
 	logger := log.L()
 
+	// Lightning allocates too many transient objects and heap size is small,
+	// so garbage collections happen too frequently and lots of time is spent in GC component.
+	//
+	// In a test of loading the table `order_line.csv` of 14k TPCC.
+	// The time need of `encode kv data and write` step reduce from 52m4s to 37m30s when change
+	// GOGC from 100 to 500, the total time needed reduce near 15m too.
+	// The cost of this is the memory of lightnin at runtime grow from about 200M to 700M, but it's acceptable.
+	//
+	// So we set the gc percentage as 500 default to reduce the GC frequency instead of 100.
+	gogc := os.Getenv("GOGC")
+	if gogc == "" {
+		old := debug.SetGCPercent(500)
+		log.L().Debug("set gc percentage", zap.Int("old", old), zap.Int("new", 500))
+	}
+
 	err := app.GoServe()
 	if err != nil {
 		logger.Error("failed to start HTTP server", zap.Error(err))

go.mod

@@ -10,13 +10,13 @@ require (
 	github.com/gogo/protobuf v1.3.1
 	github.com/golang/mock v1.4.0
 	github.com/joho/sqltocsv v0.0.0-20190824231449-5650f27fd5b6
-	github.com/pingcap/check v0.0.0-20191216031241-8a5a85928f12
+	github.com/pingcap/check v0.0.0-20200212061837-5e12011dc712
 	github.com/pingcap/errors v0.11.5-0.20190809092503-95897b64e011
-	github.com/pingcap/failpoint v0.0.0-20191029060244-12f4ac2fd11d
-	github.com/pingcap/kvproto v0.0.0-20200108025604-a4dc183d2af5
-	github.com/pingcap/log v0.0.0-20191012051959-b742a5d432e9
-	github.com/pingcap/parser v0.0.0-20200109073933-a9496438d77d
-	github.com/pingcap/tidb v1.1.0-beta.0.20200110130413-8c3ee37c1938
+	github.com/pingcap/failpoint v0.0.0-20200210140405-f8f9fb234798
+	github.com/pingcap/kvproto v0.0.0-20200221125103-35b65c96516e
+	github.com/pingcap/log v0.0.0-20200117041106-d28c14d3b1cd
+	github.com/pingcap/parser v0.0.0-20200218113622-517beb2e39c2
+	github.com/pingcap/tidb v1.1.0-beta.0.20200303031609-495f8b74382f
 	github.com/pingcap/tidb-tools v4.0.0-beta+incompatible
 	github.com/prometheus/client_golang v1.4.0
 	github.com/prometheus/client_model v0.2.0

lightning/backend/importer.go

@@ -16,6 +16,7 @@ package backend
 import (
 	"context"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/pingcap/errors"
@@ -39,6 +40,8 @@ type importer struct {
 	conn   *grpc.ClientConn
 	cli    kv.ImportKVClient
 	pdAddr string
+
+	mutationPool sync.Pool
 }
 
 // NewImporter creates a new connection to tikv-importer. A single connection
@@ -50,9 +53,10 @@ func NewImporter(ctx context.Context, tls *common.TLS, importServerAddr string,
 	}
 
 	return MakeBackend(&importer{
-		conn:   conn,
-		cli:    kv.NewImportKVClient(conn),
-		pdAddr: pdAddr,
+		conn:         conn,
+		cli:          kv.NewImportKVClient(conn),
+		pdAddr:       pdAddr,
+		mutationPool: sync.Pool{New: func() interface{} { return &kv.Mutation{} }},
 	}), nil
 }
 
@@ -61,9 +65,10 @@ func NewImporter(ctx context.Context, tls *common.TLS, importServerAddr string,
 // outside of tests.
 func NewMockImporter(cli kv.ImportKVClient, pdAddr string) Backend {
 	return MakeBackend(&importer{
-		conn:   nil,
-		cli:    cli,
-		pdAddr: pdAddr,
+		conn:         nil,
+		cli:          cli,
+		pdAddr:       pdAddr,
+		mutationPool: sync.Pool{New: func() interface{} { return &kv.Mutation{} }},
 	})
 }
 
@@ -190,11 +195,10 @@ func (importer *importer) WriteRows(
 	// Send kv paris as write request content
 	mutations := make([]*kv.Mutation, len(kvs))
 	for i, pair := range kvs {
-		mutations[i] = &kv.Mutation{
-			Op:    kv.Mutation_Put,
-			Key:   pair.Key,
-			Value: pair.Val,
-		}
+		mutations[i] = importer.mutationPool.Get().(*kv.Mutation)
+		mutations[i].Op = kv.Mutation_Put
+		mutations[i].Key = pair.Key
+		mutations[i].Value = pair.Val
 	}
 
 	req.Reset()
@@ -205,7 +209,12 @@ func (importer *importer) WriteRows(
 		},
 	}
 
-	if err := wstream.Send(req); err != nil {
+	err = wstream.Send(req)
+	for _, mutation := range mutations {
+		importer.mutationPool.Put(mutation)
+	}
+
+	if err != nil {
 		return errors.Trace(err)
 	}
 

lightning/backend/importer_test.go

@@ -2,13 +2,16 @@ package backend_test
 
 import (
 	"context"
+	"sync"
+	"testing"
 
 	"github.com/golang/mock/gomock"
 	. "github.com/pingcap/check"
 	"github.com/pingcap/errors"
 	"github.com/pingcap/kvproto/pkg/import_kvpb"
 	uuid "github.com/satori/go.uuid"
 
+	kvpb "github.com/pingcap/kvproto/pkg/import_kvpb"
 	kv "github.com/pingcap/tidb-lightning/lightning/backend"
 	"github.com/pingcap/tidb-lightning/lightning/common"
 	"github.com/pingcap/tidb-lightning/mock"
@@ -193,3 +196,39 @@ func (s *importerSuite) TestCloseImportCleanupEngine(c *C) {
 	err = engine.Cleanup(s.ctx)
 	c.Assert(err, IsNil)
 }
+
+func BenchmarkMutationAlloc(b *testing.B) {
+	var g *kvpb.Mutation
+	for i := 0; i < b.N; i++ {
+		m := &kvpb.Mutation{
+			Op:    kvpb.Mutation_Put,
+			Key:   nil,
+			Value: nil,
+		}
+		g = m
+	}
+
+	var _ = g
+}
+
+func BenchmarkMutationPool(b *testing.B) {
+	p := sync.Pool{
+		New: func() interface{} {
+			return &kvpb.Mutation{}
+		},
+	}
+	var g *kvpb.Mutation
+
+	for i := 0; i < b.N; i++ {
+		m := p.Get().(*kvpb.Mutation)
+		m.Op = kvpb.Mutation_Put
+		m.Key = nil
+		m.Value = nil
+
+		g = m
+
+		p.Put(m)
+	}
+
+	var _ = g
+}

lightning/backend/session.go

@@ -106,17 +106,14 @@ func newSession(options *SessionOptions) *session {
 	vars.StmtCtx.TimeZone = vars.Location()
 	vars.SetSystemVar("timestamp", strconv.FormatInt(options.Timestamp, 10))
 	vars.SetSystemVar(variable.TiDBRowFormatVersion, options.RowFormatVersion)
+	vars.TxnCtx = nil
 
 	s := &session{
 		txn:  transaction{},
 		vars: vars,
 	}
 	s.vars.GetWriteStmtBufs().BufStore = &kv.BufferStore{MemBuffer: &s.txn}
 
-	// FIXME: potential further improvements:
-	//  - (*TransactionContext).UpdateDeltaForTable takes 14% of Encode() time.
-	//    But we cannot make `vars.TxnCtx` into nil nor make it no-op.
-
 	return s
 }
 

lightning/mydump/csv_parser.go

@@ -50,6 +50,8 @@ type CSVParser struct {
 	// fieldIndexes is an index of fields inside recordBuffer.
 	// The i'th field ends at offset fieldIndexes[i] in recordBuffer.
 	fieldIndexes []int
+
+	lastRecord []string
 }
 
 func NewCSVParser(
@@ -163,7 +165,7 @@ func (parser *CSVParser) readUntil(chars string) ([]byte, byte, error) {
 	}
 }
 
-func (parser *CSVParser) readRecord() ([]string, error) {
+func (parser *CSVParser) readRecord(dst []string) ([]string, error) {
 	parser.recordBuffer = parser.recordBuffer[:0]
 	parser.fieldIndexes = parser.fieldIndexes[:0]
 
@@ -214,7 +216,11 @@ outside:
 	// Create a single string and create slices out of it.
 	// This pins the memory of the fields together, but allocates once.
 	str := string(parser.recordBuffer) // Convert to string once to batch allocations
-	dst := make([]string, len(parser.fieldIndexes))
+	dst = dst[:0]
+	if cap(dst) < len(parser.fieldIndexes) {
+		dst = make([]string, len(parser.fieldIndexes))
+	}
+	dst = dst[:len(parser.fieldIndexes)]
 	var preIdx int
 	for i, idx := range parser.fieldIndexes {
 		dst[i] = str[preIdx:idx]
@@ -315,7 +321,7 @@ func (parser *CSVParser) ReadRow() error {
 
 	// skip the header first
 	if parser.pos == 0 && parser.cfg.Header {
-		columns, err := parser.readRecord()
+		columns, err := parser.readRecord(nil)
 		if err != nil {
 			return errors.Trace(err)
 		}
@@ -326,10 +332,11 @@ func (parser *CSVParser) ReadRow() error {
 		}
 	}
 
-	records, err := parser.readRecord()
+	records, err := parser.readRecord(parser.lastRecord)
 	if err != nil {
 		return errors.Trace(err)
 	}
+	parser.lastRecord = records
 	// remove trailing empty values
 	if parser.cfg.TrimLastSep {
 		var i int
@@ -345,7 +352,7 @@ func (parser *CSVParser) ReadRow() error {
 		if isNull {
 			datum.SetNull()
 		} else {
-			datum.SetString(unescaped)
+			datum.SetString(unescaped, "utf8mb4_bin")
 		}
 		row.Row = append(row.Row, datum)
 	}

lightning/mydump/parser.go

@@ -472,7 +472,7 @@ func (parser *ChunkParser) ReadRow() error {
 				// can't handle integers more than 64 bits anyway)
 				fallthrough
 			case tokUnquoted, tokSingleQuoted, tokDoubleQuoted:
-				value.SetString(parser.unescapeString(string(content)))
+				value.SetString(parser.unescapeString(string(content)), "utf8mb4_bin")
 			case tokHexString:
 				hexLit, err := types.ParseHexStr(string(content))
 				if err != nil {

lightning/restore/restore.go

@@ -581,7 +581,7 @@ func (rc *RestoreController) restoreTables(ctx context.Context) error {
 
 	var restoreErr common.OnceError
 
-	stopPeriodicActions := make(chan struct{}, 1)
+	stopPeriodicActions := make(chan struct{})
 	go rc.runPeriodicActions(ctx, stopPeriodicActions)
 
 	type task struct {
@@ -691,7 +691,7 @@ func (rc *RestoreController) restoreTables(ctx context.Context) error {
 	}
 
 	wg.Wait()
-	stopPeriodicActions <- struct{}{}
+	close(stopPeriodicActions)
 
 	err := restoreErr.Get()
 	logTask.End(zap.ErrorLevel, err)
@@ -948,7 +948,7 @@ func (t *TableRestore) restoreEngine(
 	totalSQLSize := int64(0)
 	for _, chunk := range cp.Chunks {
 		totalKVSize += chunk.Checksum.SumSize()
-		totalSQLSize += chunk.Chunk.EndOffset
+		totalSQLSize += chunk.Chunk.EndOffset - chunk.Chunk.Offset
 	}
 
 	err = chunkErr.Get()

tools/go.mod

@@ -5,6 +5,6 @@ go 1.12
 require (
 	github.com/golangci/golangci-lint v1.22.2
 	github.com/mgechev/revive v1.0.0
-	github.com/pingcap/failpoint v0.0.0-20191029060244-12f4ac2fd11d
+	github.com/pingcap/failpoint v0.0.0-20200210140405-f8f9fb234798
 	github.com/shurcooL/vfsgen v0.0.0-20181202132449-6a9ea43bcacd
 )

tools/go.sum

@@ -1,4 +1,5 @@
 cloud.google.com/go v0.26.0/go.mod h1:aQUYkXzVsufM+DwF1aE+0xfcU+56JwCaLick0ClmMTw=
+github.com/BurntSushi/toml v0.3.1 h1:WXkYYl6Yr3qBf1K79EBnL4mak0OimBfB0XUf9Vl28OQ=
 github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
 github.com/OneOfOne/xxhash v1.2.2/go.mod h1:HSdplMjZKSmBqAxg5vPj2TmRDmfkzw+cTzAElWljhcU=
 github.com/OpenPeeDeeP/depguard v1.0.1/go.mod h1:xsIw86fROiiwelg+jB2uM9PiKihMMmUx/1V+TNhjQvM=
@@ -26,8 +27,10 @@ github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSs
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/dgrijalva/jwt-go v3.2.0+incompatible/go.mod h1:E3ru+11k8xSBh+hMPgOLZmtrrCbhqsmaPHjLKYnJCaQ=
 github.com/dgryski/go-sip13 v0.0.0-20181026042036-e10d5fee7954/go.mod h1:vAd38F8PWV+bWy6jNmig1y/TA+kYO4g3RSRF0IAv0no=
+github.com/fatih/color v1.7.0 h1:DkWD4oS2D8LGGgTQ6IvwJJXSL5Vp2ffcQg58nFV38Ys=
 github.com/fatih/color v1.7.0/go.mod h1:Zm6kSWBoL9eyXnKyktHP6abPY2pDugNf5KwzbycvMj4=
 github.com/fatih/structtag v1.0.0/go.mod h1:IKitwq45uXL/yqi5mYghiD3w9H6eTOvI9vnk8tXMphA=
+github.com/fatih/structtag v1.2.0 h1:/OdNE99OxoI/PqaW/SuSK9uxxT3f/tcSZgon/ssNSx4=
 github.com/fatih/structtag v1.2.0/go.mod h1:mBJUNpUnHmRKrKlQQlmCrh5PuhftFbNv8Ys4/aAZl94=
 github.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo=
 github.com/ghodss/yaml v1.0.0/go.mod h1:4dBDuWmgqj2HViK6kFavaiC9ZROes6MMH2rRYeMEF04=
@@ -117,17 +120,23 @@ github.com/magiconair/properties v1.8.0/go.mod h1:PppfXfuXeibc/6YijjN8zIbojt8czP
 github.com/magiconair/properties v1.8.1/go.mod h1:PppfXfuXeibc/6YijjN8zIbojt8czPbwD3XqdrwzmxQ=
 github.com/matoous/godox v0.0.0-20190911065817-5d6d842e92eb/go.mod h1:1BELzlh859Sh1c6+90blK8lbYy0kwQf1bYlBhBysy1s=
 github.com/mattn/go-colorable v0.1.2/go.mod h1:U0ppj6V5qS13XJ6of8GYAs25YV2eR4EVcfRqFIhoBtE=
+github.com/mattn/go-colorable v0.1.4 h1:snbPLB8fVfU9iwbbo30TPtbLRzwWu6aJS6Xh4eaaviA=
 github.com/mattn/go-colorable v0.1.4/go.mod h1:U0ppj6V5qS13XJ6of8GYAs25YV2eR4EVcfRqFIhoBtE=
 github.com/mattn/go-isatty v0.0.8/go.mod h1:Iq45c/XA43vh69/j3iqttzPXn0bhXyGjM0Hdxcsrc5s=
+github.com/mattn/go-isatty v0.0.9 h1:d5US/mDsogSGW37IV293h//ZFaeajb69h+EHFsv2xGg=
 github.com/mattn/go-isatty v0.0.9/go.mod h1:YNRxwqDuOph6SZLI9vUUz6OYw3QyUt7WiY2yME+cCiQ=
 github.com/mattn/go-runewidth v0.0.4/go.mod h1:LwmH8dsx7+W8Uxz3IHJYH5QSwggIsqBzpuz5H//U1FU=
+github.com/mattn/go-runewidth v0.0.7 h1:Ei8KR0497xHyKJPAv59M1dkC+rOZCMBJ+t3fZ+twI54=
 github.com/mattn/go-runewidth v0.0.7/go.mod h1:H031xJmbD/WCDINGzjvQ9THkh0rPKHF+m2gUSrubnMI=
 github.com/mattn/goveralls v0.0.2/go.mod h1:8d1ZMHsd7fW6IRPKQh46F2WRpyib5/X4FOpevwGNQEw=
 github.com/matttproud/golang_protobuf_extensions v1.0.1/go.mod h1:D8He9yQNgCq6Z5Ld7szi9bcBfOoFv/3dc6xSMkL2PC0=
 github.com/mgechev/dots v0.0.0-20181228164730-18fa4c4b71cc/go.mod h1:KQ7+USdGKfpPjXk4Ga+5XxQM4Lm4e3gAogrreFAYpOg=
+github.com/mgechev/dots v0.0.0-20190921121421-c36f7dcfbb81 h1:QASJXOGm2RZ5Ardbc86qNFvby9AqkLDibfChMtAg5QM=
 github.com/mgechev/dots v0.0.0-20190921121421-c36f7dcfbb81/go.mod h1:KQ7+USdGKfpPjXk4Ga+5XxQM4Lm4e3gAogrreFAYpOg=
 github.com/mgechev/revive v0.0.0-20190917153825-40564c5052ae/go.mod h1:37hJOqkogcmT5nmiriskuzkdJ/YhMlZwFSg87NDZbco=
+github.com/mgechev/revive v1.0.0 h1:CaN0UEsXgAovl5WB4yI7obZG+NhBh7HeEuQX11wVmTI=
 github.com/mgechev/revive v1.0.0/go.mod h1:Vwga5U/ducQXkYcbsNPMIOA1kFBVY1NBPs5Fa4Nfceg=
+github.com/mitchellh/go-homedir v1.1.0 h1:lukF9ziXFxDFPkA1vsr5zpc1XuPDn/wFntq5mG+4E0Y=
 github.com/mitchellh/go-homedir v1.1.0/go.mod h1:SfyaCUpYCn1Vlf4IUYiD9fPX4A5wJrkLzIz1N1q0pr0=
 github.com/mitchellh/go-ps v0.0.0-20190716172923-621e5597135b/go.mod h1:r1VsdOzOPt1ZSrGZWFoNhsAedKnEd6r9Np1+5blZCWk=
 github.com/mitchellh/mapstructure v1.1.2/go.mod h1:FVVH3fgwuzCH5S8UJGiWEs2h04kUh9fWfEaFds41c1Y=
@@ -137,6 +146,7 @@ github.com/nbutton23/zxcvbn-go v0.0.0-20180912185939-ae427f1e4c1d/go.mod h1:o96d
 github.com/nicksnyder/go-i18n v1.10.0/go.mod h1:HrK7VCrbOvQoUAQ7Vpy7i87N7JZZZ7R2xBGjv0j365Q=
 github.com/oklog/ulid v1.3.1/go.mod h1:CirwcVhetQ6Lv90oh/F+FBtV6XMibvdAFo93nm5qn4U=
 github.com/olekukonko/tablewriter v0.0.1/go.mod h1:vsDQFd/mU46D+Z4whnwzcISnGGzXWMclvtLoiIKAKIo=
+github.com/olekukonko/tablewriter v0.0.4 h1:vHD/YYe1Wolo78koG299f7V/VAS08c6IpCLn+Ejf/w8=
 github.com/olekukonko/tablewriter v0.0.4/go.mod h1:zq6QwlOf5SlnkVbMSr5EoBv3636FWnp+qbPhuoO21uA=
 github.com/onsi/ginkgo v1.6.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE=
 github.com/onsi/ginkgo v1.10.1/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE=
@@ -151,6 +161,8 @@ github.com/pingcap/failpoint v0.0.0-20190708053854-e7b1061e6e81 h1:g9j4cAoWE39lh
 github.com/pingcap/failpoint v0.0.0-20190708053854-e7b1061e6e81/go.mod h1:DNS3Qg7bEDhU6EXNHF+XSv/PGznQaMJ5FWvctpm6pQI=
 github.com/pingcap/failpoint v0.0.0-20191029060244-12f4ac2fd11d h1:F8vp38kTAckN+v8Jlc98uMBvKIzr1a+UhnLyVYn8Q5Q=
 github.com/pingcap/failpoint v0.0.0-20191029060244-12f4ac2fd11d/go.mod h1:DNS3Qg7bEDhU6EXNHF+XSv/PGznQaMJ5FWvctpm6pQI=
+github.com/pingcap/failpoint v0.0.0-20200210140405-f8f9fb234798 h1:6DMbRqPI1qzQ8N1xc3+nKY8IxSACd9VqQKkRVvbyoIg=
+github.com/pingcap/failpoint v0.0.0-20200210140405-f8f9fb234798/go.mod h1:DNS3Qg7bEDhU6EXNHF+XSv/PGznQaMJ5FWvctpm6pQI=
 github.com/pkg/errors v0.8.0/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
@@ -280,6 +292,7 @@ golang.org/x/tools v0.0.0-20190910044552-dd2b5c81c578/go.mod h1:b+2E5dAYhXwXZwtn
 golang.org/x/tools v0.0.0-20190930201159-7c411dea38b0/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
 golang.org/x/tools v0.0.0-20191010075000-0337d82405ff/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
 golang.org/x/tools v0.0.0-20191113232020-e2727e816f5a/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
+golang.org/x/tools v0.0.0-20191227053925-7b8e75db28f4 h1:Toz2IK7k8rbltAXwNAxKcn9OzqyNfMUhUNjz3sL0NMk=
 golang.org/x/tools v0.0.0-20191227053925-7b8e75db28f4/go.mod h1:TB2adYChydJhpapKDTa4BR/hXlZSLoq2Wpct/0txZ28=
 golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=