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-07-07 20:19:35 +08:00
parent 747f130f7e
commit ea54897e31
5 changed files with 25 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
internal/keypool/provider.go
View File

@@ -103,7 +103,7 @@ func (p *KeyProvider) handleSuccess(keyID uint, keyHashKey, activeKeysListKey st
return
}
if err := p.store.HSet(keyHashKey, "failure_count", 0); err != nil {
if err := p.store.HSet(keyHashKey, map[string]any{"failure_count": 0}); err != nil {
logrus.WithFields(logrus.Fields{"keyID": keyID, "error": err}).Error("Failed to reset failure count in store, aborting DB update.")
return
}
@@ -112,7 +112,7 @@ func (p *KeyProvider) handleSuccess(keyID uint, keyHashKey, activeKeysListKey st
if isInvalid {
logrus.WithField("keyID", keyID).Info("Key has recovered and is being restored to active pool.")
if err := p.store.HSet(keyHashKey, "status", models.KeyStatusActive); err != nil {
if err := p.store.HSet(keyHashKey, map[string]any{"status": models.KeyStatusActive}); err != nil {
logrus.WithFields(logrus.Fields{"keyID": keyID, "error": err}).Error("Failed to update key status to active in store, aborting DB update.")
return
}
@@ -160,7 +160,7 @@ func (p *KeyProvider) handleFailure(keyID uint, keyHashKey, activeKeysListKey st
logrus.WithFields(logrus.Fields{"keyID": keyID, "error": err}).Error("Failed to LRem key from active list, aborting DB update.")
return
}
if err := p.store.HSet(keyHashKey, "status", models.KeyStatusInvalid); err != nil {
if err := p.store.HSet(keyHashKey, map[string]any{"status": models.KeyStatusInvalid}); err != nil {
logrus.WithFields(logrus.Fields{"keyID": keyID, "error": err}).Error("Failed to update key status to invalid in store, aborting DB update.")
return
}
@@ -222,10 +222,8 @@ func (p *KeyProvider) LoadKeysFromDB() error {
if pipeline != nil {
pipeline.HSet(keyHashKey, keyDetails)
} else {
for field, value := range keyDetails {
if err := p.store.HSet(keyHashKey, field, value); err != nil {
logrus.WithFields(logrus.Fields{"keyID": key.ID, "error": err}).Error("Failed to HSet key details")
}
if err := p.store.HSet(keyHashKey, keyDetails); err != nil {
logrus.WithFields(logrus.Fields{"keyID": key.ID, "error": err}).Error("Failed to HSet key details")
}
}
@@ -307,7 +305,9 @@ func (p *KeyProvider) RemoveKeys(groupID uint, keyValues []string) (int64, error
return nil
}
result := tx.Where("group_id = ? AND key_value IN ?", groupID, keyValues).Delete(&models.APIKey{})
keyIDsToDelete := pluckIDs(keysToDelete)
result := tx.Where("id IN ?", keyIDsToDelete).Delete(&models.APIKey{})
if result.Error != nil {
return result.Error
}
@@ -394,17 +394,13 @@ func (p *KeyProvider) addKeyToStore(key *models.APIKey) error {
// 1. Store key details in HASH
keyHashKey := fmt.Sprintf("key:%d", key.ID)
keyDetails := p.apiKeyToMap(key)
for field, value := range keyDetails {
if err := p.store.HSet(keyHashKey, field, value); err != nil {
return fmt.Errorf("failed to HSet key details for key %d: %w", key.ID, err)
}
if err := p.store.HSet(keyHashKey, keyDetails); err != nil {
return fmt.Errorf("failed to HSet key details for key %d: %w", key.ID, err)
}
// 2. If active, add to the active LIST
if key.Status == models.KeyStatusActive {
activeKeysListKey := fmt.Sprintf("group:%d:active_keys", key.GroupID)
// To prevent duplicates, first remove any existing instance of the key from the list.
// This makes the add operation idempotent regarding the list.
if err := p.store.LRem(activeKeysListKey, 0, key.ID); err != nil {
return fmt.Errorf("failed to LRem key %d before LPush for group %d: %w", key.ID, key.GroupID, err)
}
@@ -432,7 +428,7 @@ func (p *KeyProvider) removeKeyFromStore(keyID, groupID uint) error {
// apiKeyToMap converts an APIKey model to a map for HSET.
func (p *KeyProvider) apiKeyToMap(key *models.APIKey) map[string]any {
return map[string]any{
"id": fmt.Sprint(key.ID), // Use fmt.Sprint for consistency in pipeline
"id": fmt.Sprint(key.ID),
"key_string": key.KeyValue,
"status": key.Status,
"failure_count": key.FailureCount,

24
internal/services/key_service.go
View File

@@ -20,9 +20,9 @@ type AddKeysResult struct {
// DeleteKeysResult holds the result of deleting multiple keys.
type DeleteKeysResult struct {
DeletedCount int `json:"deleted_count"`
IgnoredCount int `json:"ignored_count"`
TotalInGroup int64 `json:"total_in_group"`
DeletedCount int `json:"deleted_count"`
IgnoredCount int `json:"ignored_count"`
TotalInGroup int64 `json:"total_in_group"`
}
// KeyService provides services related to API keys.
@@ -80,26 +80,21 @@ func (s *KeyService) AddMultipleKeys(groupID uint, keysText string) (*AddKeysRes
}
if len(newKeysToCreate) == 0 {
var totalInGroup int64
s.DB.Model(&models.APIKey{}).Where("group_id = ?", groupID).Count(&totalInGroup)
return &AddKeysResult{
AddedCount: 0,
IgnoredCount: len(keys),
TotalInGroup: totalInGroup,
TotalInGroup: int64(len(existingKeys)),
}, nil
}
// 4. Use KeyProvider to add keys, which handles DB and cache
// 4. Use KeyProvider to add keys
err := s.KeyProvider.AddKeys(groupID, newKeysToCreate)
if err != nil {
return nil, err
}
// 5. Get the new total count
var totalInGroup int64
if err := s.DB.Model(&models.APIKey{}).Where("group_id = ?", groupID).Count(&totalInGroup).Error; err != nil {
return nil, err
}
// 5. Calculate new total count
totalInGroup := int64(len(existingKeys) + len(newKeysToCreate))
return &AddKeysResult{
AddedCount: len(newKeysToCreate),
@@ -171,13 +166,11 @@ func (s *KeyService) ClearAllInvalidKeys(groupID uint) (int64, error) {
// DeleteMultipleKeys handles the business logic of deleting keys from a text block.
func (s *KeyService) DeleteMultipleKeys(groupID uint, keysText string) (*DeleteKeysResult, error) {
// 1. Parse keys from the text block
keysToDelete := s.ParseKeysFromText(keysText)
if len(keysToDelete) == 0 {
return nil, fmt.Errorf("no valid keys found in the input text")
}
// 2. Use KeyProvider to delete keys, which handles DB and cache
deletedCount, err := s.KeyProvider.RemoveKeys(groupID, keysToDelete)
if err != nil {
return nil, err
@@ -185,7 +178,6 @@ func (s *KeyService) DeleteMultipleKeys(groupID uint, keysText string) (*DeleteK
ignoredCount := len(keysToDelete) - int(deletedCount)
// 3. Get the new total count
var totalInGroup int64
if err := s.DB.Model(&models.APIKey{}).Where("group_id = ?", groupID).Count(&totalInGroup).Error; err != nil {
return nil, err
@@ -199,7 +191,6 @@ func (s *KeyService) DeleteMultipleKeys(groupID uint, keysText string) (*DeleteK
}
// ListKeysInGroupQuery builds a query to list all keys within a specific group, filtered by status.
// It returns a GORM query builder, allowing the handler to apply pagination.
func (s *KeyService) ListKeysInGroupQuery(groupID uint, statusFilter string, searchKeyword string) *gorm.DB {
query := s.DB.Model(&models.APIKey{}).Where("group_id = ?", groupID)
@@ -208,7 +199,6 @@ func (s *KeyService) ListKeysInGroupQuery(groupID uint, statusFilter string, sea
}
if searchKeyword != "" {
// Use LIKE for fuzzy search on the key_value
query = query.Where("key_value LIKE ?", "%"+searchKeyword+"%")
}

7
internal/store/memory.go
View File

@@ -37,7 +37,6 @@ func (s *MemoryStore) Close() error {
return nil
}
// Set stores a key-value pair.
func (s *MemoryStore) Set(key string, value []byte, ttl time.Duration) error {
s.mu.Lock()
@@ -147,7 +146,7 @@ func (s *MemoryStore) SetNX(key string, value []byte, ttl time.Duration) (bool,
// --- HASH operations ---
func (s *MemoryStore) HSet(key, field string, value any) error {
func (s *MemoryStore) HSet(key string, values map[string]any) error {
s.mu.Lock()
defer s.mu.Unlock()
@@ -164,7 +163,9 @@ func (s *MemoryStore) HSet(key, field string, value any) error {
}
}
hash[field] = fmt.Sprint(value)
for field, value := range values {
hash[field] = fmt.Sprint(value)
}
return nil
}

4
internal/store/redis.go
View File

@@ -61,8 +61,8 @@ func (s *RedisStore) Close() error {
// --- HASH operations ---
func (s *RedisStore) HSet(key, field string, value any) error {
return s.client.HSet(context.Background(), key, field, value).Err()
func (s *RedisStore) HSet(key string, values map[string]any) error {
return s.client.HSet(context.Background(), key, values).Err()
}
func (s *RedisStore) HGetAll(key string) (map[string]string, error) {

2
internal/store/store.go
View File

@@ -33,7 +33,7 @@ type Store interface {
SetNX(key string, value []byte, ttl time.Duration) (bool, error)
// HASH operations
HSet(key, field string, value any) error
HSet(key string, values map[string]any) error
HGetAll(key string) (map[string]string, error)
HIncrBy(key, field string, incr int64) (int64, error)