Skip to content

Commit 93c0621

Browse files
LmanTWdezhishen
LmanTW
and
dezhishen
authored
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: 我怎么就不是一只猫呢? <26274059+dezhishen@users.noreply.github.com> * fix:(local): Avoid traversing the parent's parent, which leads to an endless loop Signed-off-by: 我怎么就不是一只猫呢? <26274059+dezhishen@users.noreply.github.com> * fix(local:) refresh dir size only enabled Signed-off-by: 我怎么就不是一只猫呢? <26274059+dezhishen@users.noreply.github.com> * fix(local): logical error && add RecalculateDirSize && cleaner code for int64 * feat(local): add Benchmark for CalculateDirSize * refactor(local): 优化移动中对于错误的判断。 --------- Signed-off-by: 我怎么就不是一只猫呢? <26274059+dezhishen@users.noreply.github.com> Co-authored-by: 我怎么就不是一只猫呢? <26274059+dezhishen@users.noreply.github.com>
1 parent b9b8eed commit 93c0621

File tree

4 files changed

+434
-17
lines changed

4 files changed

+434
-17
lines changed

‎drivers/local/benchmark_calculatedirsize_test.go‎

Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,92 @@
1+
package local
2+
3+
// TestDirCalculateSize tests the directory size calculation
4+
// It should be run with the local driver enabled and directory size calculation set to true
5+
import (
6+
"os"
7+
"path/filepath"
8+
"strconv"
9+
"testing"
10+
11+
"github.com/OpenListTeam/OpenList/v4/internal/driver"
12+
)
13+
14+
func generatedTestDir(dir string, dep, filecount int) {
15+
if dep == 0 {
16+
return
17+
}
18+
for i := 0; i < dep; i++ {
19+
subDir := dir + "/dir" + strconv.Itoa(i)
20+
os.Mkdir(subDir, 0755)
21+
generatedTestDir(subDir, dep-1, filecount)
22+
generatedFiles(subDir, filecount)
23+
}
24+
}
25+
26+
func generatedFiles(path string, count int) error {
27+
for i := 0; i < count; i++ {
28+
filePath := filepath.Join(path, "file"+strconv.Itoa(i)+".txt")
29+
file, err := os.Create(filePath)
30+
if err != nil {
31+
return err
32+
}
33+
// 使用随机ascii字符填充文件
34+
content := make([]byte, 1024) // 1KB file
35+
for j := range content {
36+
content[j] = byte('a' + j%26) // Fill with 'a' to 'z'
37+
}
38+
_, err = file.Write(content)
39+
if err != nil {
40+
return err
41+
}
42+
file.Close()
43+
}
44+
return nil
45+
}
46+
47+
// performance tests for directory size calculation
48+
func BenchmarkCalculateDirSize(t *testing.B) {
49+
// 初始化t的日志
50+
t.Logf("Starting performance test for directory size calculation")
51+
// 确保测试目录存在
52+
if testing.Short() {
53+
t.Skip("Skipping performance test in short mode")
54+
}
55+
// 创建tmp directory for testing
56+
testTempDir := t.TempDir()
57+
err := os.MkdirAll(testTempDir, 0755)
58+
if err != nil {
59+
t.Fatalf("Failed to create test directory: %v", err)
60+
}
61+
defer os.RemoveAll(testTempDir) // Clean up after test
62+
// 构建一个深度为5，每层10个文件和10个目录的目录结构
63+
generatedTestDir(testTempDir, 5, 10)
64+
// Initialize the local driver with directory size calculation enabled
65+
d := &Local{
66+
directoryMap: DirectoryMap{
67+
root: testTempDir,
68+
},
69+
Addition: Addition{
70+
DirectorySize: true,
71+
RootPath: driver.RootPath{
72+
RootFolderPath: testTempDir,
73+
},
74+
},
75+
}
76+
//record the start time
77+
t.StartTimer()
78+
// Calculate the directory size
79+
err = d.directoryMap.RecalculateDirSize()
80+
if err != nil {
81+
t.Fatalf("Failed to calculate directory size: %v", err)
82+
}
83+
//record the end time
84+
t.StopTimer()
85+
// Print the size and duration
86+
node, ok := d.directoryMap.Get(d.directoryMap.root)
87+
if !ok {
88+
t.Fatalf("Failed to get root node from directory map")
89+
}
90+
t.Logf("Directory size: %d bytes", node.fileSum+node.directorySum)
91+
t.Logf("Performance test completed successfully")
92+
}

