Skip to content

Commit 1a800dd

Browse files
tbphptbphp
committed
feat: 优化key验证流程
1 parent 411f8f2 commit 1a800dd

File tree

6 files changed

+103
-44
lines changed

6 files changed

+103
-44
lines changed

‎internal/handler/key_handler.go‎

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -201,11 +201,17 @@ func (s *Server) TestMultipleKeys(c *gin.Context) {
201201
return
202202
}
203203

204-
group, ok := s.findGroupByID(c, req.GroupID)
204+
groupDB, ok := s.findGroupByID(c, req.GroupID)
205205
if !ok {
206206
return
207207
}
208208

209+
group, err := s.GroupManager.GetGroupByName(groupDB.Name)
210+
if err != nil {
211+
response.Error(c, app_errors.NewAPIError(app_errors.ErrResourceNotFound, fmt.Sprintf("Group '%s' not found", groupDB.Name)))
212+
return
213+
}
214+
209215
if err := validateKeysText(req.KeysText); err != nil {
210216
response.Error(c, app_errors.NewAPIError(app_errors.ErrValidation, err.Error()))
211217
return
@@ -234,11 +240,17 @@ func (s *Server) ValidateGroupKeys(c *gin.Context) {
234240
return
235241
}
236242

237-
group, ok := s.findGroupByID(c, req.GroupID)
243+
groupDB, ok := s.findGroupByID(c, req.GroupID)
238244
if !ok {
239245
return
240246
}
241247

248+
group, err := s.GroupManager.GetGroupByName(groupDB.Name)
249+
if err != nil {
250+
response.Error(c, app_errors.NewAPIError(app_errors.ErrResourceNotFound, fmt.Sprintf("Group '%s' not found", groupDB.Name)))
251+
return
252+
}
253+
242254
taskStatus, err := s.KeyManualValidationService.StartValidationTask(group)
243255
if err != nil {
244256
response.Error(c, app_errors.NewAPIError(app_errors.ErrTaskInProgress, err.Error()))

‎internal/keypool/cron_checker.go‎

Lines changed: 75 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import (
66
"gpt-load/internal/models"
77
"gpt-load/internal/store"
88
"sync"
9+
"sync/atomic"
910
"time"
1011

1112
"github.com/sirupsen/logrus"
@@ -89,7 +90,7 @@ func (s *CronChecker) runLoop() {
8990
}
9091
}
9192

92-
// submitValidationJobs finds groups whose keys need validation and validates them.
93+
// submitValidationJobs finds groups whose keys need validation and validates them concurrently.
9394
func (s *CronChecker) submitValidationJobs() {
9495
var groups []models.Group
9596
if err := s.DB.Find(&groups).Error; err != nil {
@@ -98,51 +99,93 @@ func (s *CronChecker) submitValidationJobs() {
9899
}
99100

100101
validationStartTime := time.Now()
102+
var wg sync.WaitGroup
101103

102104
for i := range groups {
103105
group := &groups[i]
104-
effectiveSettings := s.SettingsManager.GetEffectiveConfig(group.Config)
105-
interval := time.Duration(effectiveSettings.KeyValidationIntervalMinutes) * time.Minute
106+
group.EffectiveConfig = s.SettingsManager.GetEffectiveConfig(group.Config)
107+
interval := time.Duration(group.EffectiveConfig.KeyValidationIntervalMinutes) * time.Minute
106108

107109
if group.LastValidatedAt == nil || validationStartTime.Sub(*group.LastValidatedAt) > interval {
108-
groupProcessStart := time.Now()
109-
var invalidKeys []models.APIKey
110-
err := s.DB.Where("group_id = ? AND status = ?", group.ID, models.KeyStatusInvalid).Find(&invalidKeys).Error
111-
if err != nil {
112-
logrus.Errorf("CronChecker: Failed to get invalid keys for group %s: %v", group.Name, err)
113-
continue
114-
}
110+
wg.Add(1)
111+
g := group
112+
go func() {
113+
defer wg.Done()
114+
s.validateGroupKeys(g)
115+
}()
116+
}
117+
}
118+
119+
wg.Wait()
120+
}
121+
122+
// validateGroupKeys validates all invalid keys for a single group concurrently.
123+
func (s *CronChecker) validateGroupKeys(group *models.Group) {
124+
groupProcessStart := time.Now()
115125

116-
validatedCount := len(invalidKeys)
117-
becameValidCount := 0
118-
if validatedCount > 0 {
119-
for j := range invalidKeys {
120-
select {
121-
case <-s.stopChan:
126+
var invalidKeys []models.APIKey
127+
err := s.DB.Where("group_id = ? AND status = ?", group.ID, models.KeyStatusInvalid).Find(&invalidKeys).Error
128+
if err != nil {
129+
logrus.Errorf("CronChecker: Failed to get invalid keys for group %s: %v", group.Name, err)
130+
return
131+
}
132+
133+
if len(invalidKeys) == 0 {
134+
if err := s.DB.Model(group).Update("last_validated_at", time.Now()).Error; err != nil {
135+
logrus.Errorf("CronChecker: Failed to update last_validated_at for group %s: %v", group.Name, err)
136+
}
137+
logrus.Infof("CronChecker: Group '%s' has no invalid keys to check.", group.Name)
138+
return
139+
}
140+
141+
var becameValidCount int32
142+
var keyWg sync.WaitGroup
143+
jobs := make(chan *models.APIKey, len(invalidKeys))
144+
145+
concurrency := group.EffectiveConfig.KeyValidationConcurrency
146+
for range concurrency {
147+
keyWg.Add(1)
148+
go func() {
149+
defer keyWg.Done()
150+
for {
151+
select {
152+
case key, ok := <-jobs:
153+
if !ok {
122154
return
123-
default:
124155
}
125-
key := &invalidKeys[j]
126156
isValid, _ := s.Validator.ValidateSingleKey(key, group)
127-
128157
if isValid {
129-
becameValidCount++
158+
atomic.AddInt32(&becameValidCount, 1)
130159
}
160+
case <-s.stopChan:
161+
return
131162
}
132163
}
164+
}()
165+
}
133166

134-
if err := s.DB.Model(group).Update("last_validated_at", time.Now()).Error; err != nil {
135-
logrus.Errorf("CronChecker: Failed to update last_validated_at for group %s: %v", group.Name, err)
136-
}
137-
138-
duration := time.Since(groupProcessStart)
139-
logrus.Infof(
140-
"CronChecker: Group '%s' validation finished. Total checked: %d, became valid: %d. Duration: %s.",
141-
group.Name,
142-
validatedCount,
143-
becameValidCount,
144-
duration.String(),
145-
)
167+
DistributeLoop:
168+
for i := range invalidKeys {
169+
select {
170+
case jobs <- &invalidKeys[i]:
171+
case <-s.stopChan:
172+
break DistributeLoop
146173
}
147174
}
175+
close(jobs)
176+
177+
keyWg.Wait()
178+
179+
if err := s.DB.Model(group).Update("last_validated_at", time.Now()).Error; err != nil {
180+
logrus.Errorf("CronChecker: Failed to update last_validated_at for group %s: %v", group.Name, err)
181+
}
182+
183+
duration := time.Since(groupProcessStart)
184+
logrus.Infof(
185+
"CronChecker: Group '%s' validation finished. Total checked: %d, became valid: %d. Duration: %s.",
186+
group.Name,
187+
len(invalidKeys),
188+
becameValidCount,
189+
duration.String(),
190+
)
148191
}

‎internal/keypool/validator.go‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,10 @@ func NewKeyValidator(params KeyValidatorParams) *KeyValidator {
4848

4949
// ValidateSingleKey performs a validation check on a single API key.
5050
func (s *KeyValidator) ValidateSingleKey(key *models.APIKey, group *models.Group) (bool, error) {
51-
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
51+
if group.EffectiveConfig.AppUrl == "" {
52+
group.EffectiveConfig = s.SettingsManager.GetEffectiveConfig(group.Config)
53+
}
54+
ctx, cancel := context.WithTimeout(context.Background(), time.Duration(group.EffectiveConfig.KeyValidationTimeoutSeconds)*time.Second)
5255
defer cancel()
5356

5457
ch, err := s.channelFactory.GetChannel(group)
@@ -58,7 +61,6 @@ func (s *KeyValidator) ValidateSingleKey(key *models.APIKey, group *models.Group
5861

5962
isValid, validationErr := ch.ValidateKey(ctx, key.KeyValue)
6063

61-
group.EffectiveConfig = s.SettingsManager.GetEffectiveConfig(group.Config)
6264
s.keypoolProvider.UpdateStatus(key, group, isValid)
6365

6466
if !isValid {

‎internal/models/types.go‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,8 @@ type GroupConfig struct {
3434
MaxRetries *int `json:"max_retries,omitempty"`
3535
BlacklistThreshold *int `json:"blacklist_threshold,omitempty"`
3636
KeyValidationIntervalMinutes *int `json:"key_validation_interval_minutes,omitempty"`
37+
KeyValidationConcurrency *int `json:"key_validation_concurrency,omitempty"`
38+
KeyValidationTimeoutSeconds *int `json:"key_validation_timeout_seconds,omitempty"`
3739
}
3840

3941
// Group 对应 groups 表
@@ -96,10 +98,10 @@ type StatCard struct {
9698

9799
// DashboardStatsResponse 用于仪表盘基础统计的API响应
98100
type DashboardStatsResponse struct {
99-
KeyCount StatCard `json:"key_count"`
100-
GroupCount StatCard `json:"group_count"`
101-
RequestCount StatCard `json:"request_count"`
102-
ErrorRate StatCard `json:"error_rate"`
101+
KeyCount StatCard `json:"key_count"`
102+
GroupCount StatCard `json:"group_count"`
103+
RequestCount StatCard `json:"request_count"`
104+
ErrorRate StatCard `json:"error_rate"`
103105
}
104106

105107
// ChartDataset 用于图表的数据集

‎internal/services/key_manual_validation_service.go‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -70,8 +70,7 @@ func (s *KeyManualValidationService) runValidation(group *models.Group, keys []m
7070
jobs := make(chan models.APIKey, len(keys))
7171
results := make(chan bool, len(keys))
7272

73-
// 固定10并发，避免超频
74-
concurrency := 10
73+
concurrency := group.EffectiveConfig.KeyValidationConcurrency
7574

7675
var wg sync.WaitGroup
7776
for range concurrency {
@@ -84,7 +83,6 @@ func (s *KeyManualValidationService) runValidation(group *models.Group, keys []m
8483
}
8584
close(jobs)
8685

87-
// Wait for all workers to complete in a separate goroutine
8886
go func() {
8987
wg.Wait()
9088
close(results)

‎internal/types/types.go‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,9 @@ type SystemSettings struct {
3232
// 密钥配置
3333
MaxRetries int `json:"max_retries" default:"3" name:"最大重试次数" category:"密钥配置" desc:"单个请求使用不同 Key 的最大重试次数，0为不重试。" validate:"min=0"`
3434
BlacklistThreshold int `json:"blacklist_threshold" default:"3" name:"黑名单阈值" category:"密钥配置" desc:"一个 Key 连续失败多少次后进入黑名单，0为不拉黑。" validate:"min=0"`
35-
KeyValidationIntervalMinutes int `json:"key_validation_interval_minutes" default:"60" name:"定时验证周期（分钟）" category:"密钥配置" desc:"后台定时验证密钥的默认周期（分钟）。" validate:"min=30"`
35+
KeyValidationIntervalMinutes int `json:"key_validation_interval_minutes" default:"60" name:"验证间隔（分钟）" category:"密钥配置" desc:"后台验证密钥的默认间隔（分钟）。" validate:"min=30"`
36+
KeyValidationConcurrency int `json:"key_validation_concurrency" default:"10" name:"验证并发数" category:"密钥配置" desc:"后台定时验证无效 Key 时的并发数。" validate:"min=1"`
37+
KeyValidationTimeoutSeconds int `json:"key_validation_timeout_seconds" default:"20" name:"验证超时（秒）" category:"密钥配置" desc:"后台定时验证单个 Key 时的 API 请求超时时间（秒）。" validate:"min=5"`
3638
}
3739

3840
// ServerConfig represents server configuration

0 commit comments

Comments
 (0)