fix: avoid data race when setting memdb footprint hook

Signed-off-by: ekexium <eke@fastmail.com>
tikv · Nov 23, 2022 · 3bdebcf · 3bdebcf
1 parent e9db9e6
commit 3bdebcf
Showing 1 changed file with 7 additions and 4 deletions.
diff --git a/internal/unionstore/memdb.go b/internal/unionstore/memdb.go
@@ -868,11 +868,14 @@ func (db *MemDB) RemoveFromBuffer(key []byte) {
 
 // SetMemoryFootprintChangeHook sets the hook function that is triggered when memdb grows.
 func (db *MemDB) SetMemoryFootprintChangeHook(hook func(uint64)) {
-	innerHook := func() {
-		hook(db.allocator.capacity + db.vlog.capacity)
+	if db.allocator.memChangeHook == nil {
+		// always set together
+		innerHook := func() {
+			hook(db.allocator.capacity + db.vlog.capacity)
+		}
+		db.allocator.memChangeHook = innerHook
+		db.vlog.memChangeHook = innerHook
 	}
-	db.allocator.memChangeHook = innerHook
-	db.vlog.memChangeHook = innerHook
 }
 
 // Mem returns the current memory footprint