Clean tmp disk data when start milvus (#24404)

Signed-off-by: xige-16 <xi.ge@zilliz.com>

cmd/roles/roles.go

@@ -21,6 +21,7 @@ import (
 	"fmt"
 	"os"
 	"os/signal"
+	"path/filepath"
 	"strings"
 	"sync"
 	"syscall"
@@ -65,6 +66,24 @@ type component interface {
 	Stop() error
 }
 
+func cleanLocalDir(path string) {
+	_, statErr := os.Stat(path)
+	// path exist, but stat error
+	if statErr != nil && !os.IsNotExist(statErr) {
+		log.Warn("Check if path exists failed when clean local data cache", zap.Error(statErr))
+		panic(statErr)
+	}
+	// path exist, remove all
+	if statErr == nil {
+		err := os.RemoveAll(path)
+		if err != nil {
+			log.Warn("Clean local data cache failed", zap.Error(err))
+			panic(err)
+		}
+		log.Info("Clean local data cache", zap.String("path", path))
+	}
+}
+
 func runComponent[T component](ctx context.Context,
 	localMsg bool,
 	runWg *sync.WaitGroup,
@@ -141,6 +160,10 @@ func (mr *MilvusRoles) runQueryCoord(ctx context.Context, localMsg bool, wg *syn
 
 func (mr *MilvusRoles) runQueryNode(ctx context.Context, localMsg bool, wg *sync.WaitGroup) *components.QueryNode {
 	wg.Add(1)
+	rootPath := paramtable.Get().LocalStorageCfg.Path.GetValue()
+	queryDataLocalPath := filepath.Join(rootPath, typeutil.QueryNodeRole)
+	cleanLocalDir(queryDataLocalPath)
+
 	return runComponent(ctx, localMsg, wg, components.NewQueryNode, metrics.RegisterQueryNode)
 }
 
@@ -161,6 +184,10 @@ func (mr *MilvusRoles) runIndexCoord(ctx context.Context, localMsg bool, wg *syn
 
 func (mr *MilvusRoles) runIndexNode(ctx context.Context, localMsg bool, wg *sync.WaitGroup) *components.IndexNode {
 	wg.Add(1)
+	rootPath := paramtable.Get().LocalStorageCfg.Path.GetValue()
+	indexDataLocalPath := filepath.Join(rootPath, typeutil.IndexNodeRole)
+	cleanLocalDir(indexDataLocalPath)
+
 	return runComponent(ctx, localMsg, wg, components.NewIndexNode, metrics.RegisterIndexNode)
 }
 

cmd/roles/roles_test.go

@@ -17,10 +17,14 @@
 package roles
 
 import (
+	"os"
+	"path/filepath"
 	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+
+	"github.com/milvus-io/milvus/pkg/util/paramtable"
 )
 
 func TestRoles(t *testing.T) {
@@ -41,4 +45,52 @@ func TestRoles(t *testing.T) {
 	assert.Equal(t, len(ss), 1)
 	ss = strings.SplitN("adb=def", "=", 2)
 	assert.Equal(t, len(ss), 2)
+
+	paramtable.Init()
+	rootPath := paramtable.Get().LocalStorageCfg.Path.GetValue()
+	localPath := filepath.Join(rootPath, "test-dir")
+
+	err := os.RemoveAll(localPath)
+	assert.NoError(t, err)
+
+	err = os.MkdirAll(localPath, os.ModeDir)
+	assert.NoError(t, err)
+	_, err = os.Create(filepath.Join(localPath, "child"))
+	assert.NoError(t, err)
+
+	err = os.RemoveAll(localPath)
+	assert.NoError(t, err)
+	_, err = os.Stat(localPath)
+	assert.Error(t, err)
+	assert.Equal(t, true, os.IsNotExist(err))
+}
+
+func TestCleanLocalDir(t *testing.T) {
+	paramtable.Init()
+	rootPath := paramtable.Get().LocalStorageCfg.Path.GetValue()
+	localPath := filepath.Join(rootPath, "test-dir")
+
+	// clean data
+	assert.NotPanics(t, func() {
+		cleanLocalDir(localPath)
+	})
+
+	// create dir and file
+	err := os.MkdirAll(localPath, os.ModeDir)
+	assert.NoError(t, err)
+	_, err = os.Create(filepath.Join(localPath, "child"))
+	assert.NoError(t, err)
+
+	// clean with path exist
+	assert.NotPanics(t, func() {
+		cleanLocalDir(localPath)
+	})
+
+	_, err = os.Stat(localPath)
+	assert.Error(t, err)
+	assert.Equal(t, true, os.IsNotExist(err))
+	// clean with path not exist
+	assert.NotPanics(t, func() {
+		cleanLocalDir(localPath)
+	})
 }

internal/indexnode/indexnode.go

@@ -31,6 +31,7 @@ import (
 	"math/rand"
 	"os"
 	"path"
+	"path/filepath"
 	"sync"
 	"syscall"
 	"time"
@@ -157,7 +158,8 @@ func (i *IndexNode) initKnowhere() {
 	cCPUNum := C.int(hardware.GetCPUNum())
 	C.InitCpuNum(cCPUNum)
 
-	initcore.InitLocalStorageConfig(Params)
+	localDataRootPath := filepath.Join(Params.LocalStorageCfg.Path.GetValue(), typeutil.IndexNodeRole)
+	initcore.InitLocalStorageConfig(localDataRootPath)
 }
 
 func (i *IndexNode) initSession() error {

internal/querynodev2/server.go

@@ -32,6 +32,7 @@ import (
 	"fmt"
 	"os"
 	"path"
+	"path/filepath"
 	"plugin"
 	"runtime/debug"
 	"sync"
@@ -211,7 +212,8 @@ func (node *QueryNode) InitSegcore() {
 	cCPUNum := C.int(hardware.GetCPUNum())
 	C.InitCpuNum(cCPUNum)
 
-	initcore.InitLocalStorageConfig(paramtable.Get())
+	localDataRootPath := filepath.Join(paramtable.Get().LocalStorageCfg.Path.GetValue(), typeutil.QueryNodeRole)
+	initcore.InitLocalStorageConfig(localDataRootPath)
 
 	mmapDirPath := paramtable.Get().QueryNodeCfg.MmapDirPath.GetValue()
 	if len(mmapDirPath) > 0 {

internal/util/initcore/init_core.go

@@ -26,17 +26,13 @@ package initcore
 import "C"
 
 import (
-	"os"
-	"path/filepath"
 	"unsafe"
 
 	"github.com/milvus-io/milvus/pkg/util/paramtable"
 )
 
-func InitLocalStorageConfig(params *paramtable.ComponentParam) {
-	b, _ := os.Getwd()
-	LocalRootPath := filepath.Dir(b) + "/" + filepath.Base(b) + "/" + "data/"
-	CLocalRootPath := C.CString(LocalRootPath)
+func InitLocalStorageConfig(path string) {
+	CLocalRootPath := C.CString(path)
 	C.InitLocalRootPath(CLocalRootPath)
 	C.free(unsafe.Pointer(CLocalRootPath))
 }