可以设置用户每天执行缓存任务的额度

This commit is contained in:
刘祥超
2022-06-05 21:15:28 +08:00
parent 71677a8638
commit a2f98d2f25
5 changed files with 84 additions and 15 deletions

View File

@@ -25,11 +25,6 @@ const (
HTTPCacheTaskTypeFetch HTTPCacheTaskType = "fetch"
)
const (
MaxKeysPerTask = 2000 // TODO 需要可以配置
MaxKeysPerDayByUser = 2000 // TODO 需要可以配置
)
type HTTPCacheTaskDAO dbs.DAO
func init() {

View File

@@ -187,6 +187,22 @@ func (this *HTTPCacheTaskKeyDAO) ResetCacheKeysWithTaskId(tx *dbs.Tx, taskId int
UpdateQuickly()
}
// CountUserTasksInDay 读取某个用户当前数量
// day YYYYMMDD
func (this *HTTPCacheTaskKeyDAO) CountUserTasksInDay(tx *dbs.Tx, userId int64, day string, taskType HTTPCacheTaskType) (int64, error) {
if userId <= 0 {
return 0, nil
}
// 这里需要包含已删除的
return this.Query(tx).
Where("taskId IN (SELECT id FROM "+SharedHTTPCacheTaskDAO.Table+" WHERE userId=:userId AND day=:day AND type=:type)").
Param("userId", userId).
Param("day", day).
Param("type", taskType).
Count()
}
// Clean 清理以往的任务
func (this *HTTPCacheTaskKeyDAO) Clean(tx *dbs.Tx, days int) error {
if days <= 0 {

View File

@@ -5,6 +5,7 @@ import (
_ "github.com/go-sql-driver/mysql"
_ "github.com/iwind/TeaGo/bootstrap"
"github.com/iwind/TeaGo/dbs"
timeutil "github.com/iwind/TeaGo/utils/time"
"testing"
)
@@ -30,3 +31,24 @@ func TestHTTPCacheTaskKeyDAO_UpdateKeyStatus(t *testing.T) {
}
t.Log("ok")
}
func TestHTTPCacheTaskKeyDAO_CountUserTasksInDay(t *testing.T) {
dbs.NotifyReady()
var dao = models.NewHTTPCacheTaskKeyDAO()
var tx *dbs.Tx
{
count, err := dao.CountUserTasksInDay(tx, 1, timeutil.Format("Ymd"), models.HTTPCacheTaskTypePurge)
if err != nil {
t.Fatal(err)
}
t.Log("count:", count)
}
{
count, err := dao.CountUserTasksInDay(tx, 1, timeutil.Format("Ymd"), models.HTTPCacheTaskTypeFetch)
if err != nil {
t.Fatal(err)
}
t.Log("count:", count)
}
}

View File

@@ -8,7 +8,9 @@ import (
"github.com/TeaOSLab/EdgeAPI/internal/errors"
"github.com/TeaOSLab/EdgeAPI/internal/utils"
"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
"github.com/TeaOSLab/EdgeCommon/pkg/userconfigs"
"github.com/iwind/TeaGo/types"
timeutil "github.com/iwind/TeaGo/utils/time"
)
// HTTPCacheTaskService 缓存任务管理
@@ -33,9 +35,50 @@ func (this *HTTPCacheTaskService) CreateHTTPCacheTask(ctx context.Context, req *
// 检查Key数量
var clusterId int64
if userId > 0 {
// TODO 限制当日刷新总条数(配额)
if len(req.Keys) > models.MaxKeysPerTask {
return nil, errors.New("too many keys (current:" + types.String(len(req.Keys)) + ", max:" + types.String(models.MaxKeysPerTask) + ")")
// 限制单次
var maxKeysPerTask = userconfigs.MaxCacheKeysPerTask
var maxKeysPerDay = userconfigs.MaxCacheKeysPerDay
serverConfig, err := models.SharedSysSettingDAO.ReadUserServerConfig(tx)
if err != nil {
return nil, err
}
if serverConfig != nil {
switch req.Type {
case models.HTTPCacheTaskTypePurge:
if serverConfig.HTTPCacheTaskPurgeConfig != nil {
if serverConfig.HTTPCacheTaskPurgeConfig.MaxKeysPerTask > 0 {
maxKeysPerTask = serverConfig.HTTPCacheTaskPurgeConfig.MaxKeysPerTask
}
if serverConfig.HTTPCacheTaskPurgeConfig.MaxKeysPerDay > 0 {
maxKeysPerDay = serverConfig.HTTPCacheTaskPurgeConfig.MaxKeysPerDay
}
}
case models.HTTPCacheTaskTypeFetch:
if serverConfig.HTTPCacheTaskFetchConfig != nil {
if serverConfig.HTTPCacheTaskFetchConfig.MaxKeysPerTask > 0 {
maxKeysPerTask = serverConfig.HTTPCacheTaskFetchConfig.MaxKeysPerTask
}
if serverConfig.HTTPCacheTaskFetchConfig.MaxKeysPerDay > 0 {
maxKeysPerDay = serverConfig.HTTPCacheTaskFetchConfig.MaxKeysPerDay
}
}
}
}
if maxKeysPerTask > 0 && len(req.Keys) > types.Int(maxKeysPerTask) {
return nil, errors.New("too many keys in task (current:" + types.String(len(req.Keys)) + ", max:" + types.String(maxKeysPerTask) + ")")
}
if maxKeysPerDay > 0 {
countInDay, err := models.SharedHTTPCacheTaskKeyDAO.CountUserTasksInDay(tx, userId, timeutil.Format("Ymd"), req.Type)
if err != nil {
return nil, err
}
if types.Int(countInDay)+len(req.Keys) > types.Int(maxKeysPerDay) {
return nil, errors.New("too many keys in today (current:" + types.String(types.Int(countInDay)+len(req.Keys)) + ", max:" + types.String(maxKeysPerDay) + ")")
}
}
clusterId, err = models.SharedUserDAO.FindUserClusterId(tx, userId)

View File

@@ -6,11 +6,9 @@ import (
"context"
"encoding/json"
"github.com/TeaOSLab/EdgeAPI/internal/db/models"
"github.com/TeaOSLab/EdgeAPI/internal/errors"
"github.com/TeaOSLab/EdgeAPI/internal/utils"
"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
"github.com/iwind/TeaGo/dbs"
"github.com/iwind/TeaGo/types"
)
// HTTPCacheTaskKeyService 缓存任务Key管理
@@ -30,11 +28,6 @@ func (this *HTTPCacheTaskKeyService) ValidateHTTPCacheTaskKeys(ctx context.Conte
// 检查Key数量
var clusterId int64
if userId > 0 {
// TODO 限制当日刷新总条数(配额)
if len(req.Keys) > models.MaxKeysPerTask {
return nil, errors.New("too many keys (current:" + types.String(len(req.Keys)) + ", max:" + types.String(models.MaxKeysPerTask) + ")")
}
clusterId, err = models.SharedUserDAO.FindUserClusterId(tx, userId)
if err != nil {
return nil, err