feat(local): add directory size support (#624)

* feat(local): add directory size support

* fix(local): fix and improve directory size calculation

* style(local): fix code style

* style(local): fix code style

* style(local): fix code style

* fix(local): refresh directory size when force refresh

Signed-off-by: 我怎么就不是一只猫呢? <[email protected]>

* fix:(local): Avoid traversing the parent's parent, which leads to an endless loop

Signed-off-by: 我怎么就不是一只猫呢? <[email protected]>

* fix(local:) refresh dir size only enabled

Signed-off-by: 我怎么就不是一只猫呢? <[email protected]>

* fix(local): logical error && add RecalculateDirSize && cleaner code for int64

* feat(local): add Benchmark for CalculateDirSize

* refactor(local): 优化移动中对于错误的判断。

---------

Signed-off-by: 我怎么就不是一只猫呢? <[email protected]>
Co-authored-by: 我怎么就不是一只猫呢? <[email protected]>

Author: LmanTW

drivers/local/benchmark_calculatedirsize_test.go

@@ -0,0 +1,92 @@
+package local
+
+// TestDirCalculateSize tests the directory size calculation
+// It should be run with the local driver enabled and directory size calculation set to true
+import (
+	"os"
+	"path/filepath"
+	"strconv"
+	"testing"
+
+	"github.com/OpenListTeam/OpenList/v4/internal/driver"
+)
+
+func generatedTestDir(dir string, dep, filecount int) {
+	if dep == 0 {
+		return
+	}
+	for i := 0; i < dep; i++ {
+		subDir := dir + "/dir" + strconv.Itoa(i)
+		os.Mkdir(subDir, 0755)
+		generatedTestDir(subDir, dep-1, filecount)
+		generatedFiles(subDir, filecount)
+	}
+}
+
+func generatedFiles(path string, count int) error {
+	for i := 0; i < count; i++ {
+		filePath := filepath.Join(path, "file"+strconv.Itoa(i)+".txt")
+		file, err := os.Create(filePath)
+		if err != nil {
+			return err
+		}
+		// 使用随机ascii字符填充文件
+		content := make([]byte, 1024) // 1KB file
+		for j := range content {
+			content[j] = byte('a' + j%26) // Fill with 'a' to 'z'
+		}
+		_, err = file.Write(content)
+		if err != nil {
+			return err
+		}
+		file.Close()
+	}
+	return nil
+}
+
+// performance tests for directory size calculation
+func BenchmarkCalculateDirSize(t *testing.B) {
+	// 初始化t的日志
+	t.Logf("Starting performance test for directory size calculation")
+	// 确保测试目录存在
+	if testing.Short() {
+		t.Skip("Skipping performance test in short mode")
+	}
+	// 创建tmp directory for testing
+	testTempDir := t.TempDir()
+	err := os.MkdirAll(testTempDir, 0755)
+	if err != nil {
+		t.Fatalf("Failed to create test directory: %v", err)
+	}
+	defer os.RemoveAll(testTempDir) // Clean up after test
+	// 构建一个深度为5，每层10个文件和10个目录的目录结构
+	generatedTestDir(testTempDir, 5, 10)
+	// Initialize the local driver with directory size calculation enabled
+	d := &Local{
+		directoryMap: DirectoryMap{
+			root: testTempDir,
+		},
+		Addition: Addition{
+			DirectorySize: true,
+			RootPath: driver.RootPath{
+				RootFolderPath: testTempDir,
+			},
+		},
+	}
+	//record the start time
+	t.StartTimer()
+	// Calculate the directory size
+	err = d.directoryMap.RecalculateDirSize()
+	if err != nil {
+		t.Fatalf("Failed to calculate directory size: %v", err)
+	}
+	//record the end time
+	t.StopTimer()
+	// Print the size and duration
+	node, ok := d.directoryMap.Get(d.directoryMap.root)
+	if !ok {
+		t.Fatalf("Failed to get root node from directory map")
+	}
+	t.Logf("Directory size: %d bytes", node.fileSum+node.directorySum)
+	t.Logf("Performance test completed successfully")
+}

drivers/local/driver.go

@@ -33,6 +33,9 @@ type Local struct {
 	Addition
 	mkdirPerm int32
 
+	// directory size data
+	directoryMap DirectoryMap
+
 	// zero means no limit
 	thumbConcurrency int
 	thumbTokenBucket TokenBucket
@@ -66,6 +69,15 @@ func (d *Local) Init(ctx context.Context) error {
 		}
 		d.Addition.RootFolderPath = abs
 	}
+	if d.DirectorySize {
+		d.directoryMap.root = d.GetRootPath()
+		_, err := d.directoryMap.CalculateDirSize(d.GetRootPath())
+		if err != nil {
+			return err
+		}
+	} else {
+		d.directoryMap.Clear()
+	}
 	if d.ThumbCacheFolder != "" && !utils.Exists(d.ThumbCacheFolder) {
 		err := os.MkdirAll(d.ThumbCacheFolder, os.FileMode(d.mkdirPerm))
 		if err != nil {
@@ -124,6 +136,9 @@ func (d *Local) GetAddition() driver.Additional {
 func (d *Local) List(ctx context.Context, dir model.Obj, args model.ListArgs) ([]model.Obj, error) {
 	fullPath := dir.GetPath()
 	rawFiles, err := readDir(fullPath)
+	if d.DirectorySize && args.Refresh {
+		d.directoryMap.RecalculateDirSize()
+	}
 	if err != nil {
 		return nil, err
 	}
@@ -147,7 +162,12 @@ func (d *Local) FileInfoToObj(ctx context.Context, f fs.FileInfo, reqPath string
 	}
 	isFolder := f.IsDir() || isSymlinkDir(f, fullPath)
 	var size int64
-	if !isFolder {
+	if isFolder {
+		node, ok := d.directoryMap.Get(filepath.Join(fullPath, f.Name()))
+		if ok {
+			size = node.fileSum + node.directorySum
+		}
+	} else {
 		size = f.Size()
 	}
 	var ctime time.Time
@@ -186,7 +206,12 @@ func (d *Local) Get(ctx context.Context, path string) (model.Obj, error) {
 	isFolder := f.IsDir() || isSymlinkDir(f, path)
 	size := f.Size()
 	if isFolder {
-		size = 0
+		node, ok := d.directoryMap.Get(path)
+		if ok {
+			size = node.fileSum + node.directorySum
+		}
+	} else {
+		size = f.Size()
 	}
 	var ctime time.Time
 	t, err := times.Stat(path)
@@ -271,22 +296,31 @@ func (d *Local) Move(ctx context.Context, srcObj, dstDir model.Obj) error {
 	if utils.IsSubPath(srcPath, dstPath) {
 		return fmt.Errorf("the destination folder is a subfolder of the source folder")
 	}
-	if err := os.Rename(srcPath, dstPath); err != nil && strings.Contains(err.Error(), "invalid cross-device link") {
-		// Handle cross-device file move in local driver
-		if err = d.Copy(ctx, srcObj, dstDir); err != nil {
-			return err
-		} else {
-			// Directly remove file without check recycle bin if successfully copied
-			if srcObj.IsDir() {
-				err = os.RemoveAll(srcObj.GetPath())
-			} else {
-				err = os.Remove(srcObj.GetPath())
-			}
+	err := os.Rename(srcPath, dstPath)
+	if err != nil && strings.Contains(err.Error(), "invalid cross-device link") {
+		// 跨设备移动，先复制再删除
+		if err := d.Copy(ctx, srcObj, dstDir); err != nil {
 			return err
 		}
-	} else {
-		return err
+		// 复制成功后直接删除源文件/文件夹
+		if srcObj.IsDir() {
+			return os.RemoveAll(srcObj.GetPath())
+		}
+		return os.Remove(srcObj.GetPath())
+	}
+	if err == nil {
+		srcParent := filepath.Dir(srcPath)
+		dstParent := filepath.Dir(dstPath)
+		if d.directoryMap.Has(srcParent) {
+			d.directoryMap.UpdateDirSize(srcParent)
+			d.directoryMap.UpdateDirParents(srcParent)
+		}
+		if d.directoryMap.Has(dstParent) {
+			d.directoryMap.UpdateDirSize(dstParent)
+			d.directoryMap.UpdateDirParents(dstParent)
+		}
 	}
+	return err
 }
 
 func (d *Local) Rename(ctx context.Context, srcObj model.Obj, newName string) error {
@@ -296,6 +330,14 @@ func (d *Local) Rename(ctx context.Context, srcObj model.Obj, newName string) er
 	if err != nil {
 		return err
 	}
+
+	if srcObj.IsDir() {
+		if d.directoryMap.Has(srcPath) {
+			d.directoryMap.DeleteDirNode(srcPath)
+			d.directoryMap.CalculateDirSize(dstPath)
+		}
+	}
+
 	return nil
 }
 
@@ -306,11 +348,21 @@ func (d *Local) Copy(_ context.Context, srcObj, dstDir model.Obj) error {
 		return fmt.Errorf("the destination folder is a subfolder of the source folder")
 	}
 	// Copy using otiai10/copy to perform more secure & efficient copy
-	return cp.Copy(srcPath, dstPath, cp.Options{
+	err := cp.Copy(srcPath, dstPath, cp.Options{
 		Sync:          true, // Sync file to disk after copy, may have performance penalty in filesystem such as ZFS
 		PreserveTimes: true,
 		PreserveOwner: true,
 	})
+	if err != nil {
+		return err
+	}
+
+	if d.directoryMap.Has(filepath.Dir(dstPath)) {
+		d.directoryMap.UpdateDirSize(filepath.Dir(dstPath))
+		d.directoryMap.UpdateDirParents(filepath.Dir(dstPath))
+	}
+
+	return nil
 }
 
 func (d *Local) Remove(ctx context.Context, obj model.Obj) error {
@@ -331,6 +383,19 @@ func (d *Local) Remove(ctx context.Context, obj model.Obj) error {
 	if err != nil {
 		return err
 	}
+	if obj.IsDir() {
+		if d.directoryMap.Has(obj.GetPath()) {
+			d.directoryMap.DeleteDirNode(obj.GetPath())
+			d.directoryMap.UpdateDirSize(filepath.Dir(obj.GetPath()))
+			d.directoryMap.UpdateDirParents(filepath.Dir(obj.GetPath()))
+		}
+	} else {
+		if d.directoryMap.Has(filepath.Dir(obj.GetPath())) {
+			d.directoryMap.UpdateDirSize(filepath.Dir(obj.GetPath()))
+			d.directoryMap.UpdateDirParents(filepath.Dir(obj.GetPath()))
+		}
+	}
+
 	return nil
 }
 
@@ -354,6 +419,11 @@ func (d *Local) Put(ctx context.Context, dstDir model.Obj, stream model.FileStre
 	if err != nil {
 		log.Errorf("[local] failed to change time of %s: %s", fullPath, err)
 	}
+	if d.directoryMap.Has(dstDir.GetPath()) {
+		d.directoryMap.UpdateDirSize(dstDir.GetPath())
+		d.directoryMap.UpdateDirParents(dstDir.GetPath())
+	}
+
 	return nil
 }
 

drivers/local/meta.go

@@ -7,6 +7,7 @@ import (
 
 type Addition struct {
 	driver.RootPath
+	DirectorySize    bool   `json:"directory_size" default:"false" help:"This might impact host performance"`
 	Thumbnail        bool   `json:"thumbnail" required:"true" help:"enable thumbnail"`
 	ThumbCacheFolder string `json:"thumb_cache_folder"`
 	ThumbConcurrency string `json:"thumb_concurrency" default:"16" required:"false" help:"Number of concurrent thumbnail generation goroutines. This controls how many thumbnails can be generated in parallel."`
@@ -27,6 +28,8 @@ var config = driver.Config{
 
 func init() {
 	op.RegisterDriver(func() driver.Driver {
-		return &Local{}
+		return &Local{
+			directoryMap: DirectoryMap{},
+		}
 	})
 }