feat: 密钥调整为异步任务,取消数量限制
This commit is contained in:
@@ -51,6 +51,9 @@ func BuildContainer() (*dig.Container, error) {
|
||||
if err := container.Provide(services.NewKeyService); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := container.Provide(services.NewKeyImportService); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := container.Provide(services.NewLogService); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
@@ -23,6 +23,7 @@ type Server struct {
|
||||
KeyManualValidationService *services.KeyManualValidationService
|
||||
TaskService *services.TaskService
|
||||
KeyService *services.KeyService
|
||||
KeyImportService *services.KeyImportService
|
||||
LogService *services.LogService
|
||||
CommonHandler *CommonHandler
|
||||
}
|
||||
@@ -37,6 +38,7 @@ type NewServerParams struct {
|
||||
KeyManualValidationService *services.KeyManualValidationService
|
||||
TaskService *services.TaskService
|
||||
KeyService *services.KeyService
|
||||
KeyImportService *services.KeyImportService
|
||||
LogService *services.LogService
|
||||
CommonHandler *CommonHandler
|
||||
}
|
||||
@@ -51,6 +53,7 @@ func NewServer(params NewServerParams) *Server {
|
||||
KeyManualValidationService: params.KeyManualValidationService,
|
||||
TaskService: params.TaskService,
|
||||
KeyService: params.KeyService,
|
||||
KeyImportService: params.KeyImportService,
|
||||
LogService: params.LogService,
|
||||
CommonHandler: params.CommonHandler,
|
||||
}
|
||||
|
@@ -34,10 +34,6 @@ func validateKeysText(keysText string) error {
|
||||
return fmt.Errorf("keys text cannot be empty")
|
||||
}
|
||||
|
||||
if len(keysText) > 10*1024*1024 {
|
||||
return fmt.Errorf("keys text is too large (max 10MB)")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -98,6 +94,33 @@ func (s *Server) AddMultipleKeys(c *gin.Context) {
|
||||
response.Success(c, result)
|
||||
}
|
||||
|
||||
// AddMultipleKeysAsync handles creating new keys from a text block within a specific group.
|
||||
func (s *Server) AddMultipleKeysAsync(c *gin.Context) {
|
||||
var req KeyTextRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
response.Error(c, app_errors.NewAPIError(app_errors.ErrInvalidJSON, err.Error()))
|
||||
return
|
||||
}
|
||||
|
||||
group, ok := s.findGroupByID(c, req.GroupID)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
if err := validateKeysText(req.KeysText); err != nil {
|
||||
response.Error(c, app_errors.NewAPIError(app_errors.ErrValidation, err.Error()))
|
||||
return
|
||||
}
|
||||
|
||||
taskStatus, err := s.KeyImportService.StartImportTask(group, req.KeysText)
|
||||
if err != nil {
|
||||
response.Error(c, app_errors.NewAPIError(app_errors.ErrTaskInProgress, err.Error()))
|
||||
return
|
||||
}
|
||||
|
||||
response.Success(c, taskStatus)
|
||||
}
|
||||
|
||||
// ListKeysInGroup handles listing all keys within a specific group with pagination.
|
||||
func (s *Server) ListKeysInGroup(c *gin.Context) {
|
||||
groupID, err := validateGroupIDFromQuery(c)
|
||||
|
@@ -123,6 +123,7 @@ func registerProtectedAPIRoutes(api *gin.RouterGroup, serverHandler *handler.Ser
|
||||
keys.GET("", serverHandler.ListKeysInGroup)
|
||||
keys.GET("/export", serverHandler.ExportKeys)
|
||||
keys.POST("/add-multiple", serverHandler.AddMultipleKeys)
|
||||
keys.POST("/add-async", serverHandler.AddMultipleKeysAsync)
|
||||
keys.POST("/delete-multiple", serverHandler.DeleteMultipleKeys)
|
||||
keys.POST("/restore-multiple", serverHandler.RestoreMultipleKeys)
|
||||
keys.POST("/restore-all-invalid", serverHandler.RestoreAllInvalidKeys)
|
||||
|
76
internal/services/key_import_service.go
Normal file
76
internal/services/key_import_service.go
Normal file
@@ -0,0 +1,76 @@
|
||||
package services
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"gpt-load/internal/models"
|
||||
"time"
|
||||
|
||||
"github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
const (
|
||||
importChunkSize = 1000
|
||||
importTimeout = 24 * time.Hour
|
||||
)
|
||||
|
||||
// KeyImportResult holds the result of an import task.
|
||||
type KeyImportResult struct {
|
||||
AddedCount int `json:"added_count"`
|
||||
IgnoredCount int `json:"ignored_count"`
|
||||
}
|
||||
|
||||
// KeyImportService handles the asynchronous import of a large number of keys.
|
||||
type KeyImportService struct {
|
||||
TaskService *TaskService
|
||||
KeyService *KeyService
|
||||
}
|
||||
|
||||
// NewKeyImportService creates a new KeyImportService.
|
||||
func NewKeyImportService(taskService *TaskService, keyService *KeyService) *KeyImportService {
|
||||
return &KeyImportService{
|
||||
TaskService: taskService,
|
||||
KeyService: keyService,
|
||||
}
|
||||
}
|
||||
|
||||
// StartImportTask initiates a new asynchronous key import task.
|
||||
func (s *KeyImportService) StartImportTask(group *models.Group, keysText string) (*TaskStatus, error) {
|
||||
keys := s.KeyService.ParseKeysFromText(keysText)
|
||||
if len(keys) == 0 {
|
||||
return nil, fmt.Errorf("no valid keys found in the input text")
|
||||
}
|
||||
|
||||
initialStatus, err := s.TaskService.StartTask(TaskTypeKeyImport, group.Name, len(keys), importTimeout)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
go s.runImport(group, keys)
|
||||
|
||||
return initialStatus, nil
|
||||
}
|
||||
|
||||
func (s *KeyImportService) runImport(group *models.Group, keys []string) {
|
||||
progressCallback := func(processed int) {
|
||||
if err := s.TaskService.UpdateProgress(processed); err != nil {
|
||||
logrus.Warnf("Failed to update task progress for group %d: %v", group.ID, err)
|
||||
}
|
||||
}
|
||||
|
||||
addedCount, ignoredCount, err := s.KeyService.processAndCreateKeys(group.ID, keys, progressCallback)
|
||||
if err != nil {
|
||||
if endErr := s.TaskService.EndTask(nil, err); endErr != nil {
|
||||
logrus.Errorf("Failed to end task with error for group %d: %v (original error: %v)", group.ID, endErr, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
result := KeyImportResult{
|
||||
AddedCount: addedCount,
|
||||
IgnoredCount: ignoredCount,
|
||||
}
|
||||
|
||||
if endErr := s.TaskService.EndTask(result, nil); endErr != nil {
|
||||
logrus.Errorf("Failed to end task with success result for group %d: %v", group.ID, endErr)
|
||||
}
|
||||
}
|
@@ -53,7 +53,7 @@ func (s *KeyManualValidationService) StartValidationTask(group *models.Group) (*
|
||||
|
||||
timeout := 30 * time.Minute
|
||||
|
||||
taskStatus, err := s.TaskService.StartTask(group.Name, len(keys), timeout)
|
||||
taskStatus, err := s.TaskService.StartTask(TaskTypeKeyValidation, group.Name, len(keys), timeout)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
@@ -55,8 +55,8 @@ func NewKeyService(db *gorm.DB, keyProvider *keypool.KeyProvider, keyValidator *
|
||||
}
|
||||
|
||||
// AddMultipleKeys handles the business logic of creating new keys from a text block.
|
||||
// deprecated: use KeyImportService for large imports
|
||||
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))
|
||||
@@ -65,17 +65,40 @@ func (s *KeyService) AddMultipleKeys(groupID uint, keysText string) (*AddKeysRes
|
||||
return nil, fmt.Errorf("no valid keys found in the input text")
|
||||
}
|
||||
|
||||
// 2. Get existing keys in the group for deduplication
|
||||
addedCount, ignoredCount, err := s.processAndCreateKeys(groupID, keys, nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var totalInGroup int64
|
||||
if err := s.DB.Model(&models.APIKey{}).Where("group_id = ?", groupID).Count(&totalInGroup).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &AddKeysResult{
|
||||
AddedCount: addedCount,
|
||||
IgnoredCount: ignoredCount,
|
||||
TotalInGroup: totalInGroup,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// processAndCreateKeys is the lowest-level reusable function for adding keys.
|
||||
func (s *KeyService) processAndCreateKeys(
|
||||
groupID uint,
|
||||
keys []string,
|
||||
progressCallback func(processed int),
|
||||
) (addedCount int, ignoredCount int, err error) {
|
||||
// 1. Get existing keys in the group for deduplication
|
||||
var existingKeys []models.APIKey
|
||||
if err := s.DB.Where("group_id = ?", groupID).Select("key_value").Find(&existingKeys).Error; err != nil {
|
||||
return nil, err
|
||||
return 0, 0, err
|
||||
}
|
||||
existingKeyMap := make(map[string]bool)
|
||||
for _, k := range existingKeys {
|
||||
existingKeyMap[k.KeyValue] = true
|
||||
}
|
||||
|
||||
// 3. Prepare new keys for creation
|
||||
// 2. Prepare new keys for creation
|
||||
var newKeysToCreate []models.APIKey
|
||||
uniqueNewKeys := make(map[string]bool)
|
||||
|
||||
@@ -98,14 +121,10 @@ func (s *KeyService) AddMultipleKeys(groupID uint, keysText string) (*AddKeysRes
|
||||
}
|
||||
|
||||
if len(newKeysToCreate) == 0 {
|
||||
return &AddKeysResult{
|
||||
AddedCount: 0,
|
||||
IgnoredCount: len(keys),
|
||||
TotalInGroup: int64(len(existingKeys)),
|
||||
}, nil
|
||||
return 0, len(keys), nil
|
||||
}
|
||||
|
||||
// 4. Use KeyProvider to add keys in chunks
|
||||
// 3. Use KeyProvider to add keys in chunks
|
||||
for i := 0; i < len(newKeysToCreate); i += chunkSize {
|
||||
end := i + chunkSize
|
||||
if end > len(newKeysToCreate) {
|
||||
@@ -113,18 +132,16 @@ func (s *KeyService) AddMultipleKeys(groupID uint, keysText string) (*AddKeysRes
|
||||
}
|
||||
chunk := newKeysToCreate[i:end]
|
||||
if err := s.KeyProvider.AddKeys(groupID, chunk); err != nil {
|
||||
return nil, err
|
||||
return addedCount, len(keys) - addedCount, err
|
||||
}
|
||||
addedCount += len(chunk)
|
||||
|
||||
if progressCallback != nil {
|
||||
progressCallback(i + len(chunk))
|
||||
}
|
||||
}
|
||||
|
||||
// 5. Calculate new total count
|
||||
totalInGroup := int64(len(existingKeys) + len(newKeysToCreate))
|
||||
|
||||
return &AddKeysResult{
|
||||
AddedCount: len(newKeysToCreate),
|
||||
IgnoredCount: len(keys) - len(newKeysToCreate),
|
||||
TotalInGroup: totalInGroup,
|
||||
}, nil
|
||||
return addedCount, len(keys) - addedCount, nil
|
||||
}
|
||||
|
||||
// ParseKeysFromText parses a string of keys from various formats into a string slice.
|
||||
|
@@ -9,12 +9,18 @@ import (
|
||||
)
|
||||
|
||||
const (
|
||||
globalTaskKey = "global_task:key_validation"
|
||||
globalTaskKey = "global_task"
|
||||
ResultTTL = 60 * time.Minute
|
||||
)
|
||||
|
||||
const (
|
||||
TaskTypeKeyValidation = "KEY_VALIDATION"
|
||||
TaskTypeKeyImport = "KEY_IMPORT"
|
||||
)
|
||||
|
||||
// TaskStatus represents the full lifecycle of a long-running task.
|
||||
type TaskStatus struct {
|
||||
TaskType string `json:"task_type"`
|
||||
IsRunning bool `json:"is_running"`
|
||||
GroupName string `json:"group_name,omitempty"`
|
||||
Processed int `json:"processed"`
|
||||
@@ -39,7 +45,7 @@ func NewTaskService(store store.Store) *TaskService {
|
||||
}
|
||||
|
||||
// StartTask attempts to start a new task. It returns an error if a task is already running.
|
||||
func (s *TaskService) StartTask(groupName string, total int, timeout time.Duration) (*TaskStatus, error) {
|
||||
func (s *TaskService) StartTask(taskType, groupName string, total int, timeout time.Duration) (*TaskStatus, error) {
|
||||
currentStatus, err := s.GetTaskStatus()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to check current task status before starting a new one: %w", err)
|
||||
@@ -50,6 +56,7 @@ func (s *TaskService) StartTask(groupName string, total int, timeout time.Durati
|
||||
}
|
||||
|
||||
status := &TaskStatus{
|
||||
TaskType: taskType,
|
||||
IsRunning: true,
|
||||
GroupName: groupName,
|
||||
Total: total,
|
||||
|
@@ -62,7 +62,7 @@ export const keysApi = {
|
||||
return res.data;
|
||||
},
|
||||
|
||||
// 批量添加密钥
|
||||
// 批量添加密钥-已弃用
|
||||
async addMultipleKeys(
|
||||
group_id: number,
|
||||
keys_text: string
|
||||
@@ -78,6 +78,15 @@ export const keysApi = {
|
||||
return res.data;
|
||||
},
|
||||
|
||||
// 异步批量添加密钥
|
||||
async addKeysAsync(group_id: number, keys_text: string): Promise<TaskInfo> {
|
||||
const res = await http.post("/keys/add-async", {
|
||||
group_id,
|
||||
keys_text,
|
||||
});
|
||||
return res.data;
|
||||
},
|
||||
|
||||
// 测试密钥
|
||||
async testKeys(
|
||||
group_id: number,
|
||||
|
@@ -5,7 +5,7 @@ import { appState } from "@/utils/app-state";
|
||||
import { NButton, NCard, NProgress, NText, useMessage } from "naive-ui";
|
||||
import { onBeforeUnmount, onMounted, ref, watch } from "vue";
|
||||
|
||||
const taskInfo = ref<TaskInfo>({ is_running: false });
|
||||
const taskInfo = ref<TaskInfo>({ is_running: false, task_type: "KEY_VALIDATION" });
|
||||
const visible = ref(false);
|
||||
let pollTimer: number | null = null;
|
||||
let isPolling = false; // 添加标志位
|
||||
@@ -46,8 +46,15 @@ async function pollOnce() {
|
||||
if (task.result) {
|
||||
const lastTask = localStorage.getItem("last_closed_task");
|
||||
if (lastTask !== task.finished_at) {
|
||||
const { total_keys, valid_keys, invalid_keys } = task.result;
|
||||
const msg = `任务已完成,处理了 ${total_keys} 个密钥,其中 ${valid_keys} 个有效密钥,${invalid_keys} 个无效密钥。`;
|
||||
let msg = "任务已完成。";
|
||||
if (task.task_type === "KEY_VALIDATION") {
|
||||
const result = task.result as import("@/types/models").KeyValidationResult;
|
||||
msg = `密钥验证完成,处理了 ${result.total_keys} 个密钥,其中 ${result.valid_keys} 个有效,${result.invalid_keys} 个无效。`;
|
||||
} else if (task.task_type === "KEY_IMPORT") {
|
||||
const result = task.result as import("@/types/models").KeyImportResult;
|
||||
msg = `密钥导入完成,成功添加 ${result.added_count} 个密钥,忽略了 ${result.ignored_count} 个。`;
|
||||
}
|
||||
|
||||
message.info(msg, {
|
||||
closable: true,
|
||||
duration: 0,
|
||||
@@ -92,6 +99,20 @@ function getProgressText(): string {
|
||||
function handleClose() {
|
||||
visible.value = false;
|
||||
}
|
||||
|
||||
function getTaskTitle(): string {
|
||||
if (!taskInfo.value) {
|
||||
return "正在处理任务...";
|
||||
}
|
||||
switch (taskInfo.value.task_type) {
|
||||
case "KEY_VALIDATION":
|
||||
return `正在验证分组 [${taskInfo.value.group_name}] 的密钥`;
|
||||
case "KEY_IMPORT":
|
||||
return `正在向分组 [${taskInfo.value.group_name}] 导入密钥`;
|
||||
default:
|
||||
return "正在处理任务...";
|
||||
}
|
||||
}
|
||||
</script>
|
||||
|
||||
<template>
|
||||
@@ -102,7 +123,7 @@ function handleClose() {
|
||||
<span class="progress-icon">⚡</span>
|
||||
<div class="progress-details">
|
||||
<n-text strong class="progress-title">
|
||||
正在处理分组 {{ taskInfo.group_name }} 的任务
|
||||
{{ getTaskTitle() }}
|
||||
</n-text>
|
||||
<n-text depth="3" class="progress-subtitle">
|
||||
{{ getProgressText() }} ({{ getProgressPercentage() }}%)
|
||||
|
@@ -1,5 +1,6 @@
|
||||
<script setup lang="ts">
|
||||
import { keysApi } from "@/api/keys";
|
||||
import { appState } from "@/utils/app-state";
|
||||
import { Close } from "@vicons/ionicons5";
|
||||
import { NButton, NCard, NInput, NModal } from "naive-ui";
|
||||
import { ref, watch } from "vue";
|
||||
@@ -51,10 +52,11 @@ async function handleSubmit() {
|
||||
try {
|
||||
loading.value = true;
|
||||
|
||||
await keysApi.addMultipleKeys(props.groupId, keysText.value);
|
||||
|
||||
emit("success");
|
||||
await keysApi.addKeysAsync(props.groupId, keysText.value);
|
||||
resetForm();
|
||||
handleClose();
|
||||
window.$message.success("密钥导入任务已开始,请稍后在下方查看进度。");
|
||||
appState.taskPollingTrigger++;
|
||||
} finally {
|
||||
loading.value = false;
|
||||
}
|
||||
|
@@ -75,18 +75,29 @@ export interface RequestStats {
|
||||
failure_rate: number;
|
||||
}
|
||||
|
||||
export type TaskType = "KEY_VALIDATION" | "KEY_IMPORT";
|
||||
|
||||
export interface KeyValidationResult {
|
||||
invalid_keys: number;
|
||||
total_keys: number;
|
||||
valid_keys: number;
|
||||
}
|
||||
|
||||
export interface KeyImportResult {
|
||||
added_count: number;
|
||||
ignored_count: number;
|
||||
}
|
||||
|
||||
export interface TaskInfo {
|
||||
task_type: TaskType;
|
||||
is_running: boolean;
|
||||
group_name?: string;
|
||||
processed?: number;
|
||||
total?: number;
|
||||
started_at?: string;
|
||||
finished_at?: string;
|
||||
result?: {
|
||||
invalid_keys: number;
|
||||
total_keys: number;
|
||||
valid_keys: number;
|
||||
};
|
||||
result?: KeyValidationResult | KeyImportResult;
|
||||
error?: string;
|
||||
}
|
||||
|
||||
// Based on backend response
|
||||
|
Reference in New Issue
Block a user