mirror of
https://github.com/router-for-me/CLIProxyAPI.git
synced 2026-02-19 04:40:52 +08:00
feat: add DisableCooling configuration to manage quota cooldown behavior
This commit is contained in:
@@ -27,6 +27,7 @@ import (
|
|||||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/usage"
|
"github.com/router-for-me/CLIProxyAPI/v6/internal/usage"
|
||||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
|
"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
|
||||||
sdkAuth "github.com/router-for-me/CLIProxyAPI/v6/sdk/auth"
|
sdkAuth "github.com/router-for-me/CLIProxyAPI/v6/sdk/auth"
|
||||||
|
coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
|
||||||
log "github.com/sirupsen/logrus"
|
log "github.com/sirupsen/logrus"
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -377,6 +378,7 @@ func main() {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
usage.SetStatisticsEnabled(cfg.UsageStatisticsEnabled)
|
usage.SetStatisticsEnabled(cfg.UsageStatisticsEnabled)
|
||||||
|
coreauth.SetQuotaCooldownDisabled(cfg.DisableCooling)
|
||||||
|
|
||||||
if err = logging.ConfigureLogOutput(cfg.LoggingToFile); err != nil {
|
if err = logging.ConfigureLogOutput(cfg.LoggingToFile); err != nil {
|
||||||
log.Fatalf("failed to configure log output: %v", err)
|
log.Fatalf("failed to configure log output: %v", err)
|
||||||
|
|||||||
@@ -233,6 +233,7 @@ func NewServer(cfg *config.Config, authManager *auth.Manager, accessManager *sdk
|
|||||||
s.oldConfigYaml, _ = yaml.Marshal(cfg)
|
s.oldConfigYaml, _ = yaml.Marshal(cfg)
|
||||||
s.applyAccessConfig(nil, cfg)
|
s.applyAccessConfig(nil, cfg)
|
||||||
managementasset.SetCurrentConfig(cfg)
|
managementasset.SetCurrentConfig(cfg)
|
||||||
|
auth.SetQuotaCooldownDisabled(cfg.DisableCooling)
|
||||||
// Initialize management handler
|
// Initialize management handler
|
||||||
s.mgmt = managementHandlers.NewHandler(cfg, configFilePath, authManager)
|
s.mgmt = managementHandlers.NewHandler(cfg, configFilePath, authManager)
|
||||||
if optionState.localPassword != "" {
|
if optionState.localPassword != "" {
|
||||||
@@ -716,6 +717,15 @@ func (s *Server) UpdateClients(cfg *config.Config) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
if oldCfg == nil || oldCfg.DisableCooling != cfg.DisableCooling {
|
||||||
|
auth.SetQuotaCooldownDisabled(cfg.DisableCooling)
|
||||||
|
if oldCfg != nil {
|
||||||
|
log.Debugf("disable_cooling updated from %t to %t", oldCfg.DisableCooling, cfg.DisableCooling)
|
||||||
|
} else {
|
||||||
|
log.Debugf("disable_cooling toggled to %t", cfg.DisableCooling)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
// Update log level dynamically when debug flag changes
|
// Update log level dynamically when debug flag changes
|
||||||
if oldCfg == nil || oldCfg.Debug != cfg.Debug {
|
if oldCfg == nil || oldCfg.Debug != cfg.Debug {
|
||||||
util.SetLogLevel(cfg)
|
util.SetLogLevel(cfg)
|
||||||
|
|||||||
@@ -34,6 +34,9 @@ type Config struct {
|
|||||||
// UsageStatisticsEnabled toggles in-memory usage aggregation; when false, usage data is discarded.
|
// UsageStatisticsEnabled toggles in-memory usage aggregation; when false, usage data is discarded.
|
||||||
UsageStatisticsEnabled bool `yaml:"usage-statistics-enabled" json:"usage-statistics-enabled"`
|
UsageStatisticsEnabled bool `yaml:"usage-statistics-enabled" json:"usage-statistics-enabled"`
|
||||||
|
|
||||||
|
// DisableCooling disables quota cooldown scheduling when true.
|
||||||
|
DisableCooling bool `yaml:"disable-cooling" json:"disable-cooling"`
|
||||||
|
|
||||||
// QuotaExceeded defines the behavior when a quota is exceeded.
|
// QuotaExceeded defines the behavior when a quota is exceeded.
|
||||||
QuotaExceeded QuotaExceeded `yaml:"quota-exceeded" json:"quota-exceeded"`
|
QuotaExceeded QuotaExceeded `yaml:"quota-exceeded" json:"quota-exceeded"`
|
||||||
|
|
||||||
@@ -183,6 +186,7 @@ func LoadConfigOptional(configFile string, optional bool) (*Config, error) {
|
|||||||
// Set defaults before unmarshal so that absent keys keep defaults.
|
// Set defaults before unmarshal so that absent keys keep defaults.
|
||||||
cfg.LoggingToFile = false
|
cfg.LoggingToFile = false
|
||||||
cfg.UsageStatisticsEnabled = false
|
cfg.UsageStatisticsEnabled = false
|
||||||
|
cfg.DisableCooling = false
|
||||||
if err = yaml.Unmarshal(data, &cfg); err != nil {
|
if err = yaml.Unmarshal(data, &cfg); err != nil {
|
||||||
if optional {
|
if optional {
|
||||||
// In cloud deploy mode, if YAML parsing fails, return empty config instead of error.
|
// In cloud deploy mode, if YAML parsing fails, return empty config instead of error.
|
||||||
|
|||||||
@@ -1192,6 +1192,9 @@ func buildConfigChangeDetails(oldCfg, newCfg *config.Config) []string {
|
|||||||
if oldCfg.UsageStatisticsEnabled != newCfg.UsageStatisticsEnabled {
|
if oldCfg.UsageStatisticsEnabled != newCfg.UsageStatisticsEnabled {
|
||||||
changes = append(changes, fmt.Sprintf("usage-statistics-enabled: %t -> %t", oldCfg.UsageStatisticsEnabled, newCfg.UsageStatisticsEnabled))
|
changes = append(changes, fmt.Sprintf("usage-statistics-enabled: %t -> %t", oldCfg.UsageStatisticsEnabled, newCfg.UsageStatisticsEnabled))
|
||||||
}
|
}
|
||||||
|
if oldCfg.DisableCooling != newCfg.DisableCooling {
|
||||||
|
changes = append(changes, fmt.Sprintf("disable-cooling: %t -> %t", oldCfg.DisableCooling, newCfg.DisableCooling))
|
||||||
|
}
|
||||||
if oldCfg.RequestLog != newCfg.RequestLog {
|
if oldCfg.RequestLog != newCfg.RequestLog {
|
||||||
changes = append(changes, fmt.Sprintf("request-log: %t -> %t", oldCfg.RequestLog, newCfg.RequestLog))
|
changes = append(changes, fmt.Sprintf("request-log: %t -> %t", oldCfg.RequestLog, newCfg.RequestLog))
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -8,6 +8,7 @@ import (
|
|||||||
"strconv"
|
"strconv"
|
||||||
"strings"
|
"strings"
|
||||||
"sync"
|
"sync"
|
||||||
|
"sync/atomic"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
"github.com/google/uuid"
|
"github.com/google/uuid"
|
||||||
@@ -44,6 +45,13 @@ const (
|
|||||||
quotaBackoffMax = 30 * time.Minute
|
quotaBackoffMax = 30 * time.Minute
|
||||||
)
|
)
|
||||||
|
|
||||||
|
var quotaCooldownDisabled atomic.Bool
|
||||||
|
|
||||||
|
// SetQuotaCooldownDisabled toggles quota cooldown scheduling globally.
|
||||||
|
func SetQuotaCooldownDisabled(disable bool) {
|
||||||
|
quotaCooldownDisabled.Store(disable)
|
||||||
|
}
|
||||||
|
|
||||||
// Result captures execution outcome used to adjust auth state.
|
// Result captures execution outcome used to adjust auth state.
|
||||||
type Result struct {
|
type Result struct {
|
||||||
// AuthID references the auth that produced this result.
|
// AuthID references the auth that produced this result.
|
||||||
@@ -535,7 +543,10 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) {
|
|||||||
shouldSuspendModel = true
|
shouldSuspendModel = true
|
||||||
case 429:
|
case 429:
|
||||||
cooldown, nextLevel := nextQuotaCooldown(state.Quota.BackoffLevel)
|
cooldown, nextLevel := nextQuotaCooldown(state.Quota.BackoffLevel)
|
||||||
next := now.Add(cooldown)
|
var next time.Time
|
||||||
|
if cooldown > 0 {
|
||||||
|
next = now.Add(cooldown)
|
||||||
|
}
|
||||||
state.NextRetryAfter = next
|
state.NextRetryAfter = next
|
||||||
state.Quota = QuotaState{
|
state.Quota = QuotaState{
|
||||||
Exceeded: true,
|
Exceeded: true,
|
||||||
@@ -750,9 +761,13 @@ func applyAuthFailureState(auth *Auth, resultErr *Error, now time.Time) {
|
|||||||
auth.Quota.Exceeded = true
|
auth.Quota.Exceeded = true
|
||||||
auth.Quota.Reason = "quota"
|
auth.Quota.Reason = "quota"
|
||||||
cooldown, nextLevel := nextQuotaCooldown(auth.Quota.BackoffLevel)
|
cooldown, nextLevel := nextQuotaCooldown(auth.Quota.BackoffLevel)
|
||||||
auth.Quota.NextRecoverAt = now.Add(cooldown)
|
var next time.Time
|
||||||
|
if cooldown > 0 {
|
||||||
|
next = now.Add(cooldown)
|
||||||
|
}
|
||||||
|
auth.Quota.NextRecoverAt = next
|
||||||
auth.Quota.BackoffLevel = nextLevel
|
auth.Quota.BackoffLevel = nextLevel
|
||||||
auth.NextRetryAfter = auth.Quota.NextRecoverAt
|
auth.NextRetryAfter = next
|
||||||
case 408, 500, 502, 503, 504:
|
case 408, 500, 502, 503, 504:
|
||||||
auth.StatusMessage = "transient upstream error"
|
auth.StatusMessage = "transient upstream error"
|
||||||
auth.NextRetryAfter = now.Add(1 * time.Minute)
|
auth.NextRetryAfter = now.Add(1 * time.Minute)
|
||||||
@@ -768,6 +783,9 @@ func nextQuotaCooldown(prevLevel int) (time.Duration, int) {
|
|||||||
if prevLevel < 0 {
|
if prevLevel < 0 {
|
||||||
prevLevel = 0
|
prevLevel = 0
|
||||||
}
|
}
|
||||||
|
if quotaCooldownDisabled.Load() {
|
||||||
|
return 0, prevLevel
|
||||||
|
}
|
||||||
cooldown := quotaBackoffBase * time.Duration(1<<prevLevel)
|
cooldown := quotaBackoffBase * time.Duration(1<<prevLevel)
|
||||||
if cooldown < quotaBackoffBase {
|
if cooldown < quotaBackoffBase {
|
||||||
cooldown = quotaBackoffBase
|
cooldown = quotaBackoffBase
|
||||||
|
|||||||
Reference in New Issue
Block a user