Feat: user storage pack

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()

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
+}

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)
+	}
+
+}

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

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))
 }