diff --git a/internal/unionstore/memdb.go b/internal/unionstore/memdb.go
index 621e41bfe..194b8864b 100644
--- a/internal/unionstore/memdb.go
+++ b/internal/unionstore/memdb.go
@@ -111,14 +111,15 @@ func (db *MemDB) Staging() int {
 
 // Release publish all modifications in the latest staging buffer to upper level.
 func (db *MemDB) Release(h int) {
+	db.Lock()
+	defer db.Unlock()
+
 	if h != len(db.stages) {
 		// This should never happens in production environment.
 		// Use panic to make debug easier.
 		panic("cannot release staging buffer")
 	}
 
-	db.Lock()
-	defer db.Unlock()
 	if h == 1 {
 		tail := db.vlog.checkpoint()
 		if !db.stages[0].isSamePosition(&tail) {
@@ -131,6 +132,9 @@ func (db *MemDB) Release(h int) {
 // Cleanup cleanup the resources referenced by the StagingHandle.
 // If the changes are not published by `Release`, they will be discarded.
 func (db *MemDB) Cleanup(h int) {
+	db.Lock()
+	defer db.Unlock()
+
 	if h > len(db.stages) {
 		return
 	}
@@ -140,8 +144,6 @@ func (db *MemDB) Cleanup(h int) {
 		panic("cannot cleanup staging buffer")
 	}
 
-	db.Lock()
-	defer db.Unlock()
 	cp := &db.stages[h-1]
 	if !db.vlogInvalid {
 		curr := db.vlog.checkpoint()
@@ -294,6 +296,9 @@ func (db *MemDB) Dirty() bool {
 }
 
 func (db *MemDB) set(key []byte, value []byte, ops ...kv.FlagsOp) error {
+	db.Lock()
+	defer db.Unlock()
+
 	if db.vlogInvalid {
 		// panic for easier debugging.
 		panic("vlog is resetted")
@@ -308,9 +313,6 @@ func (db *MemDB) set(key []byte, value []byte, ops ...kv.FlagsOp) error {
 		}
 	}
 
-	db.Lock()
-	defer db.Unlock()
-
 	if len(db.stages) == 0 {
 		db.dirty = true
 	}