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

Signed-off-by: ekexium <eke@fastmail.com>
tikv · Nov 24, 2022 · 92f0a82 · 92f0a82
1 parent e9db9e6
commit 92f0a82
Show file tree

Hide file tree

Showing 4 changed files with 7 additions and 12 deletions.
diff --git a/internal/unionstore/memdb.go b/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
-	db.vlog.memChangeHook = innerHook
+	atomic.StorePointer((*unsafe.Pointer)(unsafe.Pointer(&db.allocator.memChangeHook)), unsafe.Pointer(&innerHook))
+	atomic.StorePointer((*unsafe.Pointer)(unsafe.Pointer(&db.vlog.memChangeHook)), unsafe.Pointer(&innerHook))
 }
 
 // Mem returns the current memory footprint

diff --git a/internal/unionstore/memdb_arena.go b/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)()
 	}
 }
 

diff --git a/tikv/kv.go b/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
 	}

diff --git a/txnkv/transaction/txn.go b/txnkv/transaction/txn.go
@@ -80,9 +80,8 @@ type SchemaAmender interface {
 // TxnOptions indicates the option when beginning a transaction.
 // TxnOptions are set by the TxnOption values passed to Begin
 type TxnOptions struct {
-	TxnScope                  string
-	StartTS                   *uint64
-	MemoryFootprintChangeHook func(uint64)
+	TxnScope string
+	StartTS  *uint64
 }
 
 // KVTxn contains methods to interact with a TiKV transaction.