‎drivers/local/driver.go‎

Lines changed: 86 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,9 @@ type Local struct {
3333
Addition
3434
mkdirPerm int32
3535

36+
// directory size data
37+
directoryMap DirectoryMap
38+
3639
// zero means no limit
3740
thumbConcurrency int
3841
thumbTokenBucket TokenBucket
@@ -66,6 +69,15 @@ func (d *Local) Init(ctx context.Context) error {
6669
}
6770
d.Addition.RootFolderPath = abs
6871
}
72+
if d.DirectorySize {
73+
d.directoryMap.root = d.GetRootPath()
74+
_, err := d.directoryMap.CalculateDirSize(d.GetRootPath())
75+
if err != nil {
76+
return err
77+
}
78+
} else {
79+
d.directoryMap.Clear()
80+
}
6981
if d.ThumbCacheFolder != "" && !utils.Exists(d.ThumbCacheFolder) {
7082
err := os.MkdirAll(d.ThumbCacheFolder, os.FileMode(d.mkdirPerm))
7183
if err != nil {
@@ -124,6 +136,9 @@ func (d *Local) GetAddition() driver.Additional {
124136
func (d *Local) List(ctx context.Context, dir model.Obj, args model.ListArgs) ([]model.Obj, error) {
125137
fullPath := dir.GetPath()
126138
rawFiles, err := readDir(fullPath)
139+
if d.DirectorySize && args.Refresh {
140+
d.directoryMap.RecalculateDirSize()
141+
}
127142
if err != nil {
128143
return nil, err
129144
}
@@ -147,7 +162,12 @@ func (d *Local) FileInfoToObj(ctx context.Context, f fs.FileInfo, reqPath string
147162
}
148163
isFolder := f.IsDir() || isSymlinkDir(f, fullPath)
149164
var size int64
150-
if !isFolder {
165+
if isFolder {
166+
node, ok := d.directoryMap.Get(filepath.Join(fullPath, f.Name()))
167+
if ok {
168+
size = node.fileSum + node.directorySum
169+
}
170+
} else {
151171
size = f.Size()
152172
}
153173
var ctime time.Time
@@ -186,7 +206,12 @@ func (d *Local) Get(ctx context.Context, path string) (model.Obj, error) {
186206
isFolder := f.IsDir() || isSymlinkDir(f, path)
187207
size := f.Size()
188208
if isFolder {
189-
size = 0
209+
node, ok := d.directoryMap.Get(path)
210+
if ok {
211+
size = node.fileSum + node.directorySum
212+
}
213+
} else {
214+
size = f.Size()
190215
}
191216
var ctime time.Time
192217
t, err := times.Stat(path)
@@ -271,22 +296,31 @@ func (d *Local) Move(ctx context.Context, srcObj, dstDir model.Obj) error {
271296
if utils.IsSubPath(srcPath, dstPath) {
272297
return fmt.Errorf("the destination folder is a subfolder of the source folder")
273298
}
274-
if err := os.Rename(srcPath, dstPath); err != nil && strings.Contains(err.Error(), "invalid cross-device link") {
275-
// Handle cross-device file move in local driver
276-
if err = d.Copy(ctx, srcObj, dstDir); err != nil {
277-
return err
278-
} else {
279-
// Directly remove file without check recycle bin if successfully copied
280-
if srcObj.IsDir() {
281-
err = os.RemoveAll(srcObj.GetPath())
282-
} else {
283-
err = os.Remove(srcObj.GetPath())
284-
}
299+
err := os.Rename(srcPath, dstPath)
300+
if err != nil && strings.Contains(err.Error(), "invalid cross-device link") {
301+
// 跨设备移动，先复制再删除
302+
if err := d.Copy(ctx, srcObj, dstDir); err != nil {
285303
return err
286304
}
287-
} else {
288-
return err
305+
// 复制成功后直接删除源文件/文件夹
306+
if srcObj.IsDir() {
307+
return os.RemoveAll(srcObj.GetPath())
308+
}
309+
return os.Remove(srcObj.GetPath())
310+
}
311+
if err == nil {
312+
srcParent := filepath.Dir(srcPath)
313+
dstParent := filepath.Dir(dstPath)
314+
if d.directoryMap.Has(srcParent) {
315+
d.directoryMap.UpdateDirSize(srcParent)
316+
d.directoryMap.UpdateDirParents(srcParent)
317+
}
318+
if d.directoryMap.Has(dstParent) {
319+
d.directoryMap.UpdateDirSize(dstParent)
320+
d.directoryMap.UpdateDirParents(dstParent)
321+
}
289322
}
323+
return err
290324
}
291325

292326
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
296330
if err != nil {
297331
return err
298332
}
333+
334+
if srcObj.IsDir() {
335+
if d.directoryMap.Has(srcPath) {
336+
d.directoryMap.DeleteDirNode(srcPath)
337+
d.directoryMap.CalculateDirSize(dstPath)
338+
}
339+
}
340+
299341
return nil
300342
}
301343

@@ -306,11 +348,21 @@ func (d *Local) Copy(_ context.Context, srcObj, dstDir model.Obj) error {
306348
return fmt.Errorf("the destination folder is a subfolder of the source folder")
307349
}
308350
// Copy using otiai10/copy to perform more secure & efficient copy
309-
return cp.Copy(srcPath, dstPath, cp.Options{
351+
err := cp.Copy(srcPath, dstPath, cp.Options{
310352
Sync: true, // Sync file to disk after copy, may have performance penalty in filesystem such as ZFS
311353
PreserveTimes: true,
312354
PreserveOwner: true,
313355
})
356+
if err != nil {
357+
return err
358+
}
359+
360+
if d.directoryMap.Has(filepath.Dir(dstPath)) {
361+
d.directoryMap.UpdateDirSize(filepath.Dir(dstPath))
362+
d.directoryMap.UpdateDirParents(filepath.Dir(dstPath))
363+
}
364+
365+
return nil
314366
}
315367

316368
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 {
331383
if err != nil {
332384
return err
333385
}
386+
if obj.IsDir() {
387+
if d.directoryMap.Has(obj.GetPath()) {
388+
d.directoryMap.DeleteDirNode(obj.GetPath())
389+
d.directoryMap.UpdateDirSize(filepath.Dir(obj.GetPath()))
390+
d.directoryMap.UpdateDirParents(filepath.Dir(obj.GetPath()))
391+
}
392+
} else {
393+
if d.directoryMap.Has(filepath.Dir(obj.GetPath())) {
394+
d.directoryMap.UpdateDirSize(filepath.Dir(obj.GetPath()))
395+
d.directoryMap.UpdateDirParents(filepath.Dir(obj.GetPath()))
396+
}
397+
}
398+
334399
return nil
335400
}
336401

@@ -354,6 +419,11 @@ func (d *Local) Put(ctx context.Context, dstDir model.Obj, stream model.FileStre
354419
if err != nil {
355420
log.Errorf("[local] failed to change time of %s: %s", fullPath, err)
356421
}
422+
if d.directoryMap.Has(dstDir.GetPath()) {
423+
d.directoryMap.UpdateDirSize(dstDir.GetPath())
424+
d.directoryMap.UpdateDirParents(dstDir.GetPath())
425+
}
426+
357427
return nil
358428
}
359429

‎drivers/local/meta.go‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import (
77

88
type Addition struct {
99
driver.RootPath
10+
DirectorySize bool `json:"directory_size" default:"false" help:"This might impact host performance"`
1011
Thumbnail bool `json:"thumbnail" required:"true" help:"enable thumbnail"`
1112
ThumbCacheFolder string `json:"thumb_cache_folder"`
1213
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{
2728

2829
func init() {
2930
op.RegisterDriver(func() driver.Driver {
30-
return &Local{}
31+
return &Local{
32+
directoryMap: DirectoryMap{},
33+
}
3134
})
3235
}

0 commit comments

Comments
 (0)