Feat: use transactions to manipulate user's used storage

pkg/filesystem/file.go

@@ -70,7 +70,7 @@ func (fs *FileSystem) AddFile(ctx context.Context, parent *model.Folder, file fs
 		newFile.PicInfo = "1,1"
 	}
 
-	_, err = newFile.Create()
+	err = newFile.Create()
 
 	if err != nil {
 		if err := fs.Trigger(ctx, "AfterValidateFailed", file); err != nil {
@@ -79,6 +79,7 @@ func (fs *FileSystem) AddFile(ctx context.Context, parent *model.Folder, file fs
 		return nil, ErrFileExisted.WithError(err)
 	}
 
+	fs.User.Storage += newFile.Size
 	return &newFile, nil
 }
 

pkg/filesystem/fsctx/context.go

@@ -33,8 +33,6 @@ const (
 	ForceUsePublicEndpointCtx
 	// CancelFuncCtx Context 取消函數
 	CancelFuncCtx
-	// ValidateCapacityOnceCtx 限定归还容量的操作只執行一次
-	ValidateCapacityOnceCtx
 	// 文件在从机节点中的路径
 	SlaveSrcPath
 )

pkg/filesystem/hooks.go

@@ -2,11 +2,6 @@ package filesystem
 
 import (
 	"context"
-	"errors"
-	"io/ioutil"
-	"strings"
-	"sync"
-
 	model "github.com/cloudreve/Cloudreve/v3/models"
 	"github.com/cloudreve/Cloudreve/v3/pkg/cache"
 	"github.com/cloudreve/Cloudreve/v3/pkg/conf"
@@ -15,6 +10,8 @@ import (
 	"github.com/cloudreve/Cloudreve/v3/pkg/request"
 	"github.com/cloudreve/Cloudreve/v3/pkg/serializer"
 	"github.com/cloudreve/Cloudreve/v3/pkg/util"
+	"io/ioutil"
+	"strings"
 )
 
 // Hook 钩子函数
@@ -115,17 +112,8 @@ func HookResetPolicy(ctx context.Context, fs *FileSystem, file fsctx.FileHeader)
 	return fs.DispatchHandler()
 }
 
-// HookValidateCapacity 验证并扣除用户容量，包含数据库操作
+// HookValidateCapacity 验证用户容量
 func HookValidateCapacity(ctx context.Context, fs *FileSystem, file fsctx.FileHeader) error {
-	// 验证并扣除容量
-	if !fs.ValidateCapacity(ctx, file.Info().Size) {
-		return ErrInsufficientCapacity
-	}
-	return nil
-}
-
-// HookValidateCapacityWithoutIncrease 验证用户容量，不扣除
-func HookValidateCapacityWithoutIncrease(ctx context.Context, fs *FileSystem, file fsctx.FileHeader) error {
 	// 验证并扣除容量
 	if fs.User.GetRemainingCapacity() < file.Info().Size {
 		return ErrInsufficientCapacity
@@ -139,7 +127,7 @@ func HookValidateCapacityDiff(ctx context.Context, fs *FileSystem, newFile fsctx
 	newFileSize := newFile.Info().Size
 
 	if newFileSize > originFile.Size {
-		return HookValidateCapacityWithoutIncrease(ctx, fs, newFile)
+		return HookValidateCapacity(ctx, fs, newFile)
 	}
 
 	return nil
@@ -184,25 +172,6 @@ func HookCancelContext(ctx context.Context, fs *FileSystem, file fsctx.FileHeade
 	return nil
 }
 
-// HookGiveBackCapacity 归还用户容量
-func HookGiveBackCapacity(ctx context.Context, fs *FileSystem, file fsctx.FileHeader) error {
-	once, ok := ctx.Value(fsctx.ValidateCapacityOnceCtx).(*sync.Once)
-	if !ok {
-		once = &sync.Once{}
-	}
-
-	// 归还用户容量
-	res := true
-	once.Do(func() {
-		res = fs.User.DeductionStorage(file.Info().Size)
-	})
-
-	if !res {
-		return errors.New("无法继续降低用户已用存储")
-	}
-	return nil
-}
-
 // HookUpdateSourceName 更新文件SourceName
 // TODO：测试
 func HookUpdateSourceName(ctx context.Context, fs *FileSystem, file fsctx.FileHeader) error {
@@ -335,6 +304,14 @@ func HookChunkUploaded(ctx context.Context, fs *FileSystem, fileHeader fsctx.Fil
 	return fileInfo.Model.(*model.File).UpdateSize(fileInfo.Model.(*model.File).GetSize() + fileInfo.Size)
 }
 
+// HookChunkUploadFailed 单个分片上传失败后
+func HookChunkUploadFailed(ctx context.Context, fs *FileSystem, fileHeader fsctx.FileHeader) error {
+	fileInfo := fileHeader.Info()
+
+	// 更新文件大小
+	return fileInfo.Model.(*model.File).UpdateSize(fileInfo.Model.(*model.File).GetSize() - fileInfo.Size)
+}
+
 // HookChunkUploadFinished 分片上传结束后处理文件
 func HookChunkUploadFinished(ctx context.Context, fs *FileSystem, fileHeader fsctx.FileHeader) error {
 	fileInfo := fileHeader.Info()

pkg/filesystem/hooks_test.go

@@ -723,14 +723,14 @@ func TestHookValidateCapacityWithoutIncrease(t *testing.T) {
 	// not enough
 	{
 		fs.User.Group.MaxStorage = 10
-		a.Error(HookValidateCapacityWithoutIncrease(ctx, fs))
+		a.Error(HookValidateCapacity(ctx, fs))
 		a.EqualValues(10, fs.User.Storage)
 	}
 
 	// enough
 	{
 		fs.User.Group.MaxStorage = 11
-		a.NoError(HookValidateCapacityWithoutIncrease(ctx, fs))
+		a.NoError(HookValidateCapacity(ctx, fs))
 		a.EqualValues(10, fs.User.Storage)
 	}
 }

pkg/filesystem/upload.go

@@ -164,6 +164,23 @@ func (fs *FileSystem) CreateUploadSession(ctx context.Context, file *fsctx.FileS
 	callbackKey := uuid.Must(uuid.NewV4()).String()
 	fileSize := file.Size
 
+	// 创建占位的文件，同时校验文件信息
+	file.Mode = fsctx.Nop
+	if callbackKey != "" {
+		file.UploadSessionID = &callbackKey
+	}
+
+	fs.Use("BeforeUpload", HookValidateFile)
+	fs.Use("BeforeUpload", HookValidateCapacity)
+	if !fs.Policy.IsUploadPlaceholderWithSize() {
+		fs.Use("AfterUpload", HookClearFileHeaderSize)
+	}
+
+	fs.Use("AfterUpload", GenericAfterUpload)
+	if err := fs.Upload(ctx, file); err != nil {
+		return nil, err
+	}
+
 	uploadSession := &serializer.UploadSession{
 		Key:          callbackKey,
 		UID:          fs.User.ID,
@@ -181,27 +198,6 @@ func (fs *FileSystem) CreateUploadSession(ctx context.Context, file *fsctx.FileS
 		return nil, err
 	}
 
-	// 创建占位的文件，同时校验文件信息
-	file.Mode = fsctx.Nop
-	if callbackKey != "" {
-		file.UploadSessionID = &callbackKey
-	}
-
-	fs.Use("BeforeUpload", HookValidateFile)
-	if !fs.Policy.IsUploadPlaceholderWithSize() {
-		fs.Use("BeforeUpload", HookValidateCapacityWithoutIncrease)
-		fs.Use("AfterUpload", HookClearFileHeaderSize)
-	} else {
-		fs.Use("BeforeUpload", HookValidateCapacity)
-		fs.Use("AfterValidateFailed", HookGiveBackCapacity)
-		fs.Use("AfterUploadFailed", HookGiveBackCapacity)
-	}
-
-	fs.Use("AfterUpload", GenericAfterUpload)
-	if err := fs.Upload(ctx, file); err != nil {
-		return nil, err
-	}
-
 	// 创建回调会话
 	err = cache.Set(
 		UploadSessionCachePrefix+callbackKey,
@@ -226,12 +222,9 @@ func (fs *FileSystem) UploadFromStream(ctx context.Context, file *fsctx.FileStre
 		fs.Use("BeforeUpload", HookValidateFile)
 		fs.Use("BeforeUpload", HookValidateCapacity)
 		fs.Use("AfterUploadCanceled", HookDeleteTempFile)
-		fs.Use("AfterUploadCanceled", HookGiveBackCapacity)
 		fs.Use("AfterUpload", GenericAfterUpload)
 		fs.Use("AfterUpload", HookGenerateThumb)
 		fs.Use("AfterValidateFailed", HookDeleteTempFile)
-		fs.Use("AfterValidateFailed", HookGiveBackCapacity)
-		fs.Use("AfterUploadFailed", HookGiveBackCapacity)
 	}
 	fs.Lock.Unlock()