don't call onMemChange before all work is done in memdb (#636)

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

internal/unionstore/memdb.go

@@ -36,6 +36,7 @@ package unionstore
 
 import (
 	"bytes"
+	"fmt"
 	"math"
 	"reflect"
 	"sync"
@@ -142,7 +143,7 @@ func (db *MemDB) Cleanup(h int) {
 	if h < len(db.stages) {
 		// This should never happens in production environment.
 		// Use panic to make debug easier.
-		panic("cannot cleanup staging buffer")
+		panic(fmt.Sprintf("cannot cleanup staging buffer, h=%v, len(db.stages)=%v", h, len(db.stages)))
 	}
 
 	cp := &db.stages[h-1]
@@ -154,6 +155,7 @@ func (db *MemDB) Cleanup(h int) {
 		}
 	}
 	db.stages = db.stages[:h-1]
+	db.vlog.onMemChange()
 }
 
 // Checkpoint returns a checkpoint of MemDB.
@@ -166,6 +168,7 @@ func (db *MemDB) Checkpoint() *MemDBCheckpoint {
 func (db *MemDB) RevertToCheckpoint(cp *MemDBCheckpoint) {
 	db.vlog.revertToCheckpoint(db, cp)
 	db.vlog.truncate(cp)
+	db.vlog.onMemChange()
 }
 
 // Reset resets the MemBuffer to initial states.

internal/unionstore/memdb_arena.go

@@ -133,6 +133,8 @@ func (a *memdbArena) enlarge(allocSize, blockSize int) {
 	// for some operations (e.g. revertToCheckpoint)
 }
 
+// onMemChange should only be called right before exiting memdb.
+// This is because the hook can lead to a panic, and leave memdb in an inconsistent state.
 func (a *memdbArena) onMemChange() {
 	hook := a.memChangeHook.Load()
 	if hook != nil {
@@ -220,7 +222,7 @@ func (a *memdbArena) truncate(snap *MemDBCheckpoint) {
 	for _, block := range a.blocks {
 		a.capacity += uint64(block.length)
 	}
-	a.onMemChange()
+	// We shall not call a.onMemChange() here, since it may cause a panic and leave memdb in an inconsistent state
 }
 
 type nodeAllocator struct {