diff --git a/internal/db/models/http_cache_task_dao.go b/internal/db/models/http_cache_task_dao.go index 240ee503..5c0cb49e 100644 --- a/internal/db/models/http_cache_task_dao.go +++ b/internal/db/models/http_cache_task_dao.go @@ -25,11 +25,6 @@ const ( HTTPCacheTaskTypeFetch HTTPCacheTaskType = "fetch" ) -const ( - MaxKeysPerTask = 2000 // TODO 需要可以配置 - MaxKeysPerDayByUser = 2000 // TODO 需要可以配置 -) - type HTTPCacheTaskDAO dbs.DAO func init() { diff --git a/internal/db/models/http_cache_task_key_dao.go b/internal/db/models/http_cache_task_key_dao.go index 9dc09b1a..80cf11b8 100644 --- a/internal/db/models/http_cache_task_key_dao.go +++ b/internal/db/models/http_cache_task_key_dao.go @@ -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 { diff --git a/internal/db/models/http_cache_task_key_dao_test.go b/internal/db/models/http_cache_task_key_dao_test.go index d2dee42e..66ff3d8e 100644 --- a/internal/db/models/http_cache_task_key_dao_test.go +++ b/internal/db/models/http_cache_task_key_dao_test.go @@ -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) + } +} diff --git a/internal/rpc/services/service_http_cache_task.go b/internal/rpc/services/service_http_cache_task.go index 450d6798..bea25fdc 100644 --- a/internal/rpc/services/service_http_cache_task.go +++ b/internal/rpc/services/service_http_cache_task.go @@ -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) diff --git a/internal/rpc/services/service_http_cache_task_key.go b/internal/rpc/services/service_http_cache_task_key.go index b2e09d6d..eb23d1dc 100644 --- a/internal/rpc/services/service_http_cache_task_key.go +++ b/internal/rpc/services/service_http_cache_task_key.go @@ -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