fix: avoid data race when setting memdb footprint hook (#621)

Signed-off-by: ekexium <eke@fastmail.com>

internal/unionstore/memdb.go

@@ -39,6 +39,7 @@ import (
 	"math"
 	"reflect"
 	"sync"
+	"sync/atomic"
 	"unsafe"
 
 	tikverr "github.com/tikv/client-go/v2/error"
@@ -871,8 +872,8 @@ func (db *MemDB) SetMemoryFootprintChangeHook(hook func(uint64)) {
 	innerHook := func() {
 		hook(db.allocator.capacity + db.vlog.capacity)
 	}
-	db.allocator.memChangeHook = innerHook
+	atomic.StorePointer((*unsafe.Pointer)(unsafe.Pointer(&db.allocator.memChangeHook)), unsafe.Pointer(&innerHook))
-	db.vlog.memChangeHook = innerHook
+	atomic.StorePointer((*unsafe.Pointer)(unsafe.Pointer(&db.vlog.memChangeHook)), unsafe.Pointer(&innerHook))
 }
 
 // Mem returns the current memory footprint

internal/unionstore/memdb_arena.go

@@ -94,7 +94,7 @@ type memdbArena struct {
 	// the total size of all blocks, also the approximate memory footprint of the arena.
 	capacity uint64
 	// when it enlarges or shrinks, call this function with the current memory footprint (in bytes)
-	memChangeHook func()
+	memChangeHook *func()
 }
 
 func (a *memdbArena) alloc(size int, align bool) (memdbArenaAddr, []byte) {
@@ -133,7 +133,7 @@ func (a *memdbArena) enlarge(allocSize, blockSize int) {
 
 func (a *memdbArena) onMemChange() {
 	if a.memChangeHook != nil {
-		a.memChangeHook()
+		(*a.memChangeHook)()
 	}
 }
 

tikv/kv.go

@@ -257,11 +257,6 @@ func (s *KVStore) Begin(opts ...TxnOption) (txn *transaction.KVTxn, err error) {
 		opt(options)
 	}
 
-	defer func() {
-		if err == nil && txn != nil && options.MemoryFootprintChangeHook != nil {
-			txn.SetMemoryFootprintChangeHook(options.MemoryFootprintChangeHook)
-		}
-	}()
 	if options.TxnScope == "" {
 		options.TxnScope = oracle.GlobalTxnScope
 	}

txnkv/transaction/txn.go

@@ -82,7 +82,6 @@ type SchemaAmender interface {
 type TxnOptions struct {
 	TxnScope string
 	StartTS  *uint64
-	MemoryFootprintChangeHook func(uint64)
 }
 
 // KVTxn contains methods to interact with a TiKV transaction.