Problem: memiavl leave tmp directories behind

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

@@ -203,9 +203,56 @@ func Load(dir string, opts Options) (*DB, error) {
 		}
 	}
 
+	if err := removeAllTmpDirs(db.dir); err != nil {
+		db.logger.Error("failed to remove tmp directories", "err", err)
+	}
+
 	return db, nil
 }
 
+func removeAllTmpDirs(rootDir string) error {
+	const maxRemovals = 1000
+	var wg sync.WaitGroup
+	var mu sync.Mutex
+	var errors []error
+	sem := make(chan struct{}, maxRemovals)
+
+	err := filepath.Walk(rootDir, func(path string, info os.FileInfo, err error) error {
+		if err != nil {
+			return err
+		}
+		if info.IsDir() && strings.HasSuffix(path, "-tmp") {
+			sem <- struct{}{}
+			wg.Add(1)
+			go func() {
+				defer func() {
+					<-sem
+					wg.Done()
+				}()
+				if err := os.RemoveAll(path); err != nil {
+					mu.Lock()
+					errors = append(errors, err)
+					mu.Unlock()
+				}
+			}()
+		}
+		return nil
+	})
+	if err != nil {
+		return err
+	}
+
+	wg.Wait()
+	if len(errors) > 0 {
+		var errMsgs []string
+		for _, e := range errors {
+			errMsgs = append(errMsgs, e.Error())
+		}
+		return fmt.Errorf("failed to remove some tmp directories: %s", strings.Join(errMsgs, "; "))
+	}
+	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,6 +511,9 @@ func (db *DB) RewriteSnapshot() error {
 		return err
 	}
 	if err := db.MultiTree.WriteSnapshot(path); err != nil {
+		if err := os.RemoveAll(path); err != nil {
+			db.logger.Error("failed to remove tmp directories", "err", err)
+		}
 		return err
 	}
 	if err := os.Rename(path, filepath.Join(db.dir, snapshotDir)); err != nil {

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,