feat: 批量分配处理数据

internal/handler/handler.go

@@ -6,7 +6,6 @@ import (
 	"time"
 
 	"gpt-load/internal/config"
-	"gpt-load/internal/keypool"
 	"gpt-load/internal/models"
 	"gpt-load/internal/services"
 	"gpt-load/internal/types"
@@ -21,7 +20,6 @@ type Server struct {
 	DB                         *gorm.DB
 	config                     types.ConfigManager
 	SettingsManager            *config.SystemSettingsManager
-	KeyValidator               *keypool.KeyValidator
 	KeyManualValidationService *services.KeyManualValidationService
 	TaskService                *services.TaskService
 	KeyService                 *services.KeyService
@@ -34,7 +32,6 @@ type NewServerParams struct {
 	DB                         *gorm.DB
 	Config                     types.ConfigManager
 	SettingsManager            *config.SystemSettingsManager
-	KeyValidator               *keypool.KeyValidator
 	KeyManualValidationService *services.KeyManualValidationService
 	TaskService                *services.TaskService
 	KeyService                 *services.KeyService
@@ -47,7 +44,6 @@ func NewServer(params NewServerParams) *Server {
 		DB:                         params.DB,
 		config:                     params.Config,
 		SettingsManager:            params.SettingsManager,
-		KeyValidator:               params.KeyValidator,
 		KeyManualValidationService: params.KeyManualValidationService,
 		TaskService:                params.TaskService,
 		KeyService:                 params.KeyService,

internal/handler/key_handler.go

@@ -84,7 +84,9 @@ func (s *Server) AddMultipleKeys(c *gin.Context) {
 
 	result, err := s.KeyService.AddMultipleKeys(req.GroupID, req.KeysText)
 	if err != nil {
-		if err.Error() == "no valid keys found in the input text" {
+		if strings.Contains(err.Error(), "batch size exceeds the limit") {
+			response.Error(c, app_errors.NewAPIError(app_errors.ErrValidation, err.Error()))
+		} else if err.Error() == "no valid keys found in the input text" {
 			response.Error(c, app_errors.NewAPIError(app_errors.ErrValidation, err.Error()))
 		} else {
 			response.Error(c, app_errors.ParseDBError(err))
@@ -146,7 +148,9 @@ func (s *Server) DeleteMultipleKeys(c *gin.Context) {
 
 	result, err := s.KeyService.DeleteMultipleKeys(req.GroupID, req.KeysText)
 	if err != nil {
-		if err.Error() == "no valid keys found in the input text" {
+		if strings.Contains(err.Error(), "batch size exceeds the limit") {
+			response.Error(c, app_errors.NewAPIError(app_errors.ErrValidation, err.Error()))
+		} else if err.Error() == "no valid keys found in the input text" {
 			response.Error(c, app_errors.NewAPIError(app_errors.ErrValidation, err.Error()))
 		} else {
 			response.Error(c, app_errors.ParseDBError(err))
@@ -176,7 +180,9 @@ func (s *Server) RestoreMultipleKeys(c *gin.Context) {
 
 	result, err := s.KeyService.RestoreMultipleKeys(req.GroupID, req.KeysText)
 	if err != nil {
-		if err.Error() == "no valid keys found in the input text" {
+		if strings.Contains(err.Error(), "batch size exceeds the limit") {
+			response.Error(c, app_errors.NewAPIError(app_errors.ErrValidation, err.Error()))
+		} else if err.Error() == "no valid keys found in the input text" {
 			response.Error(c, app_errors.NewAPIError(app_errors.ErrValidation, err.Error()))
 		} else {
 			response.Error(c, app_errors.ParseDBError(err))
@@ -205,16 +211,15 @@ func (s *Server) TestMultipleKeys(c *gin.Context) {
 		return
 	}
 
-	// Re-use the parsing logic from the key service
+	results, err := s.KeyService.TestMultipleKeys(c.Request.Context(), group, req.KeysText)
-	keysToTest := s.KeyService.ParseKeysFromText(req.KeysText)
-	if len(keysToTest) == 0 {
-		response.Error(c, app_errors.NewAPIError(app_errors.ErrValidation, "no valid keys found in the input text"))
-		return
-	}
-
-	results, err := s.KeyValidator.TestMultipleKeys(c.Request.Context(), group, keysToTest)
 	if err != nil {
-		response.Error(c, app_errors.ParseDBError(err))
+		if strings.Contains(err.Error(), "batch size exceeds the limit") {
+			response.Error(c, app_errors.NewAPIError(app_errors.ErrValidation, err.Error()))
+		} else if err.Error() == "no valid keys found in the input text" {
+			response.Error(c, app_errors.NewAPIError(app_errors.ErrValidation, err.Error()))
+		} else {
+			response.Error(c, app_errors.ParseDBError(err))
+		}
 		return
 	}
 

internal/services/key_service.go

@@ -1,6 +1,7 @@
 package services
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"gpt-load/internal/keypool"
@@ -11,6 +12,11 @@ import (
 	"gorm.io/gorm"
 )
 
+const (
+	maxRequestKeys = 5000
+	chunkSize      = 1000
+)
+
 // AddKeysResult holds the result of adding multiple keys.
 type AddKeysResult struct {
 	AddedCount   int   `json:"added_count"`
@@ -34,15 +40,17 @@ type RestoreKeysResult struct {
 
 // KeyService provides services related to API keys.
 type KeyService struct {
-	DB          *gorm.DB
+	DB           *gorm.DB
-	KeyProvider *keypool.KeyProvider
+	KeyProvider  *keypool.KeyProvider
+	KeyValidator *keypool.KeyValidator
 }
 
 // NewKeyService creates a new KeyService.
-func NewKeyService(db *gorm.DB, keyProvider *keypool.KeyProvider) *KeyService {
+func NewKeyService(db *gorm.DB, keyProvider *keypool.KeyProvider, keyValidator *keypool.KeyValidator) *KeyService {
 	return &KeyService{
-		DB:          db,
+		DB:           db,
-		KeyProvider: keyProvider,
+		KeyProvider:  keyProvider,
+		KeyValidator: keyValidator,
 	}
 }
 
@@ -50,6 +58,9 @@ func NewKeyService(db *gorm.DB, keyProvider *keypool.KeyProvider) *KeyService {
 func (s *KeyService) AddMultipleKeys(groupID uint, keysText string) (*AddKeysResult, error) {
 	// 1. Parse keys from the text block
 	keys := s.ParseKeysFromText(keysText)
+	if len(keys) > maxRequestKeys {
+		return nil, fmt.Errorf("batch size exceeds the limit of %d keys, got %d", maxRequestKeys, len(keys))
+	}
 	if len(keys) == 0 {
 		return nil, fmt.Errorf("no valid keys found in the input text")
 	}
@@ -94,10 +105,16 @@ func (s *KeyService) AddMultipleKeys(groupID uint, keysText string) (*AddKeysRes
 		}, nil
 	}
 
-	// 4. Use KeyProvider to add keys
+	// 4. Use KeyProvider to add keys in chunks
-	err := s.KeyProvider.AddKeys(groupID, newKeysToCreate)
+	for i := 0; i < len(newKeysToCreate); i += chunkSize {
-	if err != nil {
+		end := i + chunkSize
-		return nil, err
+		if end > len(newKeysToCreate) {
+			end = len(newKeysToCreate)
+		}
+		chunk := newKeysToCreate[i:end]
+		if err := s.KeyProvider.AddKeys(groupID, chunk); err != nil {
+			return nil, err
+		}
 	}
 
 	// 5. Calculate new total count
@@ -164,16 +181,28 @@ func (s *KeyService) isValidKeyFormat(key string) bool {
 // RestoreMultipleKeys handles the business logic of restoring keys from a text block.
 func (s *KeyService) RestoreMultipleKeys(groupID uint, keysText string) (*RestoreKeysResult, error) {
 	keysToRestore := s.ParseKeysFromText(keysText)
+	if len(keysToRestore) > maxRequestKeys {
+		return nil, fmt.Errorf("batch size exceeds the limit of %d keys, got %d", maxRequestKeys, len(keysToRestore))
+	}
 	if len(keysToRestore) == 0 {
 		return nil, fmt.Errorf("no valid keys found in the input text")
 	}
 
-	restoredCount, err := s.KeyProvider.RestoreMultipleKeys(groupID, keysToRestore)
+	var totalRestoredCount int64
-	if err != nil {
+	for i := 0; i < len(keysToRestore); i += chunkSize {
-		return nil, err
+		end := i + chunkSize
+		if end > len(keysToRestore) {
+			end = len(keysToRestore)
+		}
+		chunk := keysToRestore[i:end]
+		restoredCount, err := s.KeyProvider.RestoreMultipleKeys(groupID, chunk)
+		if err != nil {
+			return nil, err
+		}
+		totalRestoredCount += restoredCount
 	}
 
-	ignoredCount := len(keysToRestore) - int(restoredCount)
+	ignoredCount := len(keysToRestore) - int(totalRestoredCount)
 
 	var totalInGroup int64
 	if err := s.DB.Model(&models.APIKey{}).Where("group_id = ?", groupID).Count(&totalInGroup).Error; err != nil {
@@ -181,7 +210,7 @@ func (s *KeyService) RestoreMultipleKeys(groupID uint, keysText string) (*Restor
 	}
 
 	return &RestoreKeysResult{
-		RestoredCount: int(restoredCount),
+		RestoredCount: int(totalRestoredCount),
 		IgnoredCount:  ignoredCount,
 		TotalInGroup:  totalInGroup,
 	}, nil
@@ -200,16 +229,28 @@ 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) {
 	keysToDelete := s.ParseKeysFromText(keysText)
+	if len(keysToDelete) > maxRequestKeys {
+		return nil, fmt.Errorf("batch size exceeds the limit of %d keys, got %d", maxRequestKeys, len(keysToDelete))
+	}
 	if len(keysToDelete) == 0 {
 		return nil, fmt.Errorf("no valid keys found in the input text")
 	}
 
-	deletedCount, err := s.KeyProvider.RemoveKeys(groupID, keysToDelete)
+	var totalDeletedCount int64
-	if err != nil {
+	for i := 0; i < len(keysToDelete); i += chunkSize {
-		return nil, err
+		end := i + chunkSize
+		if end > len(keysToDelete) {
+			end = len(keysToDelete)
+		}
+		chunk := keysToDelete[i:end]
+		deletedCount, err := s.KeyProvider.RemoveKeys(groupID, chunk)
+		if err != nil {
+			return nil, err
+		}
+		totalDeletedCount += deletedCount
 	}
 
-	ignoredCount := len(keysToDelete) - int(deletedCount)
+	ignoredCount := len(keysToDelete) - int(totalDeletedCount)
 
 	var totalInGroup int64
 	if err := s.DB.Model(&models.APIKey{}).Where("group_id = ?", groupID).Count(&totalInGroup).Error; err != nil {
@@ -217,7 +258,7 @@ func (s *KeyService) DeleteMultipleKeys(groupID uint, keysText string) (*DeleteK
 	}
 
 	return &DeleteKeysResult{
-		DeletedCount: int(deletedCount),
+		DeletedCount: int(totalDeletedCount),
 		IgnoredCount: ignoredCount,
 		TotalInGroup: totalInGroup,
 	}, nil
@@ -237,3 +278,32 @@ func (s *KeyService) ListKeysInGroupQuery(groupID uint, statusFilter string, sea
 
 	return query
 }
+
+// TestMultipleKeys handles a one-off validation test for multiple keys.
+func (s *KeyService) TestMultipleKeys(ctx context.Context, group *models.Group, keysText string) ([]keypool.KeyTestResult, error) {
+	keysToTest := s.ParseKeysFromText(keysText)
+	if len(keysToTest) > maxRequestKeys {
+		return nil, fmt.Errorf("batch size exceeds the limit of %d keys, got %d", maxRequestKeys, len(keysToTest))
+	}
+	if len(keysToTest) == 0 {
+		return nil, fmt.Errorf("no valid keys found in the input text")
+	}
+
+	var allResults []keypool.KeyTestResult
+	for i := 0; i < len(keysToTest); i += chunkSize {
+		end := i + chunkSize
+		if end > len(keysToTest) {
+			end = len(keysToTest)
+		}
+		chunk := keysToTest[i:end]
+		results, err := s.KeyValidator.TestMultipleKeys(ctx, group, chunk)
+		if err != nil {
+			// If one chunk fails, we might want to stop or collect partial results.
+			// For now, let's stop and return the error.
+			return nil, err
+		}
+		allResults = append(allResults, results...)
+	}
+
+	return allResults, nil
+}