Problem: memiavl leave tmp directories behind

CHANGELOG.md

@@ -57,6 +57,7 @@
 - [#1067](https://github.com/crypto-org-chain/cronos/pull/1067) memiavl: only export state-sync snapshots on an exist snapshot
 - [#1082](https://github.com/crypto-org-chain/cronos/pull/1082) Make memiavl setup code reusable.
 - [#1092](https://github.com/crypto-org-chain/cronos/pull/1092) memiavl disable sdk address cache if zero-copy enabled, and disable zero-copy by default.
+- [#1099](https://github.com/crypto-org-chain/cronos/pull/1099) clean up memiavl tmp directories left behind.
 
 *April 13, 2023*
 

memiavl/README.md

@@ -116,8 +116,6 @@ IAVL snapshot is composed by four files:
   *repeat*
   ```
 
-- `kvs.index`, optional MPHF(Minimal-Perfect-Hash-Function) hash index build from `kvs`, support query key-values as a hash map.
-
 #### Compression
 
 The items in snapshot reference with each other by file offsets, we can apply some block compression techniques to compress keys and values files while maintain random accessbility by uncompressed file offset, for example zstd's experimental seekable format[^1].

memiavl/db.go

@@ -124,6 +124,11 @@ func Load(dir string, opts Options) (*DB, error) {
 		}
 	}
 
+	// cleanup any temporary directories left by interrupted snapshot rewrite
+	if err := removeTmpDirs(dir); err != nil {
+		return nil, fmt.Errorf("fail to cleanup tmp directories: %w", err)
+	}
+
 	wal, err := OpenWAL(walPath(dir), &wal.Options{NoCopy: true, NoSync: true})
 	if err != nil {
 		return nil, err
@@ -206,6 +211,25 @@ func Load(dir string, opts Options) (*DB, error) {
 	return db, nil
 }
 
+func removeTmpDirs(rootDir string) error {
+	entries, err := os.ReadDir(rootDir)
+	if err != nil {
+		return err
+	}
+
+	for _, entry := range entries {
+		if !entry.IsDir() || !strings.HasSuffix(entry.Name(), "-tmp") {
+			continue
+		}
+
+		if err := os.RemoveAll(filepath.Join(rootDir, entry.Name())); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
 // SetInitialVersion wraps `MultiTree.SetInitialVersion`.
 // it do an immediate snapshot rewrite, because we can't use wal log to record this change,
 // because we need it to convert versions to wal index in the first place.
@@ -464,7 +488,7 @@ func (db *DB) RewriteSnapshot() error {
 		return err
 	}
 	if err := db.MultiTree.WriteSnapshot(path); err != nil {
-		return err
+		return errors.Join(err, os.RemoveAll(path))
 	}
 	if err := os.Rename(path, filepath.Join(db.dir, snapshotDir)); err != nil {
 		return err

memiavl/db_test.go

@@ -4,6 +4,7 @@ import (
 	"encoding/hex"
 	"errors"
 	"os"
+	"path/filepath"
 	"runtime/debug"
 	"strconv"
 	"testing"
@@ -38,6 +39,27 @@ func TestRewriteSnapshot(t *testing.T) {
 	}
 }
 
+func TestRemoveSnapshotDir(t *testing.T) {
+	dbDir := t.TempDir()
+	defer os.RemoveAll(dbDir)
+
+	snapshotDir := filepath.Join(dbDir, snapshotName(0))
+	tmpDir := snapshotDir + "-tmp"
+	if err := os.MkdirAll(tmpDir, os.ModePerm); err != nil {
+		t.Fatalf("Failed to create dummy snapshot directory: %v", err)
+	}
+	_, err := Load(dbDir, Options{
+		CreateIfMissing:    true,
+		InitialStores:      []string{"test"},
+		SnapshotKeepRecent: 0,
+	})
+	require.NoError(t, err)
+
+	if _, err := os.Stat(tmpDir); !os.IsNotExist(err) {
+		t.Errorf("Expected temporary snapshot directory to be deleted, but it still exists")
+	}
+}
+
 func TestRewriteSnapshotBackground(t *testing.T) {
 	db, err := Load(t.TempDir(), Options{
 		CreateIfMissing:    true,