store/tikv: set memdb's size limits after create

store/driver/txn/txn_driver.go

@@ -15,6 +15,7 @@ package txn
 
 import (
 	"context"
+	"sync/atomic"
 
 	"github.com/pingcap/errors"
 	"github.com/pingcap/parser/model"
@@ -33,6 +34,11 @@ type tikvTxn struct {
 // NewTiKVTxn returns a new Transaction.
 func NewTiKVTxn(txn *tikv.KVTxn) kv.Transaction {
 	txn.SetOption(tikvstore.KVFilter, TiDBKVFilter{})
+
+	entryLimit := atomic.LoadUint64(&kv.TxnEntrySizeLimit)
+	totalLimit := atomic.LoadUint64(&kv.TxnTotalSizeLimit)
+	txn.GetUnionStore().SetEntrySizeLimit(entryLimit, totalLimit)
+
 	return &tikvTxn{txn, make(map[int64]*model.TableInfo)}
 }
 

store/tikv/2pc.go

@@ -29,7 +29,6 @@ import (
 	"github.com/pingcap/failpoint"
 	pb "github.com/pingcap/kvproto/pkg/kvrpcpb"
 	"github.com/pingcap/parser/terror"
-	tidbkv "github.com/pingcap/tidb/kv"
 	"github.com/pingcap/tidb/store/tikv/config"
 	"github.com/pingcap/tidb/store/tikv/kv"
 	"github.com/pingcap/tidb/store/tikv/logutil"
@@ -397,9 +396,6 @@ func (c *twoPhaseCommitter) initKeysAndMutations() error {
 	}
 	c.txnSize = size
 
-	if size > int(tidbkv.TxnTotalSizeLimit) {
-		return tidbkv.ErrTxnTooLarge.GenWithStackByArgs(size)
-	}
 	const logEntryCount = 10000
 	const logSize = 4 * 1024 * 1024 // 4MB
 	if c.mutations.Len() > logEntryCount || size > logSize {

store/tikv/unionstore/memdb.go

@@ -15,9 +15,9 @@ package unionstore
 
 import (
 	"bytes"
+	"math"
 	"reflect"
 	"sync"
-	"sync/atomic"
 	"unsafe"
 
 	tidbkv "github.com/pingcap/tidb/kv"
@@ -72,8 +72,8 @@ func newMemDB() *MemDB {
 	db.allocator.init()
 	db.root = nullAddr
 	db.stages = make([]memdbCheckpoint, 0, 2)
-	db.entrySizeLimit = atomic.LoadUint64(&tidbkv.TxnEntrySizeLimit)
-	db.bufferSizeLimit = atomic.LoadUint64(&tidbkv.TxnTotalSizeLimit)
+	db.entrySizeLimit = math.MaxUint64
+	db.bufferSizeLimit = math.MaxUint64
 	return db
 }
 

store/tikv/unionstore/memdb_bench_test.go

@@ -15,7 +15,6 @@ package unionstore
 
 import (
 	"encoding/binary"
-	"math"
 	"math/rand"
 	"testing"
 )
@@ -25,19 +24,12 @@ const (
 	valueSize = 128
 )
 
-func newMemDBForBench() *MemDB {
-	db := newMemDB()
-	db.bufferSizeLimit = math.MaxUint64
-	db.entrySizeLimit = math.MaxUint64
-	return db
-}
-
 func BenchmarkLargeIndex(b *testing.B) {
 	buf := make([][valueSize]byte, 10000000)
 	for i := range buf {
 		binary.LittleEndian.PutUint32(buf[i][:], uint32(i))
 	}
-	db := newMemDBForBench()
+	db := newMemDB()
 	b.ResetTimer()
 
 	for i := range buf {
@@ -51,7 +43,7 @@ func BenchmarkPut(b *testing.B) {
 		binary.BigEndian.PutUint32(buf[i][:], uint32(i))
 	}
 
-	p := newMemDBForBench()
+	p := newMemDB()
 	b.ResetTimer()
 
 	for i := range buf {
@@ -65,7 +57,7 @@ func BenchmarkPutRandom(b *testing.B) {
 		binary.LittleEndian.PutUint32(buf[i][:], uint32(rand.Int()))
 	}
 
-	p := newMemDBForBench()
+	p := newMemDB()
 	b.ResetTimer()
 
 	for i := range buf {
@@ -79,7 +71,7 @@ func BenchmarkGet(b *testing.B) {
 		binary.BigEndian.PutUint32(buf[i][:], uint32(i))
 	}
 
-	p := newMemDBForBench()
+	p := newMemDB()
 	for i := range buf {
 		p.Set(buf[i][:keySize], buf[i][:])
 	}
@@ -96,7 +88,7 @@ func BenchmarkGetRandom(b *testing.B) {
 		binary.LittleEndian.PutUint32(buf[i][:], uint32(rand.Int()))
 	}
 
-	p := newMemDBForBench()
+	p := newMemDB()
 	for i := range buf {
 		p.Set(buf[i][:keySize], buf[i][:])
 	}

store/tikv/unionstore/memdb_test.go

@@ -19,7 +19,6 @@ import (
 	"context"
 	"encoding/binary"
 	"fmt"
-	"math"
 	"testing"
 
 	. "github.com/pingcap/check"
@@ -71,8 +70,6 @@ func (s *testMemDBSuite) TestGetSet(c *C) {
 
 func (s *testMemDBSuite) TestBigKV(c *C) {
 	db := newMemDB()
-	db.entrySizeLimit = math.MaxUint64
-	db.bufferSizeLimit = math.MaxUint64
 	db.Set([]byte{1}, make([]byte, 80<<20))
 	c.Assert(db.vlog.blockSize, Equals, maxBlockSize)
 	c.Assert(len(db.vlog.blocks), Equals, 1)

store/tikv/unionstore/union_store.go

@@ -139,6 +139,12 @@ func (us *KVUnionStore) GetOption(opt int) interface{} {
 	return us.opts[opt]
 }
 
+// SetEntrySizeLimit sets the size limit for each entry and total buffer.
+func (us *KVUnionStore) SetEntrySizeLimit(entryLimit, bufferLimit uint64) {
+	us.memBuffer.entrySizeLimit = entryLimit
+	us.memBuffer.bufferSizeLimit = bufferLimit
+}
+
 type options map[int]interface{}
 
 func (opts options) Get(opt int) (interface{}, bool) {