lightning: add send-kv-size to avoid oom when each kv is large on default config

br/pkg/lightning/backend/local/local.go

@@ -394,8 +394,10 @@ type BackendConfig struct {
 	// compress type when write or ingest into tikv
 	ConnCompressType config.CompressionType
 	// concurrency of generateJobForRange and import(write & ingest) workers
-	WorkerConcurrency      int
-	KVWriteBatchSize       int
+	WorkerConcurrency int
+	// batch kv count and size when writing to TiKV
+	KVWriteBatchCount      int
+	KVWriteBatchSize       int64
 	RegionSplitBatchSize   int
 	RegionSplitConcurrency int
 	CheckpointEnabled      bool
@@ -430,7 +432,8 @@ func NewBackendConfig(cfg *config.Config, maxOpenFiles int, keyspaceName string)
 		MaxConnPerStore:         cfg.TikvImporter.RangeConcurrency,
 		ConnCompressType:        cfg.TikvImporter.CompressKVPairs,
 		WorkerConcurrency:       cfg.TikvImporter.RangeConcurrency * 2,
-		KVWriteBatchSize:        cfg.TikvImporter.SendKVPairs,
+		KVWriteBatchCount:       cfg.TikvImporter.SendKVPairs,
+		KVWriteBatchSize:        int64(cfg.TikvImporter.SendKVSize),
 		RegionSplitBatchSize:    cfg.TikvImporter.RegionSplitBatchSize,
 		RegionSplitConcurrency:  cfg.TikvImporter.RegionSplitConcurrency,
 		CheckpointEnabled:       cfg.Checkpoint.Enable,

br/pkg/lightning/backend/local/region_job.go

@@ -17,6 +17,7 @@ package local
 import (
 	"container/heap"
 	"context"
+	"fmt"
 	"strings"
 	"sync"
 	"time"
@@ -177,6 +178,7 @@ func (local *Backend) writeToTiKV(ctx context.Context, j *regionJob) error {
 
 	apiVersion := local.tikvCodec.GetAPIVersion()
 	clientFactory := local.importClientFactory
+	kvBatchCount := local.KVWriteBatchCount
 	kvBatchSize := local.KVWriteBatchSize
 	bufferPool := local.bufferPool
 	writeLimiter := local.writeLimiter
@@ -254,7 +256,7 @@ func (local *Backend) writeToTiKV(ctx context.Context, j *regionJob) error {
 
 	bytesBuf := bufferPool.NewBuffer()
 	defer bytesBuf.Destroy()
-	pairs := make([]*sst.Pair, 0, kvBatchSize)
+	pairs := make([]*sst.Pair, 0, kvBatchCount)
 	count := 0
 	size := int64(0)
 	totalSize := int64(0)
@@ -265,8 +267,6 @@ func (local *Backend) writeToTiKV(ctx context.Context, j *regionJob) error {
 	if j.regionSplitSize <= int64(config.SplitRegionSize) {
 		regionMaxSize = j.regionSplitSize * 4 / 3
 	}
-	// Set a lower flush limit to make the speed of write more smooth.
-	flushLimit := int64(writeLimiter.Limit() / 10)
 
 	flushKVs := func() error {
 		for i := range clients {
@@ -278,6 +278,9 @@ func (local *Backend) writeToTiKV(ctx context.Context, j *regionJob) error {
 				return annotateErr(err, allPeers[i])
 			}
 		}
+		failpoint.Inject("afterFlushKVs", func() {
+			log.FromContext(ctx).Info(fmt.Sprintf("afterFlushKVs count=%d,size=%d", count, size))
+		})
 		return nil
 	}
 
@@ -305,7 +308,7 @@ func (local *Backend) writeToTiKV(ctx context.Context, j *regionJob) error {
 		size += kvSize
 		totalSize += kvSize
 
-		if count >= kvBatchSize || size >= flushLimit {
+		if count >= kvBatchCount || size >= kvBatchSize {
 			if err := flushKVs(); err != nil {
 				return errors.Trace(err)
 			}

br/pkg/lightning/config/config.go

@@ -70,7 +70,10 @@ const (
 	// ErrorOnDup indicates using INSERT INTO to insert data, which would violate PK or UNIQUE constraint
 	ErrorOnDup = "error"
 
-	KVWriteBatchSize        = 32768
+	KVWriteBatchCount = 32768
+	// KVWriteBatchSize batch size when write to TiKV.
+	// this is the default value of linux send buffer size(net.ipv4.tcp_wmem) too.
+	KVWriteBatchSize        = 16 * units.KiB
 	DefaultRangeConcurrency = 16
 
 	defaultDistSQLScanConcurrency     = 15
@@ -745,6 +748,7 @@ type TikvImporter struct {
 	OnDuplicate             string                       `toml:"on-duplicate" json:"on-duplicate"`
 	MaxKVPairs              int                          `toml:"max-kv-pairs" json:"max-kv-pairs"`
 	SendKVPairs             int                          `toml:"send-kv-pairs" json:"send-kv-pairs"`
+	SendKVSize              ByteSize                     `toml:"send-kv-size" json:"send-kv-size"`
 	CompressKVPairs         CompressionType              `toml:"compress-kv-pairs" json:"compress-kv-pairs"`
 	RegionSplitSize         ByteSize                     `toml:"region-split-size" json:"region-split-size"`
 	RegionSplitKeys         int                          `toml:"region-split-keys" json:"region-split-keys"`
@@ -952,7 +956,8 @@ func NewConfig() *Config {
 			Backend:                 "",
 			OnDuplicate:             ReplaceOnDup,
 			MaxKVPairs:              4096,
-			SendKVPairs:             KVWriteBatchSize,
+			SendKVPairs:             KVWriteBatchCount,
+			SendKVSize:              KVWriteBatchSize,
 			RegionSplitSize:         0,
 			RegionSplitBatchSize:    DefaultRegionSplitBatchSize,
 			RegionSplitConcurrency:  runtime.GOMAXPROCS(0),

br/pkg/lightning/config/config_test.go

@@ -865,6 +865,9 @@ func TestDefaultCouldBeOverwritten(t *testing.T) {
 	require.Equal(t, 20, cfg.App.IndexConcurrency)
 	require.Equal(t, 60, cfg.App.TableConcurrency)
 
+	require.Equal(t, config.KVWriteBatchCount, cfg.TikvImporter.SendKVPairs)
+	require.Equal(t, config.ByteSize(config.KVWriteBatchSize), cfg.TikvImporter.SendKVSize)
+
 	cfg.TikvImporter.RegionSplitConcurrency = 1
 	// backoff can be 0
 	cfg.TikvImporter.RegionCheckBackoffLimit = 0

br/tests/lightning_write_batch/kv-count.toml

@@ -0,0 +1,6 @@
+[tikv-importer]
+range-concurrency = 1
+send-kv-pairs = 20
+
+[mydumper.csv]
+header = false

br/tests/lightning_write_batch/kv-size.toml

@@ -0,0 +1,6 @@
+[tikv-importer]
+range-concurrency = 1
+send-kv-size = "10b"
+
+[mydumper.csv]
+header = false

br/tests/lightning_write_batch/run.sh

@@ -0,0 +1,66 @@
+#!/bin/bash
+#
+# Copyright 2023 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,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+set -eux
+
+mkdir -p "$TEST_DIR/data"
+
+run_sql "DROP DATABASE IF EXISTS test;"
+run_sql "DROP TABLE IF EXISTS test.t;"
+
+cat <<EOF >"$TEST_DIR/data/test-schema-create.sql"
+CREATE DATABASE test;
+EOF
+cat <<EOF >"$TEST_DIR/data/test.t-schema.sql"
+CREATE TABLE test.t (a varchar(1024));
+EOF
+
+#
+# test send-kv-pairs
+#
+set +x
+for i in {1..100}; do
+  echo "$i" >>"$TEST_DIR/data/test.t.0.csv"
+done
+set -x
+
+rm -rf $TEST_DIR/lightning.log
+export GO_FAILPOINTS="github.com/pingcap/tidb/br/pkg/lightning/backend/local/afterFlushKVs=return(true)"
+run_lightning --backend local -d "$TEST_DIR/data" --config "tests/$TEST_NAME/kv-count.toml"
+check_contains 'afterFlushKVs count=20,' $TEST_DIR/lightning.log
+check_not_contains 'afterFlushKVs count=1,' $TEST_DIR/lightning.log
+check_contains 'send-kv-pairs\":20,' $TEST_DIR/lightning.log
+check_contains 'send-kv-size\":16384,' $TEST_DIR/lightning.log
+
+#
+# test send-kv-size
+#
+rm -rf $TEST_DIR/data/test.t.0.csv
+run_sql "truncate table test.t;"
+set +x
+for i in {1..5}; do
+  echo "abcdefghijklmnopqrstuvwxyz0123456789" >>"$TEST_DIR/data/test.t.0.csv"
+done
+set -x
+
+rm -rf $TEST_DIR/lightning.log
+export GO_FAILPOINTS="github.com/pingcap/tidb/br/pkg/lightning/backend/local/afterFlushKVs=return(true)"
+run_lightning --backend local -d "$TEST_DIR/data" --config "tests/$TEST_NAME/kv-size.toml"
+# each kv is 64b, so each kv is a batch
+check_contains 'afterFlushKVs count=1,' $TEST_DIR/lightning.log
+check_not_contains 'afterFlushKVs count=20,' $TEST_DIR/lightning.log
+check_contains 'send-kv-pairs\":32768,' $TEST_DIR/lightning.log
+check_contains 'send-kv-size\":10,' $TEST_DIR/lightning.log

executor/importer/table_import.go

@@ -116,6 +116,7 @@ func NewTableImporter(param *JobImportParam, e *LoadDataController) (ti *TableIm
 		MaxConnPerStore:        config.DefaultRangeConcurrency,
 		ConnCompressType:       config.CompressionNone,
 		WorkerConcurrency:      config.DefaultRangeConcurrency * 2,
+		KVWriteBatchCount:      config.KVWriteBatchCount,
 		KVWriteBatchSize:       config.KVWriteBatchSize,
 		RegionSplitBatchSize:   config.DefaultRegionSplitBatchSize,
 		RegionSplitConcurrency: runtime.GOMAXPROCS(0),