fork-it/gpt-load
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: 优化

Browse Source
This commit is contained in:
tbphp
2025-06-06 21:56:40 +08:00
parent 219c068dbf
commit 1faa0b1b73
5 changed files with 193 additions and 156 deletions
Download Patch File Download Diff File Expand all files Collapse all files

181
internal/keymanager/keymanager.go
View File

@@ -27,13 +27,13 @@ type KeyInfo struct {
// Stats 统计信息
type Stats struct {
CurrentIndex int64 `json:"currentIndex"`
TotalKeys int `json:"totalKeys"`
HealthyKeys int `json:"healthyKeys"`
BlacklistedKeys int `json:"blacklistedKeys"`
SuccessCount int64 `json:"successCount"`
FailureCount int64 `json:"failureCount"`
MemoryUsage MemoryUsage `json:"memoryUsage"`
CurrentIndex int64 `json:"currentIndex"`
TotalKeys int `json:"totalKeys"`
HealthyKeys int `json:"healthyKeys"`
BlacklistedKeys int `json:"blacklistedKeys"`
SuccessCount int64 `json:"successCount"`
FailureCount int64 `json:"failureCount"`
MemoryUsage MemoryUsage `json:"memoryUsage"`
}
// MemoryUsage 内存使用情况
@@ -60,22 +60,22 @@ type BlacklistInfo struct {
// KeyManager 密钥管理器
type KeyManager struct {
keysFilePath string
keys []string
keyPreviews []string
currentIndex int64
blacklistedKeys sync.Map
successCount int64
failureCount int64
keysFilePath string
keys []string
keyPreviews []string
currentIndex int64
blacklistedKeys sync.Map
successCount int64
failureCount int64
keyFailureCounts sync.Map
// 性能优化：预编译正则表达式
permanentErrorPatterns []*regexp.Regexp
// 内存管理
cleanupTicker *time.Ticker
stopCleanup chan bool
// 读写锁保护密钥列表
keysMutex sync.RWMutex
}
@@ -85,12 +85,12 @@ func NewKeyManager(keysFilePath string) *KeyManager {
if keysFilePath == "" {
keysFilePath = config.AppConfig.Keys.FilePath
}
km := &KeyManager{
keysFilePath: keysFilePath,
currentIndex: int64(config.AppConfig.Keys.StartIndex),
stopCleanup: make(chan bool),
// 预编译正则表达式
permanentErrorPatterns: []*regexp.Regexp{
regexp.MustCompile(`(?i)invalid api key`),
@@ -101,10 +101,10 @@ func NewKeyManager(keysFilePath string) *KeyManager {
regexp.MustCompile(`(?i)billing`),
},
}
// 启动内存清理
km.setupMemoryCleanup()
return km
}
@@ -115,25 +115,25 @@ func (km *KeyManager) LoadKeys() error {
return fmt.Errorf("无法打开密钥文件: %w", err)
}
defer file.Close()
var keys []string
scanner := bufio.NewScanner(file)
for scanner.Scan() {
line := strings.TrimSpace(scanner.Text())
if line != "" && strings.HasPrefix(line, "sk-") {
keys = append(keys, line)
}
}
if err := scanner.Err(); err != nil {
return fmt.Errorf("读取密钥文件失败: %w", err)
}
if len(keys) == 0 {
return fmt.Errorf("密钥文件中没有有效的API密钥")
}
km.keysMutex.Lock()
km.keys = keys
// 预生成密钥预览，避免运行时重复计算
@@ -146,7 +146,7 @@ func (km *KeyManager) LoadKeys() error {
}
}
km.keysMutex.Unlock()
logrus.Infof("✅ 成功加载 %d 个 API 密钥", len(keys))
return nil
}
@@ -155,36 +155,40 @@ func (km *KeyManager) LoadKeys() error {
func (km *KeyManager) GetNextKey() (*KeyInfo, error) {
km.keysMutex.RLock()
keysLen := len(km.keys)
km.keysMutex.RUnlock()
if keysLen == 0 {
km.keysMutex.RUnlock()
return nil, fmt.Errorf("没有可用的 API 密钥")
}
// 检查是否所有密钥都被拉黑
blacklistedCount := 0
km.blacklistedKeys.Range(func(key, value interface{}) bool {
blacklistedCount++
return true
})
if blacklistedCount >= keysLen {
logrus.Warn("⚠️ 所有密钥都被拉黑，重置黑名单")
km.blacklistedKeys = sync.Map{}
km.keyFailureCounts = sync.Map{}
// 快速路径：直接获取下一个密钥，避免黑名单检查的开销
currentIdx := atomic.AddInt64(&km.currentIndex, 1) - 1
keyIndex := int(currentIdx) % keysLen
selectedKey := km.keys[keyIndex]
keyPreview := km.keyPreviews[keyIndex]
km.keysMutex.RUnlock()
// 检查是否被拉黑
if _, blacklisted := km.blacklistedKeys.Load(selectedKey); !blacklisted {
return &KeyInfo{
Key: selectedKey,
Index: keyIndex,
Preview: keyPreview,
}, nil
}
// 使用原子操作避免竞态条件
// 慢速路径：寻找可用密钥
attempts := 0
for attempts < keysLen {
currentIdx := atomic.AddInt64(&km.currentIndex, 1) - 1
keyIndex := int(currentIdx) % keysLen
maxAttempts := keysLen * 2 // 最多尝试两轮
for attempts < maxAttempts {
currentIdx = atomic.AddInt64(&km.currentIndex, 1) - 1
keyIndex = int(currentIdx) % keysLen
km.keysMutex.RLock()
selectedKey := km.keys[keyIndex]
keyPreview := km.keyPreviews[keyIndex]
selectedKey = km.keys[keyIndex]
keyPreview = km.keyPreviews[keyIndex]
km.keysMutex.RUnlock()
if _, blacklisted := km.blacklistedKeys.Load(selectedKey); !blacklisted {
return &KeyInfo{
Key: selectedKey,
@@ -192,21 +196,36 @@ func (km *KeyManager) GetNextKey() (*KeyInfo, error) {
Preview: keyPreview,
}, nil
}
attempts++
}
// 兜底：返回第一个密钥
km.keysMutex.RLock()
firstKey := km.keys[0]
firstPreview := km.keyPreviews[0]
km.keysMutex.RUnlock()
return &KeyInfo{
Key: firstKey,
Index: 0,
Preview: firstPreview,
}, nil
// 检查是否所有密钥都被拉黑，如果是则重置
blacklistedCount := 0
km.blacklistedKeys.Range(func(key, value interface{}) bool {
blacklistedCount++
return blacklistedCount < keysLen // 提前退出优化
})
if blacklistedCount >= keysLen {
logrus.Warn("⚠️ 所有密钥都被拉黑，重置黑名单")
km.blacklistedKeys = sync.Map{}
km.keyFailureCounts = sync.Map{}
// 重置后返回第一个密钥
km.keysMutex.RLock()
firstKey := km.keys[0]
firstPreview := km.keyPreviews[0]
km.keysMutex.RUnlock()
return &KeyInfo{
Key: firstKey,
Index: 0,
Preview: firstPreview,
}, nil
}
return nil, fmt.Errorf("暂时没有可用的 API 密钥")
}
// RecordSuccess 记录密钥使用成功
@@ -219,7 +238,7 @@ func (km *KeyManager) RecordSuccess(key string) {
// RecordFailure 记录密钥使用失败
func (km *KeyManager) RecordFailure(key string, err error) {
atomic.AddInt64(&km.failureCount, 1)
// 检查是否是永久性错误
if km.isPermanentError(err) {
km.blacklistedKeys.Store(key, true)
@@ -227,7 +246,7 @@ func (km *KeyManager) RecordFailure(key string, err error) {
logrus.Warnf("🚫 密钥已被拉黑（永久性错误）: %s (%s)", key[:20]+"...", err.Error())
return
}
// 临时性错误：增加失败计数
currentFailures := 0
if val, exists := km.keyFailureCounts.Load(key); exists {
@@ -235,7 +254,7 @@ func (km *KeyManager) RecordFailure(key string, err error) {
}
newFailures := currentFailures + 1
km.keyFailureCounts.Store(key, newFailures)
threshold := config.AppConfig.Keys.BlacklistThreshold
if newFailures >= threshold {
km.blacklistedKeys.Store(key, true)
@@ -262,19 +281,19 @@ func (km *KeyManager) GetStats() *Stats {
km.keysMutex.RLock()
totalKeys := len(km.keys)
km.keysMutex.RUnlock()
blacklistedCount := 0
km.blacklistedKeys.Range(func(key, value interface{}) bool {
blacklistedCount++
return true
})
failureCountsSize := 0
km.keyFailureCounts.Range(func(key, value interface{}) bool {
failureCountsSize++
return true
})
return &Stats{
CurrentIndex: atomic.LoadInt64(&km.currentIndex),
TotalKeys: totalKeys,
@@ -296,16 +315,16 @@ func (km *KeyManager) ResetKeys() map[string]interface{} {
beforeCount++
return true
})
km.blacklistedKeys = sync.Map{}
km.keyFailureCounts = sync.Map{}
logrus.Infof("🔄 密钥状态已重置，清除了 %d 个黑名单密钥", beforeCount)
km.keysMutex.RLock()
totalKeys := len(km.keys)
km.keysMutex.RUnlock()
return map[string]interface{}{
"success": true,
"message": fmt.Sprintf("已清除 %d 个黑名单密钥", beforeCount),
@@ -317,12 +336,12 @@ func (km *KeyManager) ResetKeys() map[string]interface{} {
// GetBlacklistDetails 获取黑名单详情
func (km *KeyManager) GetBlacklistDetails() *BlacklistInfo {
var blacklistDetails []BlacklistDetail
km.keysMutex.RLock()
keys := km.keys
keyPreviews := km.keyPreviews
km.keysMutex.RUnlock()
for i, key := range keys {
if _, blacklisted := km.blacklistedKeys.Load(key); blacklisted {
blacklistDetails = append(blacklistDetails, BlacklistDetail{
@@ -333,7 +352,7 @@ func (km *KeyManager) GetBlacklistDetails() *BlacklistInfo {
})
}
}
return &BlacklistInfo{
TotalBlacklisted: len(blacklistDetails),
TotalKeys: len(keys),
@@ -345,7 +364,7 @@ func (km *KeyManager) GetBlacklistDetails() *BlacklistInfo {
// setupMemoryCleanup 设置内存清理机制
func (km *KeyManager) setupMemoryCleanup() {
km.cleanupTicker = time.NewTicker(10 * time.Minute)
go func() {
for {
select {
@@ -367,20 +386,20 @@ func (km *KeyManager) performMemoryCleanup() {
maxSize = 1000
}
km.keysMutex.RUnlock()
currentSize := 0
km.keyFailureCounts.Range(func(key, value interface{}) bool {
currentSize++
return true
})
if currentSize > maxSize {
logrus.Infof("🧹 清理失败计数缓存 (%d -> %d)", currentSize, maxSize)
// 简单策略：清理一半的失败计数
cleared := 0
target := currentSize - maxSize
km.keyFailureCounts.Range(func(key, value interface{}) bool {
if cleared < target {
km.keyFailureCounts.Delete(key)