feat: add channel quota reset controls
Add per-channel quota limits with reset scheduling and runtime checks so channel capacity can be capped and restored on configured periods.
This commit is contained in:
zizi
parent
0a0616a370
commit
f245e9bdf5
@ -510,6 +510,20 @@ func validateChannel(channel *model.Channel, isAdd bool) error {
|
||||
}
|
||||
}
|
||||
|
||||
// 校验 quota_reset_period
|
||||
if channel.QuotaResetPeriod != "" {
|
||||
normalized := model.NormalizeChannelResetPeriod(channel.QuotaResetPeriod)
|
||||
if normalized == model.ChannelResetNever && channel.QuotaResetPeriod != "never" && channel.QuotaResetPeriod != "" {
|
||||
return fmt.Errorf("invalid quota_reset_period: %s", channel.QuotaResetPeriod)
|
||||
}
|
||||
channel.QuotaResetPeriod = normalized
|
||||
}
|
||||
|
||||
// custom 模式必须有 customSeconds > 0
|
||||
if model.NormalizeChannelResetPeriod(channel.QuotaResetPeriod) == model.ChannelResetCustom && channel.QuotaResetCustomSeconds <= 0 {
|
||||
return fmt.Errorf("quota_reset_custom_seconds must be > 0 when quota_reset_period is 'custom'")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
@ -602,6 +616,7 @@ func AddChannel(c *gin.Context) {
|
||||
}
|
||||
|
||||
addChannelRequest.Channel.CreatedTime = common.GetTimestamp()
|
||||
model.InitChannelNextResetTime(addChannelRequest.Channel)
|
||||
keys := make([]string, 0)
|
||||
switch addChannelRequest.Mode {
|
||||
case "multi_to_single":
|
||||
@ -974,6 +989,7 @@ func UpdateChannel(c *gin.Context) {
|
||||
// 覆盖模式：直接使用新密钥（默认行为，不需要特殊处理）
|
||||
}
|
||||
}
|
||||
model.InitChannelNextResetTime(&channel.Channel)
|
||||
err = channel.Update()
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
|
||||
@ -0,0 +1,524 @@
|
||||
# 渠道订阅类型额度实现方案
|
||||
|
||||
**版本：** v2.0
|
||||
**日期：** 2026-05-19
|
||||
**状态：** 执行版
|
||||
|
||||
---
|
||||
|
||||
## 1. 目标与范围
|
||||
|
||||
参照 `model/subscription.go` 的订阅额度重置机制，为 Channel 增加**额度上限 + 周期重置**功能。详细设计决策见审阅版。
|
||||
|
||||
---
|
||||
|
||||
## 2. 参考实现索引
|
||||
|
||||
| 功能 | 文件:行号 |
|
||||
|------|-----------|
|
||||
| 重置周期常量 | `model/subscription.go:26-33` |
|
||||
| `NormalizeResetPeriod` | `model/subscription.go:300-307` |
|
||||
| `calcNextResetTime` | `model/subscription.go:309-348` |
|
||||
| `maybeResetUserSubscriptionWithPlanTx` | `model/subscription.go:933-967` |
|
||||
| `ResetDueSubscriptions` | `model/subscription.go:1100-1140` |
|
||||
| 定时任务 | `service/subscription_reset_task.go:29-93` |
|
||||
| 定时任务启动 | `main.go:120` |
|
||||
| 渠道额度扣减 | `model/channel.go:824-837` |
|
||||
| 请求完成后扣减调用点 | `service/text_quota.go:370` |
|
||||
| 渠道校验 | `controller/channel.go:457-514` |
|
||||
| 渠道缓存 | `model/channel_cache.go:193-205` |
|
||||
|
||||
---
|
||||
|
||||
## 3. Channel 模型新增字段（修改 `model/channel.go`）
|
||||
|
||||
```go
|
||||
// 新增字段，接在 UsedQuota 之后
|
||||
QuotaLimit int64 `json:"quota_limit" gorm:"bigint;default:0"`
|
||||
QuotaResetPeriod string `json:"quota_reset_period" gorm:"type:varchar(16);default:'never'"`
|
||||
QuotaResetCustomSeconds int `json:"quota_reset_custom_seconds" gorm:"type:int;default:0"`
|
||||
LastResetTime int64 `json:"last_reset_time" gorm:"bigint;default:0"`
|
||||
NextResetTime int64 `json:"next_reset_time" gorm:"bigint;default:0"`
|
||||
```
|
||||
|
||||
**语义**：
|
||||
- `QuotaLimit = 0` → 不启用额度管控（原有行为）
|
||||
- `QuotaLimit > 0` + `QuotaResetPeriod != 'never'` → 启用额度管控 + 周期重置
|
||||
- `QuotaLimit > 0` + `QuotaResetPeriod = 'never'` → 启用额度管控，永不重置
|
||||
- `UsedQuota` 沿用现有字段，重置时清零
|
||||
|
||||
---
|
||||
|
||||
## 4. 新增 `model/channel_quota.go`
|
||||
|
||||
### 4.1 重置周期常量
|
||||
|
||||
```go
|
||||
const (
|
||||
ChannelResetNever = "never"
|
||||
ChannelResetDaily = "daily"
|
||||
ChannelResetWeekly = "weekly"
|
||||
ChannelResetMonthly = "monthly"
|
||||
ChannelResetCustom = "custom"
|
||||
)
|
||||
```
|
||||
|
||||
与 `subscription.go` 的常量值一致，命名前缀改为 `Channel` 避免冲突。
|
||||
|
||||
### 4.2 NormalizeChannelResetPeriod
|
||||
|
||||
```go
|
||||
func NormalizeChannelResetPeriod(period string) string {
|
||||
switch strings.TrimSpace(period) {
|
||||
case ChannelResetDaily, ChannelResetWeekly, ChannelResetMonthly, ChannelResetCustom:
|
||||
return strings.TrimSpace(period)
|
||||
default:
|
||||
return ChannelResetNever
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
### 4.3 calcChannelNextResetTime
|
||||
|
||||
与 `subscription.go:calcNextResetTime` 逻辑一致，但接收 `period` 和 `customSeconds` 参数而非 `*SubscriptionPlan`：
|
||||
|
||||
```go
|
||||
func calcChannelNextResetTime(base time.Time, period string, customSeconds int) int64 {
|
||||
period = NormalizeChannelResetPeriod(period)
|
||||
if period == ChannelResetNever {
|
||||
return 0
|
||||
}
|
||||
var next time.Time
|
||||
switch period {
|
||||
case ChannelResetDaily:
|
||||
next = time.Date(base.Year(), base.Month(), base.Day(), 0, 0, 0, 0, base.Location()).
|
||||
AddDate(0, 0, 1)
|
||||
case ChannelResetWeekly:
|
||||
weekday := int(base.Weekday())
|
||||
if weekday == 0 {
|
||||
weekday = 7
|
||||
}
|
||||
daysUntil := 8 - weekday
|
||||
next = time.Date(base.Year(), base.Month(), base.Day(), 0, 0, 0, 0, base.Location()).
|
||||
AddDate(0, 0, daysUntil)
|
||||
case ChannelResetMonthly:
|
||||
next = time.Date(base.Year(), base.Month(), 1, 0, 0, 0, 0, base.Location()).
|
||||
AddDate(0, 1, 0)
|
||||
case ChannelResetCustom:
|
||||
if customSeconds <= 0 {
|
||||
return 0
|
||||
}
|
||||
next = base.Add(time.Duration(customSeconds) * time.Second)
|
||||
default:
|
||||
return 0
|
||||
}
|
||||
return next.Unix()
|
||||
}
|
||||
```
|
||||
|
||||
与 `calcNextResetTime` 的区别：无 `endUnix` 参数（渠道无到期时间）。
|
||||
|
||||
### 4.4 maybeResetChannelQuotaTx
|
||||
|
||||
与 `subscription.go:maybeResetUserSubscriptionWithPlanTx` 逻辑一致，针对 Channel：
|
||||
|
||||
```go
|
||||
func maybeResetChannelQuotaTx(tx *gorm.DB, channel *Channel, now int64) error {
|
||||
if tx == nil || channel == nil {
|
||||
return errors.New("invalid reset args")
|
||||
}
|
||||
if channel.NextResetTime > 0 && channel.NextResetTime > now {
|
||||
return nil
|
||||
}
|
||||
if NormalizeChannelResetPeriod(channel.QuotaResetPeriod) == ChannelResetNever {
|
||||
return nil
|
||||
}
|
||||
baseUnix := channel.LastResetTime
|
||||
if baseUnix <= 0 {
|
||||
baseUnix = channel.CreatedTime
|
||||
}
|
||||
base := time.Unix(baseUnix, 0)
|
||||
next := calcChannelNextResetTime(base, channel.QuotaResetPeriod, channel.QuotaResetCustomSeconds)
|
||||
|
||||
advanced := false
|
||||
for next > 0 && next <= now {
|
||||
advanced = true
|
||||
base = time.Unix(next, 0)
|
||||
next = calcChannelNextResetTime(base, channel.QuotaResetPeriod, channel.QuotaResetCustomSeconds)
|
||||
}
|
||||
|
||||
if !advanced {
|
||||
if channel.NextResetTime == 0 && next > 0 {
|
||||
return tx.Model(&Channel{}).Where("id = ?", channel.Id).Updates(map[string]interface{}{
|
||||
"next_reset_time": next,
|
||||
"last_reset_time": base.Unix(),
|
||||
}).Error
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// 重置额度，只用 Updates 不用 Save，避免覆盖并发更新
|
||||
return tx.Model(&Channel{}).Where("id = ?", channel.Id).Updates(map[string]interface{}{
|
||||
"used_quota": 0,
|
||||
"last_reset_time": base.Unix(),
|
||||
"next_reset_time": next,
|
||||
}).Error
|
||||
}
|
||||
```
|
||||
|
||||
**与订阅实现的区别**：用 `tx.Model(&Channel{}).Where("id = ?").Updates(...)` 替代 `tx.Save(channel)`，只更新重置相关字段，避免覆盖 `UsedQuota` 的并发扣减。
|
||||
|
||||
### 4.5 ResetDueChannels
|
||||
|
||||
与 `subscription.go:ResetDueSubscriptions` 逻辑一致：
|
||||
|
||||
```go
|
||||
func ResetDueChannels(limit int) (int, error) {
|
||||
if limit <= 0 {
|
||||
limit = 200
|
||||
}
|
||||
now := GetDBTimestamp()
|
||||
var channels []Channel
|
||||
if err := DB.Where("next_reset_time > 0 AND next_reset_time <= ? AND status = ? AND quota_limit > 0",
|
||||
now, common.ChannelStatusEnabled).
|
||||
Order("next_reset_time asc").
|
||||
Limit(limit).
|
||||
Find(&channels).Error; err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if len(channels) == 0 {
|
||||
return 0, nil
|
||||
}
|
||||
resetCount := 0
|
||||
for _, ch := range channels {
|
||||
chCopy := ch
|
||||
err := DB.Transaction(func(tx *gorm.DB) error {
|
||||
var locked Channel
|
||||
if err := tx.Set("gorm:query_option", "FOR UPDATE").
|
||||
Where("id = ? AND next_reset_time > 0 AND next_reset_time <= ?", chCopy.Id, now).
|
||||
First(&locked).Error; err != nil {
|
||||
return nil
|
||||
}
|
||||
if err := maybeResetChannelQuotaTx(tx, &locked, now); err != nil {
|
||||
return err
|
||||
}
|
||||
resetCount++
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
return resetCount, err
|
||||
}
|
||||
}
|
||||
return resetCount, nil
|
||||
}
|
||||
```
|
||||
|
||||
**与订阅实现的区别**：
|
||||
- 查询条件增加 `AND quota_limit > 0`，跳过未启用额度管控的渠道
|
||||
- 无需加载关联的 Plan（渠道的重置参数直接在 Channel 上）
|
||||
|
||||
### 4.6 CheckChannelQuota
|
||||
|
||||
```go
|
||||
func CheckChannelQuota(channelId int) error {
|
||||
var channel Channel
|
||||
if err := DB.Select("id, used_quota, quota_limit, quota_reset_period, next_reset_time").
|
||||
Where("id = ?", channelId).First(&channel).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
// 不启用额度管控
|
||||
if channel.QuotaLimit <= 0 {
|
||||
return nil
|
||||
}
|
||||
// 懒重置：在额度检查之前
|
||||
if channel.NextResetTime > 0 {
|
||||
now := GetDBTimestamp()
|
||||
if channel.NextResetTime <= now {
|
||||
err := DB.Transaction(func(tx *gorm.DB) error {
|
||||
var locked Channel
|
||||
if err := tx.Set("gorm:query_option", "FOR UPDATE").
|
||||
Where("id = ? AND next_reset_time > 0 AND next_reset_time <= ?", channelId, now).
|
||||
First(&locked).Error; err != nil {
|
||||
return nil // 已被其他事务处理
|
||||
}
|
||||
return maybeResetChannelQuotaTx(tx, &locked, now)
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
// 重新读取重置后的值
|
||||
if err := DB.Select("used_quota, quota_limit").
|
||||
Where("id = ?", channelId).First(&channel).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
if channel.UsedQuota >= channel.QuotaLimit {
|
||||
return fmt.Errorf("channel quota exceeded (used=%d, limit=%d)", channel.UsedQuota, channel.QuotaLimit)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
```
|
||||
|
||||
**设计要点**：
|
||||
- 直接查 DB，不依赖缓存（缓存 `UsedQuota` 可能过期）
|
||||
- 懒重置在额度检查之前执行（修复原方案的问题 1）
|
||||
- 重置后重新读取确保拿到最新值
|
||||
|
||||
### 4.7 InitChannelNextResetTime
|
||||
|
||||
创建/更新渠道时初始化 `NextResetTime`：
|
||||
|
||||
```go
|
||||
func InitChannelNextResetTime(channel *Channel) {
|
||||
if channel.QuotaLimit <= 0 {
|
||||
channel.NextResetTime = 0
|
||||
channel.LastResetTime = 0
|
||||
return
|
||||
}
|
||||
period := NormalizeChannelResetPeriod(channel.QuotaResetPeriod)
|
||||
if period == ChannelResetNever {
|
||||
channel.NextResetTime = 0
|
||||
channel.LastResetTime = 0
|
||||
return
|
||||
}
|
||||
baseUnix := channel.LastResetTime
|
||||
if baseUnix <= 0 {
|
||||
baseUnix = channel.CreatedTime
|
||||
}
|
||||
if baseUnix <= 0 {
|
||||
baseUnix = GetDBTimestamp()
|
||||
}
|
||||
base := time.Unix(baseUnix, 0)
|
||||
next := calcChannelNextResetTime(base, channel.QuotaResetPeriod, channel.QuotaResetCustomSeconds)
|
||||
channel.NextResetTime = next
|
||||
if channel.LastResetTime <= 0 {
|
||||
channel.LastResetTime = baseUnix
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 5. 新增 `service/channel_quota_reset_task.go`
|
||||
|
||||
与 `service/subscription_reset_task.go` 结构一致，仅替换调用：
|
||||
|
||||
```go
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/logger"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
|
||||
"github.com/bytedance/gopkg/util/gopool"
|
||||
)
|
||||
|
||||
const (
|
||||
channelQuotaResetTickInterval = 1 * time.Minute
|
||||
channelQuotaResetBatchSize = 300
|
||||
)
|
||||
|
||||
var (
|
||||
channelQuotaResetOnce sync.Once
|
||||
channelQuotaResetRunning atomic.Bool
|
||||
)
|
||||
|
||||
func StartChannelQuotaResetTask() {
|
||||
channelQuotaResetOnce.Do(func() {
|
||||
if !common.IsMasterNode {
|
||||
return
|
||||
}
|
||||
gopool.Go(func() {
|
||||
logger.LogInfo(context.Background(), fmt.Sprintf("channel quota reset task started: tick=%s", channelQuotaResetTickInterval))
|
||||
ticker := time.NewTicker(channelQuotaResetTickInterval)
|
||||
defer ticker.Stop()
|
||||
|
||||
runChannelQuotaResetOnce()
|
||||
for range ticker.C {
|
||||
runChannelQuotaResetOnce()
|
||||
}
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
func runChannelQuotaResetOnce() {
|
||||
if !channelQuotaResetRunning.CompareAndSwap(false, true) {
|
||||
return
|
||||
}
|
||||
defer channelQuotaResetRunning.Store(false)
|
||||
|
||||
ctx := context.Background()
|
||||
totalReset := 0
|
||||
for {
|
||||
n, err := model.ResetDueChannels(channelQuotaResetBatchSize)
|
||||
if err != nil {
|
||||
logger.LogWarn(ctx, fmt.Sprintf("channel quota reset task failed: %v", err))
|
||||
return
|
||||
}
|
||||
if n == 0 {
|
||||
break
|
||||
}
|
||||
totalReset += n
|
||||
if n < channelQuotaResetBatchSize {
|
||||
break
|
||||
}
|
||||
}
|
||||
if common.DebugEnabled && totalReset > 0 {
|
||||
logger.LogDebug(ctx, "channel quota maintenance: reset_count=%d", totalReset)
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 6. 修改 `service/text_quota.go`
|
||||
|
||||
在 `PostTextConsumeQuota` 中，请求扣减前增加额度检查：
|
||||
|
||||
```go
|
||||
// 原代码 (line 369-370):
|
||||
model.UpdateUserUsedQuotaAndRequestCount(relayInfo.UserId, summary.Quota)
|
||||
model.UpdateChannelUsedQuota(relayInfo.ChannelId, summary.Quota)
|
||||
|
||||
// 修改后:
|
||||
model.UpdateUserUsedQuotaAndRequestCount(relayInfo.UserId, summary.Quota)
|
||||
// 渠道额度管控：检查额度是否超限
|
||||
if err := model.CheckChannelQuota(relayInfo.ChannelId); err != nil {
|
||||
// 额度已超限，仍扣减已消耗的额度，但记录警告
|
||||
logger.LogWarn(ctx, fmt.Sprintf("channel quota exceeded: channel_id=%d, %v", relayInfo.ChannelId, err))
|
||||
}
|
||||
model.UpdateChannelUsedQuota(relayInfo.ChannelId, summary.Quota)
|
||||
```
|
||||
|
||||
**注意**：额度检查放在**请求完成后扣减前**，而非请求前。原因：
|
||||
1. 请求前的 `estimatedQuota` 不可靠（流式请求无法预估）
|
||||
2. 请求已完成，必须扣减实际消耗
|
||||
3. 超限后下次请求会被 `distributor` 拒绝（见第 7 节）
|
||||
|
||||
这样设计：**当前请求正常完成并扣费，但下次请求会被拒绝**。这是"实时扣减"模式的自然语义——不做预扣，只做事后拦截。
|
||||
|
||||
---
|
||||
|
||||
## 7. 修改 `middleware/distributor.go`
|
||||
|
||||
在渠道选中后、`c.Next()` 前，检查渠道额度：
|
||||
|
||||
```go
|
||||
// 原代码 (line 158-160):
|
||||
common.SetContextKey(c, constant.ContextKeyRequestStartTime, time.Now())
|
||||
SetupContextForSelectedChannel(c, channel, modelRequest.Model)
|
||||
c.Next()
|
||||
|
||||
// 修改后:
|
||||
common.SetContextKey(c, constant.ContextKeyRequestStartTime, time.Now())
|
||||
SetupContextForSelectedChannel(c, channel, modelRequest.Model)
|
||||
// 渠道额度管控：请求前检查
|
||||
if err := model.CheckChannelQuota(channel.Id); err != nil {
|
||||
abortWithOpenAiMessage(c, http.StatusTooManyRequests, err.Error())
|
||||
return
|
||||
}
|
||||
c.Next()
|
||||
```
|
||||
|
||||
**此检查是真正的拦截点**：
|
||||
- `CheckChannelQuota` 内部会先执行懒重置，再检查 `UsedQuota >= QuotaLimit`
|
||||
- 额度超限时直接返回 429，不转发请求到上游
|
||||
- `text_quota.go` 中的检查是兜底（防止并发场景下漏过）
|
||||
|
||||
---
|
||||
|
||||
## 8. 修改 `controller/channel.go`
|
||||
|
||||
### 8.1 validateChannel 增加校验
|
||||
|
||||
```go
|
||||
// 在 validateChannel 末尾 return nil 之前添加：
|
||||
|
||||
// 校验 quota_reset_period
|
||||
if channel.QuotaResetPeriod != "" {
|
||||
normalized := model.NormalizeChannelResetPeriod(channel.QuotaResetPeriod)
|
||||
if normalized == model.ChannelResetNever && channel.QuotaResetPeriod != "never" && channel.QuotaResetPeriod != "" {
|
||||
return fmt.Errorf("invalid quota_reset_period: %s", channel.QuotaResetPeriod)
|
||||
}
|
||||
channel.QuotaResetPeriod = normalized
|
||||
}
|
||||
|
||||
// custom 模式必须有 customSeconds > 0
|
||||
if model.NormalizeChannelResetPeriod(channel.QuotaResetPeriod) == model.ChannelResetCustom && channel.QuotaResetCustomSeconds <= 0 {
|
||||
return fmt.Errorf("quota_reset_custom_seconds must be > 0 when quota_reset_period is 'custom'")
|
||||
}
|
||||
```
|
||||
|
||||
### 8.2 创建/更新渠道时初始化 NextResetTime
|
||||
|
||||
在 `AddChannel` 和 `UpdateChannel` 中，保存前调用：
|
||||
|
||||
```go
|
||||
model.InitChannelNextResetTime(channel)
|
||||
```
|
||||
|
||||
这确保 `NextResetTime` 不为 0（修复原方案的问题 2），定时任务能正确选中需要重置的渠道。
|
||||
|
||||
---
|
||||
|
||||
## 9. 修改 `main.go`
|
||||
|
||||
在 `StartSubscriptionQuotaResetTask()` 后添加：
|
||||
|
||||
```go
|
||||
service.StartChannelQuotaResetTask()
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 10. 数据库迁移
|
||||
|
||||
GORM `AutoMigrate` 自动处理字段新增。`Channel{}` 已在 `model/main.go` 的 AutoMigrate 列表中，新增字段会自动添加。
|
||||
|
||||
无需手动 SQL。
|
||||
|
||||
---
|
||||
|
||||
## 11. 文件清单
|
||||
|
||||
### 11.1 新增
|
||||
|
||||
| 文件 | 说明 |
|
||||
|------|------|
|
||||
| `model/channel_quota.go` | 常量、重置时间计算、重置事务、ResetDueChannels、CheckChannelQuota、InitChannelNextResetTime |
|
||||
| `service/channel_quota_reset_task.go` | 定时任务 |
|
||||
|
||||
### 11.2 修改
|
||||
|
||||
| 文件 | 修改内容 |
|
||||
|------|----------|
|
||||
| `model/channel.go` | 新增 5 个字段 |
|
||||
| `service/text_quota.go` | 扣减前额度检查（兜底） |
|
||||
| `middleware/distributor.go` | 请求前额度检查（主拦截点） |
|
||||
| `controller/channel.go` | 校验 + 初始化 NextResetTime |
|
||||
| `main.go` | 启动定时任务 |
|
||||
|
||||
---
|
||||
|
||||
## 12. 前端（后续迭代）
|
||||
|
||||
本次实现不含前端修改。前端渠道编辑表单增加 `QuotaLimit` 和 `QuotaResetPeriod` 字段、列表展示额度使用率，列为后续迭代。
|
||||
|
||||
---
|
||||
|
||||
## 13. 测试要点
|
||||
|
||||
1. **重置周期正确性**：daily/weekly/monthly/custom，验证 `NextResetTime` 计算正确
|
||||
2. **懒重置**：定时任务未触发时，请求前懒检查能补救
|
||||
3. **额度拦截**：`UsedQuota >= QuotaLimit` 时请求返回 429
|
||||
4. **兼容性**：未设置 `QuotaLimit` 的历史渠道行为不变
|
||||
5. **NextResetTime 初始化**：创建/更新渠道后 `NextResetTime` 正确
|
||||
6. **三库兼容**：SQLite / MySQL / PostgreSQL 均通过
|
||||
258
docs/agent-specs/2026-05-19-channel-subscription-quota-review.md
Normal file
258
docs/agent-specs/2026-05-19-channel-subscription-quota-review.md
Normal file
@ -0,0 +1,258 @@
|
||||
# 渠道订阅类型额度设计文档
|
||||
|
||||
**版本：** v2.0
|
||||
**日期：** 2026-05-19
|
||||
**状态：** 审阅版
|
||||
|
||||
---
|
||||
|
||||
## 1. 背景与问题
|
||||
|
||||
### 1.1 现状
|
||||
|
||||
Channel 模型仅有两种额度相关字段：
|
||||
|
||||
| 字段 | 语义 | 来源 |
|
||||
|------|------|------|
|
||||
| `Balance` | 上游 API 实际余额（USD） | 上游 API 查询 |
|
||||
| `UsedQuota` | 累计消耗量 | 本地扣减，永不重置 |
|
||||
|
||||
两者都无法满足"固定额度上限 + 周期重置"的运营需求。
|
||||
|
||||
### 1.2 需求场景
|
||||
|
||||
运营者希望对渠道配额管控：
|
||||
|
||||
- 渠道 A：每天最多 1 万 quota，超限拒绝请求
|
||||
- 渠道 B：每月最多 50 万 quota，超限拒绝请求
|
||||
- 渠道 C：不设上限，仅记录消耗（与现有 `UsedQuota` 等价）
|
||||
|
||||
### 1.3 已有参考实现
|
||||
|
||||
`model/subscription.go` 中已实现完整的用户订阅额度重置机制：
|
||||
|
||||
```go
|
||||
// 订阅套餐可配置重置周期
|
||||
QuotaResetPeriod // daily/weekly/monthly/custom/never
|
||||
|
||||
// 用户订阅记录
|
||||
AmountTotal // 总额度
|
||||
AmountUsed // 已消耗
|
||||
LastResetTime // 上次重置
|
||||
NextResetTime // 下次重置
|
||||
```
|
||||
|
||||
渠道额度直接复用此设计模式，但消费模式不同（见 3.1）。
|
||||
|
||||
---
|
||||
|
||||
## 2. 目标与非目标
|
||||
|
||||
### 2.1 目标
|
||||
|
||||
- 渠道可配置额度上限（`QuotaLimit`）
|
||||
- 渠道可配置重置周期（`QuotaResetPeriod`：daily / weekly / monthly / custom / never）
|
||||
- 周期到达时自动重置 `UsedQuota` 为 0
|
||||
- 额度耗尽时拒绝请求（返回 429）
|
||||
|
||||
### 2.2 非目标
|
||||
|
||||
- 不支持预扣 + 结算模式（与订阅的 `PreConsume` 不同）
|
||||
- 不支持按渠道内单个 key 分别管控
|
||||
- 不支持按模型分组分别设置额度
|
||||
- 不支持额度耗尽自动降级到备用渠道（仅拒绝）
|
||||
- 不支持额度使用率告警通知
|
||||
|
||||
---
|
||||
|
||||
## 3. 核心决策
|
||||
|
||||
### 3.1 消费模式：实时扣减（非预扣 + 结算）
|
||||
|
||||
订阅采用"预扣 + 结算"双阶段，因为订阅钱包需要精确余额。
|
||||
|
||||
渠道采用**实时扣减**，更简单：
|
||||
|
||||
1. 请求前检查：`UsedQuota >= QuotaLimit` → 拒绝
|
||||
2. 请求完成：`UsedQuota += actualQuota`
|
||||
|
||||
不做预估算额。`UsedQuota >= QuotaLimit` 时拒绝，意味着当前周期内已消耗量达到上限。由于请求完成才扣减，最后一个被放行的请求可能实际消耗后超过 `QuotaLimit`，但这是可接受的——与 `UpdateChannelUsedQuota` 的原子递增语义一致。
|
||||
|
||||
### 3.2 重置时机：定时任务 + 懒检查兜底
|
||||
|
||||
与订阅一致：
|
||||
|
||||
```go
|
||||
// 定时任务（每分钟）
|
||||
ResetDueChannels()
|
||||
|
||||
// 请求前懒检查（定时任务未及时触发时补救）
|
||||
maybeResetChannelQuotaTx()
|
||||
```
|
||||
|
||||
**关键**：懒检查必须在额度检查之前执行，否则重置周期刚到时旧 `UsedQuota` 仍会拒绝合法请求。
|
||||
|
||||
### 3.3 并发安全：FOR UPDATE 行锁 + Updates 精确更新
|
||||
|
||||
```go
|
||||
tx.Set("gorm:query_option", "FOR UPDATE").
|
||||
Where("id = ? AND next_reset_time <= ?", channelId, now).
|
||||
First(&locked)
|
||||
```
|
||||
|
||||
重置时只更新相关字段，不用 `Save` 避免覆盖并发更新：
|
||||
|
||||
```go
|
||||
tx.Model(&Channel{}).Where("id = ?", channel.Id).Updates(map[string]interface{}{
|
||||
"used_quota": 0,
|
||||
"last_reset_time": base.Unix(),
|
||||
"next_reset_time": next,
|
||||
})
|
||||
```
|
||||
|
||||
### 3.4 额度耗尽行为：返回 429
|
||||
|
||||
额度耗尽时返回 HTTP 429，不自动降级。降级涉及渠道选择逻辑改动，复杂度高，列为非目标。
|
||||
|
||||
### 3.5 Balance 与 QuotaLimit 的关系
|
||||
|
||||
两者独立共存：
|
||||
|
||||
| 字段 | 语义 | 来源 |
|
||||
|------|------|------|
|
||||
| `Balance` | 上游 API 实际余额 | 上游 API 查询 |
|
||||
| `UsedQuota` | 本地累计消耗 | 本地扣减 |
|
||||
| `QuotaLimit` | 运营者设定的额度上限 | 配置 |
|
||||
|
||||
`Balance` 是上游真实余额，`QuotaLimit` 是本地配额管控，两者互不影响。运营者可同时使用或只选其一。
|
||||
|
||||
### 3.6 多 Key 模式
|
||||
|
||||
`UsedQuota` 是整个渠道的累计，不区分单个 key。多 key 模式下仅全局管控，与现有行为一致。
|
||||
|
||||
---
|
||||
|
||||
## 4. 影响范围
|
||||
|
||||
### 4.1 新增文件
|
||||
|
||||
| 文件 | 说明 |
|
||||
|------|------|
|
||||
| `model/channel_quota.go` | 渠道额度核心逻辑（重置时间计算、重置事务、ResetDueChannels、CheckChannelQuota） |
|
||||
| `service/channel_quota_reset_task.go` | 渠道额度重置定时任务 |
|
||||
|
||||
### 4.2 修改文件
|
||||
|
||||
| 文件 | 修改内容 |
|
||||
|------|----------|
|
||||
| `model/channel.go` | 新增 5 个字段 |
|
||||
| `service/text_quota.go` | 请求前懒重置 + 额度检查 |
|
||||
| `controller/channel.go` | 校验新字段 + 创建/更新时初始化 NextResetTime |
|
||||
| `main.go` | 启动定时任务 |
|
||||
|
||||
### 4.3 数据库字段变更
|
||||
|
||||
```sql
|
||||
ALTER TABLE channels ADD COLUMN quota_limit BIGINT DEFAULT 0;
|
||||
ALTER TABLE channels ADD COLUMN quota_reset_period VARCHAR(16) DEFAULT 'never';
|
||||
ALTER TABLE channels ADD COLUMN quota_reset_custom_seconds INT DEFAULT 0;
|
||||
ALTER TABLE channels ADD COLUMN last_reset_time BIGINT DEFAULT 0;
|
||||
ALTER TABLE channels ADD COLUMN next_reset_time BIGINT DEFAULT 0;
|
||||
```
|
||||
|
||||
`used_quota` 已存在，沿用。现有渠道 `quota_reset_period = 'never'` 即为原有行为。
|
||||
|
||||
---
|
||||
|
||||
## 5. 数据流
|
||||
|
||||
```mermaid
|
||||
sequenceDiagram
|
||||
participant Req as 请求
|
||||
participant Relay as Relay 层 (text_quota.go)
|
||||
participant Model as channel_quota.go
|
||||
participant DB as Database
|
||||
|
||||
Req->>Relay: 请求完成，准备扣减
|
||||
Relay->>Model: 懒重置检查 (channelId)
|
||||
Model->>DB: SELECT ... FOR UPDATE
|
||||
alt NextResetTime <= now
|
||||
Model->>DB: UPDATE used_quota=0, last/next_reset_time
|
||||
end
|
||||
Relay->>Model: 额度检查 (channelId)
|
||||
alt UsedQuota >= QuotaLimit (且 QuotaLimit > 0)
|
||||
Relay-->>Req: 429 Channel quota exceeded
|
||||
else
|
||||
Relay->>DB: UPDATE used_quota += actual
|
||||
end
|
||||
```
|
||||
|
||||
```mermaid
|
||||
sequenceDiagram
|
||||
participant Cron as 定时任务 (每分钟)
|
||||
participant Model as channel_quota.go
|
||||
participant DB as Database
|
||||
|
||||
Cron->>Model: ResetDueChannels(300)
|
||||
Model->>DB: SELECT next_reset_time <= now AND quota_limit > 0
|
||||
loop 每条渠道
|
||||
Model->>DB: FOR UPDATE 锁行
|
||||
Model->>DB: UPDATE used_quota=0, last/next_reset_time
|
||||
end
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 6. 风险与兼容性
|
||||
|
||||
| 风险 | 影响 | 缓解措施 |
|
||||
|------|------|----------|
|
||||
| `quota_reset_period = 'never'` 时 `UsedQuota` 行为不变 | 无 | 默认值即为此 |
|
||||
| 历史渠道 `next_reset_time = 0` | 无 | 定时任务跳过，懒检查也跳过 |
|
||||
| 重置时刻有进行中请求 | 低 | `FOR UPDATE` 锁保证串行 |
|
||||
| 最后一个放行请求可能超用 | 低 | 可接受，与现有行为一致 |
|
||||
| SQLite 下 `FOR UPDATE` 无效 | 低 | SQLite 自带数据库级锁，事务串行化 |
|
||||
| 缓存中 `UsedQuota` 可能过期 | 低 | 额度检查直接查 DB，不依赖缓存 |
|
||||
|
||||
---
|
||||
|
||||
## 7. 验收标准
|
||||
|
||||
### 7.1 功能验收
|
||||
|
||||
- [ ] 新增/编辑渠道时可配置 `quota_limit` 和 `quota_reset_period`
|
||||
- [ ] `QuotaLimit > 0` 且 `UsedQuota >= QuotaLimit` 时请求返回 429
|
||||
- [ ] `QuotaLimit = 0` 时不做额度检查（原有行为）
|
||||
- [ ] 每日重置：次日 00:00 `UsedQuota` 清零
|
||||
- [ ] 每月重置：每月 1 日 00:00 清零
|
||||
- [ ] 自定义秒数重置：按配置精确重置
|
||||
- [ ] 永不重置：`UsedQuota` 持续累计
|
||||
- [ ] 创建/更新渠道时正确初始化 `NextResetTime`
|
||||
|
||||
### 7.2 兼容性验收
|
||||
|
||||
- [ ] 未设置 `QuotaLimit` 的历史渠道行为不变
|
||||
- [ ] GORM AutoMigrate 正确添加新字段
|
||||
- [ ] SQLite / MySQL / PostgreSQL 均兼容
|
||||
|
||||
---
|
||||
|
||||
## 8. 回滚与降级
|
||||
|
||||
### 8.1 回滚
|
||||
|
||||
部署后如需回滚：新字段由 GORM AutoMigrate 添加，旧代码忽略未知字段，向前兼容。回滚后定时任务不运行，`UsedQuota` 恢复持续累计行为。
|
||||
|
||||
### 8.2 降级
|
||||
|
||||
关闭单个渠道的额度管控：设 `QuotaLimit = 0` + `QuotaResetPeriod = 'never'`，等价于原有行为。
|
||||
|
||||
---
|
||||
|
||||
## 9. 参考实现
|
||||
|
||||
- `model/subscription.go:300-307` — `NormalizeResetPeriod`
|
||||
- `model/subscription.go:309-348` — `calcNextResetTime`
|
||||
- `model/subscription.go:933-967` — `maybeResetUserSubscriptionWithPlanTx`
|
||||
- `model/subscription.go:1100-1140` — `ResetDueSubscriptions`
|
||||
- `service/subscription_reset_task.go` — 定时任务
|
||||
2
main.go
2
main.go
@ -118,6 +118,8 @@ func main() {
|
||||
|
||||
// Subscription quota reset task (daily/weekly/monthly/custom)
|
||||
service.StartSubscriptionQuotaResetTask()
|
||||
// Channel quota reset task (daily/weekly/monthly/custom)
|
||||
service.StartChannelQuotaResetTask()
|
||||
|
||||
// Wire task polling adaptor factory (breaks service -> relay import cycle)
|
||||
service.GetTaskAdaptorFunc = func(platform constant.TaskPlatform) service.TaskPollingAdaptor {
|
||||
|
||||
@ -157,6 +157,12 @@ func Distribute() func(c *gin.Context) {
|
||||
}
|
||||
common.SetContextKey(c, constant.ContextKeyRequestStartTime, time.Now())
|
||||
SetupContextForSelectedChannel(c, channel, modelRequest.Model)
|
||||
if channel != nil {
|
||||
if err := model.CheckChannelQuota(channel.Id); err != nil {
|
||||
abortWithOpenAiMessage(c, http.StatusTooManyRequests, err.Error())
|
||||
return
|
||||
}
|
||||
}
|
||||
c.Next()
|
||||
if channel != nil && c.Writer != nil && c.Writer.Status() < http.StatusBadRequest {
|
||||
service.RecordChannelAffinity(c, channel.Id)
|
||||
|
||||
@ -20,25 +20,30 @@ import (
|
||||
)
|
||||
|
||||
type Channel struct {
|
||||
Id int `json:"id"`
|
||||
Type int `json:"type" gorm:"default:0"`
|
||||
Key string `json:"key" gorm:"not null"`
|
||||
OpenAIOrganization *string `json:"openai_organization"`
|
||||
TestModel *string `json:"test_model"`
|
||||
Status int `json:"status" gorm:"default:1"`
|
||||
Name string `json:"name" gorm:"index"`
|
||||
Weight *uint `json:"weight" gorm:"default:0"`
|
||||
CreatedTime int64 `json:"created_time" gorm:"bigint"`
|
||||
TestTime int64 `json:"test_time" gorm:"bigint"`
|
||||
ResponseTime int `json:"response_time"` // in milliseconds
|
||||
BaseURL *string `json:"base_url" gorm:"column:base_url;default:''"`
|
||||
Other string `json:"other"`
|
||||
Balance float64 `json:"balance"` // in USD
|
||||
BalanceUpdatedTime int64 `json:"balance_updated_time" gorm:"bigint"`
|
||||
Models string `json:"models"`
|
||||
Group string `json:"group" gorm:"type:varchar(64);default:'default'"`
|
||||
UsedQuota int64 `json:"used_quota" gorm:"bigint;default:0"`
|
||||
ModelMapping *string `json:"model_mapping" gorm:"type:text"`
|
||||
Id int `json:"id"`
|
||||
Type int `json:"type" gorm:"default:0"`
|
||||
Key string `json:"key" gorm:"not null"`
|
||||
OpenAIOrganization *string `json:"openai_organization"`
|
||||
TestModel *string `json:"test_model"`
|
||||
Status int `json:"status" gorm:"default:1"`
|
||||
Name string `json:"name" gorm:"index"`
|
||||
Weight *uint `json:"weight" gorm:"default:0"`
|
||||
CreatedTime int64 `json:"created_time" gorm:"bigint"`
|
||||
TestTime int64 `json:"test_time" gorm:"bigint"`
|
||||
ResponseTime int `json:"response_time"` // in milliseconds
|
||||
BaseURL *string `json:"base_url" gorm:"column:base_url;default:''"`
|
||||
Other string `json:"other"`
|
||||
Balance float64 `json:"balance"` // in USD
|
||||
BalanceUpdatedTime int64 `json:"balance_updated_time" gorm:"bigint"`
|
||||
Models string `json:"models"`
|
||||
Group string `json:"group" gorm:"type:varchar(64);default:'default'"`
|
||||
UsedQuota int64 `json:"used_quota" gorm:"bigint;default:0"`
|
||||
QuotaLimit int64 `json:"quota_limit" gorm:"bigint;default:0"`
|
||||
QuotaResetPeriod string `json:"quota_reset_period" gorm:"type:varchar(16);default:'never'"`
|
||||
QuotaResetCustomSeconds int `json:"quota_reset_custom_seconds" gorm:"type:int;default:0"`
|
||||
LastResetTime int64 `json:"last_reset_time" gorm:"bigint;default:0"`
|
||||
NextResetTime int64 `json:"next_reset_time" gorm:"bigint;default:0"`
|
||||
ModelMapping *string `json:"model_mapping" gorm:"type:text"`
|
||||
//MaxInputTokens *int `json:"max_input_tokens" gorm:"default:0"`
|
||||
StatusCodeMapping *string `json:"status_code_mapping" gorm:"type:varchar(1024);default:''"`
|
||||
Priority *int64 `json:"priority" gorm:"bigint;default:0"`
|
||||
|
||||
200
model/channel_quota.go
Normal file
200
model/channel_quota.go
Normal file
@ -0,0 +1,200 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
const (
|
||||
ChannelResetNever = "never"
|
||||
ChannelResetDaily = "daily"
|
||||
ChannelResetWeekly = "weekly"
|
||||
ChannelResetMonthly = "monthly"
|
||||
ChannelResetCustom = "custom"
|
||||
)
|
||||
|
||||
func NormalizeChannelResetPeriod(period string) string {
|
||||
switch strings.TrimSpace(period) {
|
||||
case ChannelResetDaily, ChannelResetWeekly, ChannelResetMonthly, ChannelResetCustom:
|
||||
return strings.TrimSpace(period)
|
||||
default:
|
||||
return ChannelResetNever
|
||||
}
|
||||
}
|
||||
|
||||
func calcChannelNextResetTime(base time.Time, period string, customSeconds int) int64 {
|
||||
period = NormalizeChannelResetPeriod(period)
|
||||
if period == ChannelResetNever {
|
||||
return 0
|
||||
}
|
||||
var next time.Time
|
||||
switch period {
|
||||
case ChannelResetDaily:
|
||||
next = time.Date(base.Year(), base.Month(), base.Day(), 0, 0, 0, 0, base.Location()).
|
||||
AddDate(0, 0, 1)
|
||||
case ChannelResetWeekly:
|
||||
weekday := int(base.Weekday())
|
||||
if weekday == 0 {
|
||||
weekday = 7
|
||||
}
|
||||
daysUntil := 8 - weekday
|
||||
next = time.Date(base.Year(), base.Month(), base.Day(), 0, 0, 0, 0, base.Location()).
|
||||
AddDate(0, 0, daysUntil)
|
||||
case ChannelResetMonthly:
|
||||
next = time.Date(base.Year(), base.Month(), 1, 0, 0, 0, 0, base.Location()).
|
||||
AddDate(0, 1, 0)
|
||||
case ChannelResetCustom:
|
||||
if customSeconds <= 0 {
|
||||
return 0
|
||||
}
|
||||
next = base.Add(time.Duration(customSeconds) * time.Second)
|
||||
default:
|
||||
return 0
|
||||
}
|
||||
return next.Unix()
|
||||
}
|
||||
|
||||
func maybeResetChannelQuotaTx(tx *gorm.DB, channel *Channel, now int64) error {
|
||||
if tx == nil || channel == nil {
|
||||
return errors.New("invalid reset args")
|
||||
}
|
||||
if channel.NextResetTime > 0 && channel.NextResetTime > now {
|
||||
return nil
|
||||
}
|
||||
if NormalizeChannelResetPeriod(channel.QuotaResetPeriod) == ChannelResetNever {
|
||||
return nil
|
||||
}
|
||||
baseUnix := channel.LastResetTime
|
||||
if baseUnix <= 0 {
|
||||
baseUnix = channel.CreatedTime
|
||||
}
|
||||
base := time.Unix(baseUnix, 0)
|
||||
next := calcChannelNextResetTime(base, channel.QuotaResetPeriod, channel.QuotaResetCustomSeconds)
|
||||
|
||||
advanced := false
|
||||
for next > 0 && next <= now {
|
||||
advanced = true
|
||||
base = time.Unix(next, 0)
|
||||
next = calcChannelNextResetTime(base, channel.QuotaResetPeriod, channel.QuotaResetCustomSeconds)
|
||||
}
|
||||
|
||||
if !advanced {
|
||||
if channel.NextResetTime == 0 && next > 0 {
|
||||
return tx.Model(&Channel{}).Where("id = ?", channel.Id).Updates(map[string]interface{}{
|
||||
"next_reset_time": next,
|
||||
"last_reset_time": base.Unix(),
|
||||
}).Error
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
return tx.Model(&Channel{}).Where("id = ?", channel.Id).Updates(map[string]interface{}{
|
||||
"used_quota": 0,
|
||||
"last_reset_time": base.Unix(),
|
||||
"next_reset_time": next,
|
||||
}).Error
|
||||
}
|
||||
|
||||
func ResetDueChannels(limit int) (int, error) {
|
||||
if limit <= 0 {
|
||||
limit = 200
|
||||
}
|
||||
now := GetDBTimestamp()
|
||||
var channels []Channel
|
||||
if err := DB.Where("next_reset_time > 0 AND next_reset_time <= ? AND status = ? AND quota_limit > 0",
|
||||
now, common.ChannelStatusEnabled).
|
||||
Order("next_reset_time asc").
|
||||
Limit(limit).
|
||||
Find(&channels).Error; err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if len(channels) == 0 {
|
||||
return 0, nil
|
||||
}
|
||||
resetCount := 0
|
||||
for _, ch := range channels {
|
||||
chCopy := ch
|
||||
err := DB.Transaction(func(tx *gorm.DB) error {
|
||||
var locked Channel
|
||||
if err := tx.Set("gorm:query_option", "FOR UPDATE").
|
||||
Where("id = ? AND next_reset_time > 0 AND next_reset_time <= ?", chCopy.Id, now).
|
||||
First(&locked).Error; err != nil {
|
||||
return nil
|
||||
}
|
||||
if err := maybeResetChannelQuotaTx(tx, &locked, now); err != nil {
|
||||
return err
|
||||
}
|
||||
resetCount++
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
return resetCount, err
|
||||
}
|
||||
}
|
||||
return resetCount, nil
|
||||
}
|
||||
|
||||
func CheckChannelQuota(channelId int) error {
|
||||
var channel Channel
|
||||
if err := DB.Select("id, used_quota, quota_limit, quota_reset_period, next_reset_time").
|
||||
Where("id = ?", channelId).First(&channel).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if channel.QuotaLimit <= 0 {
|
||||
return nil
|
||||
}
|
||||
if channel.NextResetTime > 0 {
|
||||
now := GetDBTimestamp()
|
||||
if channel.NextResetTime <= now {
|
||||
_ = DB.Transaction(func(tx *gorm.DB) error {
|
||||
var locked Channel
|
||||
if err := tx.Set("gorm:query_option", "FOR UPDATE").
|
||||
Where("id = ? AND next_reset_time > 0 AND next_reset_time <= ?", channelId, now).
|
||||
First(&locked).Error; err != nil {
|
||||
return nil
|
||||
}
|
||||
return maybeResetChannelQuotaTx(tx, &locked, now)
|
||||
})
|
||||
if err := DB.Select("used_quota, quota_limit").
|
||||
Where("id = ?", channelId).First(&channel).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
if channel.UsedQuota >= channel.QuotaLimit {
|
||||
return fmt.Errorf("channel quota exceeded (used=%d, limit=%d)", channel.UsedQuota, channel.QuotaLimit)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func InitChannelNextResetTime(channel *Channel) {
|
||||
if channel.QuotaLimit <= 0 {
|
||||
channel.NextResetTime = 0
|
||||
channel.LastResetTime = 0
|
||||
return
|
||||
}
|
||||
period := NormalizeChannelResetPeriod(channel.QuotaResetPeriod)
|
||||
if period == ChannelResetNever {
|
||||
channel.NextResetTime = 0
|
||||
channel.LastResetTime = 0
|
||||
return
|
||||
}
|
||||
baseUnix := channel.LastResetTime
|
||||
if baseUnix <= 0 {
|
||||
baseUnix = channel.CreatedTime
|
||||
}
|
||||
if baseUnix <= 0 {
|
||||
baseUnix = GetDBTimestamp()
|
||||
}
|
||||
base := time.Unix(baseUnix, 0)
|
||||
next := calcChannelNextResetTime(base, channel.QuotaResetPeriod, channel.QuotaResetCustomSeconds)
|
||||
channel.NextResetTime = next
|
||||
if channel.LastResetTime <= 0 {
|
||||
channel.LastResetTime = baseUnix
|
||||
}
|
||||
}
|
||||
70
service/channel_quota_reset_task.go
Normal file
70
service/channel_quota_reset_task.go
Normal file
@ -0,0 +1,70 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/logger"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
|
||||
"github.com/bytedance/gopkg/util/gopool"
|
||||
)
|
||||
|
||||
const (
|
||||
channelQuotaResetTickInterval = 1 * time.Minute
|
||||
channelQuotaResetBatchSize = 300
|
||||
)
|
||||
|
||||
var (
|
||||
channelQuotaResetOnce sync.Once
|
||||
channelQuotaResetRunning atomic.Bool
|
||||
)
|
||||
|
||||
func StartChannelQuotaResetTask() {
|
||||
channelQuotaResetOnce.Do(func() {
|
||||
if !common.IsMasterNode {
|
||||
return
|
||||
}
|
||||
gopool.Go(func() {
|
||||
logger.LogInfo(context.Background(), fmt.Sprintf("channel quota reset task started: tick=%s", channelQuotaResetTickInterval))
|
||||
ticker := time.NewTicker(channelQuotaResetTickInterval)
|
||||
defer ticker.Stop()
|
||||
|
||||
runChannelQuotaResetOnce()
|
||||
for range ticker.C {
|
||||
runChannelQuotaResetOnce()
|
||||
}
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
func runChannelQuotaResetOnce() {
|
||||
if !channelQuotaResetRunning.CompareAndSwap(false, true) {
|
||||
return
|
||||
}
|
||||
defer channelQuotaResetRunning.Store(false)
|
||||
|
||||
ctx := context.Background()
|
||||
totalReset := 0
|
||||
for {
|
||||
n, err := model.ResetDueChannels(channelQuotaResetBatchSize)
|
||||
if err != nil {
|
||||
logger.LogWarn(ctx, fmt.Sprintf("channel quota reset task failed: %v", err))
|
||||
return
|
||||
}
|
||||
if n == 0 {
|
||||
break
|
||||
}
|
||||
totalReset += n
|
||||
if n < channelQuotaResetBatchSize {
|
||||
break
|
||||
}
|
||||
}
|
||||
if common.DebugEnabled && totalReset > 0 {
|
||||
logger.LogDebug(ctx, "channel quota maintenance: reset_count=%d", totalReset)
|
||||
}
|
||||
}
|
||||
@ -368,6 +368,9 @@ func PostTextConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, us
|
||||
} else {
|
||||
model.UpdateUserUsedQuotaAndRequestCount(relayInfo.UserId, summary.Quota)
|
||||
model.UpdateChannelUsedQuota(relayInfo.ChannelId, summary.Quota)
|
||||
if err := model.CheckChannelQuota(relayInfo.ChannelId); err != nil {
|
||||
logger.LogWarn(ctx, fmt.Sprintf("channel quota exceeded: channel_id=%d, %v", relayInfo.ChannelId, err))
|
||||
}
|
||||
}
|
||||
|
||||
if err := SettleBilling(ctx, relayInfo, summary.Quota); err != nil {
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user