[MemIAVL] Create snapshot whenever height diff exceed interval (#109)

* Snapshot should be created whenever height diff exceed interval

* Fix unit test

* Fix unit test

* Fix lint

* Add log for closing db

Author: yzang2019

sc/memiavl/db.go

@@ -548,12 +548,19 @@ func (db *DB) reloadMultiTree(mtree *MultiTree) error {
 
 // rewriteIfApplicable execute the snapshot rewrite strategy according to current height
 func (db *DB) rewriteIfApplicable(height int64) {
-	if height%int64(db.snapshotInterval) != 0 {
+	if db.snapshotRewriteChan != nil {
+		return
+	}
+
+	if db.snapshotInterval <= 0 || height <= 0 || height < db.MultiTree.SnapshotVersion() {
 		return
 	}
 
-	if err := db.rewriteSnapshotBackground(); err != nil {
-		db.logger.Error("failed to rewrite snapshot in background", "err", err)
+	// create snapshot when current height - last snapshot height > interval
+	if height-db.MultiTree.SnapshotVersion() >= int64(db.snapshotInterval) {
+		if err := db.rewriteSnapshotBackground(); err != nil {
+			db.logger.Error("failed to rewrite snapshot in background", "err", err)
+		}
 	}
 }
 
@@ -616,19 +623,22 @@ func (db *DB) rewriteSnapshotBackground() error {
 }
 
 func (db *DB) Close() error {
+	db.logger.Info("Closing memiavl db...")
 	db.mtx.Lock()
 	defer db.mtx.Unlock()
 	errs := []error{}
 	db.pruneSnapshotLock.Lock()
 	defer db.pruneSnapshotLock.Unlock()
 	// Close stream handler
+	db.logger.Info("Closing stream handler...")
 	if db.streamHandler != nil {
 		err := db.streamHandler.Close()
 		errs = append(errs, err)
 		db.streamHandler = nil
 	}
 
 	// Close rewrite channel
+	db.logger.Info("Closing rewrite channel...")
 	if db.snapshotRewriteChan != nil {
 		db.snapshotRewriteCancelFunc()
 		<-db.snapshotRewriteChan
@@ -639,12 +649,13 @@ func (db *DB) Close() error {
 	errs = append(errs, db.MultiTree.Close())
 
 	// Close file lock
+	db.logger.Info("Closing file lock...")
 	if db.fileLock != nil {
 		errs = append(errs, db.fileLock.Unlock())
 		errs = append(errs, db.fileLock.Destroy())
 		db.fileLock = nil
 	}
-
+	db.logger.Info("Closed memiavl db.")
 	return errorutils.Join(errs...)
 }
 

sc/memiavl/db_test.go

@@ -137,6 +137,61 @@ func TestRewriteSnapshotBackground(t *testing.T) {
 	stopped = true
 }
 
+// helper to commit one change to bump height
+func RequireCommitWithNoError(t *testing.T, db *DB, key, val string) int64 {
+	pairs := []*iavl.KVPair{{Key: []byte(key), Value: []byte(val)}}
+	cs := []*proto.NamedChangeSet{
+		{Name: "test", Changeset: iavl.ChangeSet{Pairs: pairs}},
+	}
+	require.NoError(t, db.ApplyChangeSets(cs))
+	v, err := db.Commit()
+	require.NoError(t, err)
+	return v
+}
+
+// Ensures snapshot rewrite is triggered when current height minus last snapshot height
+// exceeds the configured snapshot interval (strictly greater).
+func TestSnapshotTriggerOnIntervalDiff(t *testing.T) {
+	dir := t.TempDir()
+	db, err := OpenDB(logger.NewNopLogger(), 0, Options{
+		Dir:                dir,
+		CreateIfMissing:    true,
+		InitialStores:      []string{"test"},
+		SnapshotInterval:   5,
+		SnapshotKeepRecent: 0,
+	})
+	require.NoError(t, err)
+
+	// Heights 1..4 should NOT trigger because diff<=interval
+	for i := 1; i < 5; i++ {
+		v := RequireCommitWithNoError(t, db, "k"+strconv.Itoa(i), "v")
+		require.EqualValues(t, i, v)
+		// allow any background processing
+		time.Sleep(10 * time.Millisecond)
+		require.Nil(t, db.snapshotRewriteChan, "rewrite should not start at height %d", i)
+		// snapshot version should remain 0 until rewrite
+		require.EqualValues(t, 0, db.MultiTree.SnapshotVersion())
+	}
+
+	// Height 5 should trigger rewrite
+	v := RequireCommitWithNoError(t, db, "k6", "v")
+	require.Equal(t, int64(5), v)
+
+	// wait briefly for background rewrite to start
+	require.Eventually(t, func() bool {
+		return db.snapshotRewriteChan != nil
+	}, 3*time.Second, 100*time.Millisecond)
+	require.Eventually(t, func() bool {
+		require.NoError(t, db.checkAsyncTasks())
+		return db.snapshotRewriteChan == nil
+	}, 5*time.Second, 100*time.Millisecond)
+
+	// After completion, snapshot version should be 6
+	require.EqualValues(t, 5, db.MultiTree.SnapshotVersion())
+
+	require.NoError(t, db.Close())
+}
+
 func TestRlog(t *testing.T) {
 	dir := t.TempDir()
 	db, err := OpenDB(logger.NewNopLogger(), 0, Options{