From 5be7ec98c1ba9453ab5aa61d59c1b68a2526f46b Mon Sep 17 00:00:00 2001 From: HFO4 <912394456@qq.com> Date: Tue, 14 Jan 2020 10:32:54 +0800 Subject: [PATCH] Feat: user storage pack --- models/migration.go | 2 +- models/storage_pack.go | 61 ++++++++++++++++++++++++++++++++ models/storage_pack_test.go | 39 ++++++++++++++++++++ models/user.go | 5 +-- models/user_test.go | 12 +++++++ pkg/filesystem/file.go | 2 +- pkg/filesystem/hooks_test.go | 3 ++ pkg/filesystem/manage_test.go | 4 +++ pkg/filesystem/validator_test.go | 7 ++++ pkg/serializer/user.go | 7 ++-- pkg/serializer/user_test.go | 13 +++++++ 11 files changed, 148 insertions(+), 7 deletions(-) create mode 100644 models/storage_pack.go create mode 100644 models/storage_pack_test.go diff --git a/models/migration.go b/models/migration.go index a4bdc89..3051f70 100644 --- a/models/migration.go +++ b/models/migration.go @@ -29,7 +29,7 @@ func migration() { if conf.DatabaseConfig.Type == "mysql" { DB = DB.Set("gorm:table_options", "ENGINE=InnoDB") } - DB.AutoMigrate(&User{}, &Setting{}, &Group{}, &Policy{}, &Folder{}, &File{}) + DB.AutoMigrate(&User{}, &Setting{}, &Group{}, &Policy{}, &Folder{}, &File{}, &StoragePack{}) // 创建初始存储策略 addDefaultPolicy() diff --git a/models/storage_pack.go b/models/storage_pack.go new file mode 100644 index 0000000..89a1930 --- /dev/null +++ b/models/storage_pack.go @@ -0,0 +1,61 @@ +package model + +import ( + "github.com/HFO4/cloudreve/pkg/cache" + "github.com/jinzhu/gorm" + "strconv" + "time" +) + +// StoragePack 容量包模型 +type StoragePack struct { + // 表字段 + gorm.Model + Name string + UserID uint + ActiveTime *time.Time + ExpiredTime *time.Time `gorm:"index:expired"` + Size uint64 +} + +// GetAvailablePackSize 返回给定用户当前可用的容量包总容量 +func (user *User) GetAvailablePackSize() uint64 { + var ( + packs []StoragePack + total uint64 + firstExpire *time.Time + timeNow = time.Now() + ttl int64 + ) + + // 尝试从缓存中读取 + cacheKey := "pack_size_" + strconv.FormatUint(uint64(user.ID), 10) + if total, ok := cache.Get(cacheKey); ok { + return total.(uint64) + } + + // 查找所有有效容量包 + DB.Where("expired_time > ? AND user_id = ?", timeNow, user.ID).Find(&packs) + + // 计算总容量, 并找到其中最早的过期时间 + for _, v := range packs { + total += v.Size + if firstExpire == nil { + firstExpire = v.ExpiredTime + continue + } + if v.ExpiredTime != nil && firstExpire.After(*v.ExpiredTime) { + firstExpire = v.ExpiredTime + } + } + + // 用最早的过期时间计算缓存TTL，并写入缓存 + if firstExpire != nil { + ttl = firstExpire.Unix() - timeNow.Unix() + if ttl > 0 { + _ = cache.Set(cacheKey, total, int(ttl)) + } + } + + return total +} diff --git a/models/storage_pack_test.go b/models/storage_pack_test.go new file mode 100644 index 0000000..d4a3856 --- /dev/null +++ b/models/storage_pack_test.go @@ -0,0 +1,39 @@ +package model + +import ( + "github.com/DATA-DOG/go-sqlmock" + "github.com/jinzhu/gorm" + "github.com/stretchr/testify/assert" + "testing" + "time" +) + +func TestUser_GetAvailablePackSize(t *testing.T) { + asserts := assert.New(t) + user := User{ + Model: gorm.Model{ID: 1}, + } + + // 未命中缓存 + { + mock.ExpectQuery("SELECT(.+)"). + WithArgs(sqlmock.AnyArg(), 1). + WillReturnRows( + sqlmock.NewRows([]string{"id", "size", "expired_time"}). + AddRow(1, 10, time.Now().Add(time.Second*time.Duration(20))). + AddRow(3, 0, nil). + AddRow(2, 20, time.Now().Add(time.Second*time.Duration(10))), + ) + total := user.GetAvailablePackSize() + asserts.NoError(mock.ExpectationsWereMet()) + asserts.EqualValues(30, total) + } + + // 命中缓存 + { + total := user.GetAvailablePackSize() + asserts.NoError(mock.ExpectationsWereMet()) + asserts.EqualValues(30, total) + } + +} diff --git a/models/user.go b/models/user.go index ef266ed..d45f909 100644 --- a/models/user.go +++ b/models/user.go @@ -105,10 +105,11 @@ func (user *User) IncreaseStorageWithoutCheck(size uint64) { // GetRemainingCapacity 获取剩余配额 func (user *User) GetRemainingCapacity() uint64 { - if user.Group.MaxStorage <= user.Storage { + total := user.Group.MaxStorage + user.GetAvailablePackSize() + if total <= user.Storage { return 0 } - return user.Group.MaxStorage - user.Storage + return total - user.Storage } // GetPolicyID 获取用户当前的存储策略ID diff --git a/models/user_test.go b/models/user_test.go index d5f1023..63d7212 100644 --- a/models/user_test.go +++ b/models/user_test.go @@ -3,6 +3,7 @@ package model import ( "encoding/json" "github.com/DATA-DOG/go-sqlmock" + "github.com/HFO4/cloudreve/pkg/cache" "github.com/jinzhu/gorm" "github.com/pkg/errors" "github.com/stretchr/testify/assert" @@ -171,6 +172,7 @@ func TestUser_GetPolicyID(t *testing.T) { func TestUser_GetRemainingCapacity(t *testing.T) { asserts := assert.New(t) newUser := NewUser() + cache.Set("pack_size_0", uint64(0), 0) newUser.Group.MaxStorage = 100 asserts.Equal(uint64(100), newUser.GetRemainingCapacity()) @@ -186,6 +188,11 @@ func TestUser_GetRemainingCapacity(t *testing.T) { newUser.Group.MaxStorage = 100 newUser.Storage = 200 asserts.Equal(uint64(0), newUser.GetRemainingCapacity()) + + cache.Set("pack_size_0", uint64(10), 0) + newUser.Group.MaxStorage = 100 + newUser.Storage = 101 + asserts.Equal(uint64(9), newUser.GetRemainingCapacity()) } func TestUser_DeductionCapacity(t *testing.T) { @@ -204,6 +211,7 @@ func TestUser_DeductionCapacity(t *testing.T) { newUser, err := GetUserByID(1) newUser.Group.MaxStorage = 100 + cache.Set("pack_size_1", uint64(0), 0) asserts.NoError(err) asserts.NoError(mock.ExpectationsWereMet()) @@ -219,6 +227,10 @@ func TestUser_DeductionCapacity(t *testing.T) { asserts.Equal(false, newUser.IncreaseStorage(1)) asserts.Equal(uint64(100), newUser.Storage) + cache.Set("pack_size_1", uint64(1), 0) + asserts.Equal(true, newUser.IncreaseStorage(1)) + asserts.Equal(uint64(101), newUser.Storage) + asserts.True(newUser.IncreaseStorage(0)) } diff --git a/pkg/filesystem/file.go b/pkg/filesystem/file.go index efa17a8..6fb2265 100644 --- a/pkg/filesystem/file.go +++ b/pkg/filesystem/file.go @@ -99,7 +99,7 @@ func (fs *FileSystem) Preview(ctx context.Context, path string, isText bool) (*r // 如果是文本文件预览，需要检查大小限制 sizeLimit := model.GetIntSetting("maxEditSize", 2<<20) - if fs.FileTarget[0].Size > uint64(sizeLimit) { + if isText && fs.FileTarget[0].Size > uint64(sizeLimit) { return nil, ErrFileSizeTooBig } diff --git a/pkg/filesystem/hooks_test.go b/pkg/filesystem/hooks_test.go index 3d6d739..d21456f 100644 --- a/pkg/filesystem/hooks_test.go +++ b/pkg/filesystem/hooks_test.go @@ -28,6 +28,7 @@ func TestGenericBeforeUpload(t *testing.T) { Size: 5, Name: "1.txt", } + cache.Set("pack_size_0", uint64(0), 0) ctx := context.WithValue(context.Background(), fsctx.FileHeaderCtx, file) fs := FileSystem{ User: &model.User{ @@ -266,6 +267,7 @@ func TestHookIsFileExist(t *testing.T) { func TestHookValidateCapacity(t *testing.T) { asserts := assert.New(t) + cache.Set("pack_size_1", uint64(0), 0) fs := &FileSystem{User: &model.User{ Model: gorm.Model{ID: 1}, Storage: 0, @@ -313,6 +315,7 @@ func TestHookResetPolicy(t *testing.T) { func TestHookChangeCapacity(t *testing.T) { asserts := assert.New(t) + cache.Set("pack_size_1", uint64(0), 0) // 容量增加 失败 { diff --git a/pkg/filesystem/manage_test.go b/pkg/filesystem/manage_test.go index 68a7d89..c54ce56 100644 --- a/pkg/filesystem/manage_test.go +++ b/pkg/filesystem/manage_test.go @@ -5,6 +5,7 @@ import ( "errors" "github.com/DATA-DOG/go-sqlmock" model "github.com/HFO4/cloudreve/models" + "github.com/HFO4/cloudreve/pkg/cache" "github.com/HFO4/cloudreve/pkg/conf" "github.com/HFO4/cloudreve/pkg/serializer" "github.com/jinzhu/gorm" @@ -275,6 +276,7 @@ func TestFileSystem_ListDeleteDirs(t *testing.T) { func TestFileSystem_Delete(t *testing.T) { conf.DatabaseConfig.Type = "mysql" asserts := assert.New(t) + cache.Set("pack_size_1", uint64(0), 0) fs := &FileSystem{User: &model.User{ Model: gorm.Model{ ID: 1, @@ -381,6 +383,7 @@ func TestFileSystem_Delete(t *testing.T) { func TestFileSystem_Copy(t *testing.T) { asserts := assert.New(t) + cache.Set("pack_size_1", uint64(0), 0) fs := &FileSystem{User: &model.User{ Model: gorm.Model{ ID: 1, @@ -431,6 +434,7 @@ func TestFileSystem_Copy(t *testing.T) { func TestFileSystem_Move(t *testing.T) { asserts := assert.New(t) + cache.Set("pack_size_1", uint64(0), 0) fs := &FileSystem{User: &model.User{ Model: gorm.Model{ ID: 1, diff --git a/pkg/filesystem/validator_test.go b/pkg/filesystem/validator_test.go index 6d9f626..2dffc7a 100644 --- a/pkg/filesystem/validator_test.go +++ b/pkg/filesystem/validator_test.go @@ -5,6 +5,7 @@ import ( "database/sql" "github.com/DATA-DOG/go-sqlmock" model "github.com/HFO4/cloudreve/models" + "github.com/HFO4/cloudreve/pkg/cache" "github.com/jinzhu/gorm" "github.com/stretchr/testify/assert" "testing" @@ -42,6 +43,7 @@ func TestFileSystem_ValidateLegalName(t *testing.T) { func TestFileSystem_ValidateCapacity(t *testing.T) { asserts := assert.New(t) ctx := context.Background() + cache.Set("pack_size_0", uint64(0), 0) fs := FileSystem{ User: &model.User{ Storage: 10, @@ -57,6 +59,11 @@ func TestFileSystem_ValidateCapacity(t *testing.T) { fs.User.Storage = 5 asserts.False(fs.ValidateCapacity(ctx, 10)) asserts.Equal(uint64(5), fs.User.Storage) + + fs.User.Storage = 5 + cache.Set("pack_size_0", uint64(15), 0) + asserts.True(fs.ValidateCapacity(ctx, 10)) + asserts.Equal(uint64(15), fs.User.Storage) } func TestFileSystem_ValidateFileSize(t *testing.T) { diff --git a/pkg/serializer/user.go b/pkg/serializer/user.go index 0086d1c..8908de5 100644 --- a/pkg/serializer/user.go +++ b/pkg/serializer/user.go @@ -87,13 +87,14 @@ func BuildUserResponse(user model.User) Response { // BuildUserStorageResponse 序列化用户存储概况响应 func BuildUserStorageResponse(user model.User) Response { + total := user.Group.MaxStorage + user.GetAvailablePackSize() storageResp := storage{ Used: user.Storage, - Free: user.Group.MaxStorage - user.Storage, - Total: user.Group.MaxStorage, + Free: total - user.Storage, + Total: total, } - if user.Group.MaxStorage < user.Storage { + if total < user.Storage { storageResp.Free = 0 } diff --git a/pkg/serializer/user_test.go b/pkg/serializer/user_test.go index fa3df73..06eafb4 100644 --- a/pkg/serializer/user_test.go +++ b/pkg/serializer/user_test.go @@ -2,6 +2,7 @@ package serializer import ( model "github.com/HFO4/cloudreve/models" + "github.com/HFO4/cloudreve/pkg/cache" "github.com/stretchr/testify/assert" "testing" ) @@ -27,6 +28,7 @@ func TestBuildUserResponse(t *testing.T) { func TestBuildUserStorageResponse(t *testing.T) { asserts := assert.New(t) + cache.Set("pack_size_0", uint64(0), 0) { user := model.User{ @@ -58,4 +60,15 @@ func TestBuildUserStorageResponse(t *testing.T) { asserts.Equal(uint64(10), res.Data.(storage).Total) asserts.Equal(uint64(0), res.Data.(storage).Free) } + { + cache.Set("pack_size_0", uint64(1), 0) + user := model.User{ + Storage: 6, + Group: model.Group{MaxStorage: 10}, + } + res := BuildUserStorageResponse(user) + asserts.Equal(uint64(6), res.Data.(storage).Used) + asserts.Equal(uint64(11), res.Data.(storage).Total) + asserts.Equal(uint64(5), res.Data.(storage).Free) + } }