feat: key检测重试功能

This commit is contained in:
tbphp
2025-07-05 20:54:57 +08:00
parent 8f2a0c1afc
commit adfaf3c1e7
3 changed files with 41 additions and 42 deletions

View File

@@ -86,7 +86,7 @@ func main() {
taskService := services.NewTaskService(storage) taskService := services.NewTaskService(storage)
channelFactory := channel.NewFactory(settingsManager) channelFactory := channel.NewFactory(settingsManager)
keyValidatorService := services.NewKeyValidatorService(database, channelFactory) keyValidatorService := services.NewKeyValidatorService(database, channelFactory, settingsManager)
keyManualValidationService := services.NewKeyManualValidationService(database, keyValidatorService, taskService, settingsManager) keyManualValidationService := services.NewKeyManualValidationService(database, keyValidatorService, taskService, settingsManager)
keyCronService := services.NewKeyCronService(database, keyValidatorService, settingsManager) keyCronService := services.NewKeyCronService(database, keyValidatorService, settingsManager)

View File

@@ -56,7 +56,7 @@ func (s *KeyCronService) run() {
// Dynamically get the interval for the next run // Dynamically get the interval for the next run
intervalMinutes := s.SettingsManager.GetInt("key_validation_interval_minutes", 60) intervalMinutes := s.SettingsManager.GetInt("key_validation_interval_minutes", 60)
if intervalMinutes <= 0 { if intervalMinutes <= 0 {
intervalMinutes = 60 // Fallback to a safe default intervalMinutes = 60
} }
nextRunTimer := time.NewTimer(time.Duration(intervalMinutes) * time.Minute) nextRunTimer := time.NewTimer(time.Duration(intervalMinutes) * time.Minute)
@@ -84,6 +84,7 @@ func (s *KeyCronService) validateAllGroups(ctx context.Context) {
// Get effective settings for the group // Get effective settings for the group
effectiveSettings := s.SettingsManager.GetEffectiveConfig(g.Config) effectiveSettings := s.SettingsManager.GetEffectiveConfig(g.Config)
interval := time.Duration(effectiveSettings.KeyValidationIntervalMinutes) * time.Minute interval := time.Duration(effectiveSettings.KeyValidationIntervalMinutes) * time.Minute
logrus.Infof("KeyCronService: Validating group %s with interval %s", g.Name, interval)
// Check if it's time to validate this group // Check if it's time to validate this group
if g.LastValidatedAt == nil || time.Since(*g.LastValidatedAt) > interval { if g.LastValidatedAt == nil || time.Since(*g.LastValidatedAt) > interval {

View File

@@ -4,6 +4,7 @@ import (
"context" "context"
"fmt" "fmt"
"gpt-load/internal/channel" "gpt-load/internal/channel"
"gpt-load/internal/config"
"gpt-load/internal/models" "gpt-load/internal/models"
"github.com/sirupsen/logrus" "github.com/sirupsen/logrus"
@@ -21,19 +22,20 @@ type KeyTestResult struct {
type KeyValidatorService struct { type KeyValidatorService struct {
DB *gorm.DB DB *gorm.DB
channelFactory *channel.Factory channelFactory *channel.Factory
SettingsManager *config.SystemSettingsManager
} }
// NewKeyValidatorService creates a new KeyValidatorService. // NewKeyValidatorService creates a new KeyValidatorService.
func NewKeyValidatorService(db *gorm.DB, factory *channel.Factory) *KeyValidatorService { func NewKeyValidatorService(db *gorm.DB, factory *channel.Factory, settingsManager *config.SystemSettingsManager) *KeyValidatorService {
return &KeyValidatorService{ return &KeyValidatorService{
DB: db, DB: db,
channelFactory: factory, channelFactory: factory,
SettingsManager: settingsManager,
} }
} }
// ValidateSingleKey performs a validation check on a single API key. // ValidateSingleKey performs a validation check on a single API key.
func (s *KeyValidatorService) ValidateSingleKey(ctx context.Context, key *models.APIKey, group *models.Group) (bool, error) { func (s *KeyValidatorService) ValidateSingleKey(ctx context.Context, key *models.APIKey, group *models.Group) (bool, error) {
// 添加超时保护
if ctx.Err() != nil { if ctx.Err() != nil {
return false, fmt.Errorf("context cancelled or timed out: %w", ctx.Err()) return false, fmt.Errorf("context cancelled or timed out: %w", ctx.Err())
} }
@@ -49,56 +51,52 @@ func (s *KeyValidatorService) ValidateSingleKey(ctx context.Context, key *models
return false, fmt.Errorf("failed to get channel for group %s: %w", group.Name, err) return false, fmt.Errorf("failed to get channel for group %s: %w", group.Name, err)
} }
// 记录验证开始 effectiveSettings := s.SettingsManager.GetEffectiveConfig(group.Config)
logrus.WithFields(logrus.Fields{ retries := effectiveSettings.BlacklistThreshold
"key_id": key.ID, if retries <= 0 {
"group_id": group.ID, retries = 1
"group_name": group.Name, }
}).Debug("Starting key validation")
var lastErr error
isValid, validationErr := ch.ValidateKey(ctx, key.KeyValue) for range retries {
if validationErr != nil { isValid, validationErr := ch.ValidateKey(ctx, key.KeyValue)
logrus.WithFields(logrus.Fields{ if validationErr == nil && isValid {
"key_id": key.ID, logrus.WithFields(logrus.Fields{
"group_id": group.ID, "key_id": key.ID,
"group_name": group.Name, "is_valid": isValid,
"error": validationErr, }).Debug("Key validation successful")
}).Debug("Key validation failed") return true, nil
return false, validationErr }
lastErr = validationErr
} }
// 记录验证结果
logrus.WithFields(logrus.Fields{ logrus.WithFields(logrus.Fields{
"key_id": key.ID, "key_id": key.ID,
"group_id": group.ID, "group_id": group.ID,
"group_name": group.Name, "max_retries": retries,
"is_valid": isValid, }).Debug("Key validation failed after all retries")
}).Debug("Key validation completed")
return isValid, nil return false, lastErr
} }
// TestMultipleKeys performs a synchronous validation for a list of key values within a specific group. // TestMultipleKeys performs a synchronous validation for a list of key values within a specific group.
func (s *KeyValidatorService) TestMultipleKeys(ctx context.Context, group *models.Group, keyValues []string) ([]KeyTestResult, error) { func (s *KeyValidatorService) TestMultipleKeys(ctx context.Context, group *models.Group, keyValues []string) ([]KeyTestResult, error) {
results := make([]KeyTestResult, len(keyValues)) results := make([]KeyTestResult, len(keyValues))
ch, err := s.channelFactory.GetChannel(group)
if err != nil {
return nil, fmt.Errorf("failed to get channel for group %s: %w", group.Name, err)
}
// Find which of the provided keys actually exist in the database for this group // Find which of the provided keys actually exist in the database for this group
var existingKeys []models.APIKey var existingKeys []models.APIKey
if err := s.DB.Where("group_id = ? AND key_value IN ?", group.ID, keyValues).Find(&existingKeys).Error; err != nil { if err := s.DB.Where("group_id = ? AND key_value IN ?", group.ID, keyValues).Find(&existingKeys).Error; err != nil {
return nil, fmt.Errorf("failed to query keys from DB: %w", err) return nil, fmt.Errorf("failed to query keys from DB: %w", err)
} }
existingKeyMap := make(map[string]bool) existingKeyMap := make(map[string]models.APIKey)
for _, k := range existingKeys { for _, k := range existingKeys {
existingKeyMap[k.KeyValue] = true existingKeyMap[k.KeyValue] = k
} }
for i, kv := range keyValues { for i, kv := range keyValues {
// Pre-check: ensure the key belongs to the group to prevent unnecessary API calls apiKey, exists := existingKeyMap[kv]
if !existingKeyMap[kv] { if !exists {
results[i] = KeyTestResult{ results[i] = KeyTestResult{
KeyValue: kv, KeyValue: kv,
IsValid: false, IsValid: false,
@@ -107,11 +105,11 @@ func (s *KeyValidatorService) TestMultipleKeys(ctx context.Context, group *model
continue continue
} }
isValid, validationErr := ch.ValidateKey(ctx, kv) isValid, validationErr := s.ValidateSingleKey(ctx, &apiKey, group)
results[i] = KeyTestResult{ results[i] = KeyTestResult{
KeyValue: kv, KeyValue: kv,
IsValid: isValid, IsValid: isValid,
Error: "", // Explicitly set error to empty string on success Error: "",
} }
if validationErr != nil { if validationErr != nil {
results[i].Error = validationErr.Error() results[i].Error = validationErr.Error()