Fix basepool & 优化 repo

2025-11-23 22:42:58 +08:00
parent 2b0b9b67dc
commit 6c7283d51b
16 changed files with 1312 additions and 723 deletions
--- a/internal/repository/key_cache.go
+++ b/internal/repository/key_cache.go
@@ -1,4 +1,4 @@
-// Filename: internal/repository/key_cache.go
+// Filename: internal/repository/key_cache.go (最终定稿)
 package repository

 import (
@@ -9,6 +9,7 @@ import (
 	"strconv"
 )

+// --- Redis Key 常量定义 ---
 const (
 	KeyGroup               = "group:%d:keys:active"
 	KeyDetails             = "key:%d:details"
@@ -23,13 +24,16 @@ const (
 	BasePoolRandomCooldown = "basepool:%s:keys:random:cooldown"
 )

+// LoadAllKeysToStore 从数据库加载所有密钥和映射关系，并完整重建Redis缓存。
 func (r *gormKeyRepository) LoadAllKeysToStore(ctx context.Context) error {
-	r.logger.Info("Starting to load all keys and associations into cache, including polling structures...")
+	r.logger.Info("Starting full cache rebuild for all keys and polling structures.")
+
 	var allMappings []*models.GroupAPIKeyMapping
 	if err := r.db.Preload("APIKey").Find(&allMappings).Error; err != nil {
-		return fmt.Errorf("failed to load all mappings with APIKeys from DB: %w", err)
+		return fmt.Errorf("failed to load mappings with preloaded APIKeys: %w", err)
 	}

+	// 1. 批量解密所有涉及的密钥
 	keyMap := make(map[uint]*models.APIKey)
 	for _, m := range allMappings {
 		if m.APIKey != nil {
@@ -41,16 +45,16 @@ func (r *gormKeyRepository) LoadAllKeysToStore(ctx context.Context) error {
 		keysToDecrypt = append(keysToDecrypt, *k)
 	}
 	if err := r.decryptKeys(keysToDecrypt); err != nil {
-		r.logger.WithError(err).Error("Critical error during cache preload: batch decryption failed.")
+		r.logger.WithError(err).Error("Batch decryption failed during cache rebuild.")
+		// 即使解密失败，也继续尝试加载未加密或已解密的部分
 	}
 	decryptedKeyMap := make(map[uint]models.APIKey)
 	for _, k := range keysToDecrypt {
 		decryptedKeyMap[k.ID] = k
 	}

-	activeKeysByGroup := make(map[uint][]*models.GroupAPIKeyMapping)
-	pipe := r.store.Pipeline(context.Background())
-	detailsToSet := make(map[string][]byte)
+	// 2. 清理所有分组的旧轮询结构
+	pipe := r.store.Pipeline(ctx)
 	var allGroups []*models.KeyGroup
 	if err := r.db.Find(&allGroups).Error; err == nil {
 		for _, group := range allGroups {
@@ -63,26 +67,41 @@ func (r *gormKeyRepository) LoadAllKeysToStore(ctx context.Context) error {
 			)
 		}
 	} else {
-		r.logger.WithError(err).Error("Failed to get all groups for cache cleanup")
+		r.logger.WithError(err).Error("Failed to get groups for cache cleanup; proceeding with rebuild.")
 	}

+	// 3. 准备批量更新数据
+	activeKeysByGroup := make(map[uint][]*models.GroupAPIKeyMapping)
+	detailsToSet := make(map[string]any)
+
 	for _, mapping := range allMappings {
 		if mapping.APIKey == nil {
 			continue
 		}
 		decryptedKey, ok := decryptedKeyMap[mapping.APIKeyID]
 		if !ok {
-			continue
+			continue // 跳过解密失败的密钥
 		}
+
+		// 准备 KeyDetails 和 KeyMapping 的 MSet 数据
 		keyJSON, _ := json.Marshal(decryptedKey)
 		detailsToSet[fmt.Sprintf(KeyDetails, decryptedKey.ID)] = keyJSON
 		mappingJSON, _ := json.Marshal(mapping)
 		detailsToSet[fmt.Sprintf(KeyMapping, mapping.KeyGroupID, decryptedKey.ID)] = mappingJSON
+
 		if mapping.Status == models.StatusActive {
 			activeKeysByGroup[mapping.KeyGroupID] = append(activeKeysByGroup[mapping.KeyGroupID], mapping)
 		}
 	}

+	// 4. 使用 MSet 批量写入详情和映射缓存
+	if len(detailsToSet) > 0 {
+		if err := r.store.MSet(ctx, detailsToSet); err != nil {
+			r.logger.WithError(err).Error("Failed to MSet key details and mappings during cache rebuild.")
+		}
+	}
+
+	// 5. 在Pipeline中重建所有分组的轮询结构
 	for groupID, activeMappings := range activeKeysByGroup {
 		if len(activeMappings) == 0 {
 			continue
@@ -101,22 +120,19 @@ func (r *gormKeyRepository) LoadAllKeysToStore(ctx context.Context) error {
 		pipe.SAdd(fmt.Sprintf(KeyGroup, groupID), activeKeyIDs...)
 		pipe.LPush(fmt.Sprintf(KeyGroupSequential, groupID), activeKeyIDs...)
 		pipe.SAdd(fmt.Sprintf(KeyGroupRandomMain, groupID), activeKeyIDs...)
-		go r.store.ZAdd(context.Background(), fmt.Sprintf(KeyGroupLRU, groupID), lruMembers)
+		pipe.ZAdd(fmt.Sprintf(KeyGroupLRU, groupID), lruMembers)
 	}

+	// 6. 执行Pipeline
 	if err := pipe.Exec(); err != nil {
-		return fmt.Errorf("failed to execute pipeline for cache rebuild: %w", err)
-	}
-	for key, value := range detailsToSet {
-		if err := r.store.Set(context.Background(), key, value, 0); err != nil {
-			r.logger.WithError(err).Warnf("Failed to set key detail in cache: %s", key)
-		}
+		return fmt.Errorf("pipeline execution for polling structures failed: %w", err)
 	}

-	r.logger.Info("Cache rebuild complete, including all polling structures.")
+	r.logger.Info("Full cache rebuild completed successfully.")
 	return nil
 }

+// updateStoreCacheForKey 更新单个APIKey的详情缓存 (K-V)。
 func (r *gormKeyRepository) updateStoreCacheForKey(key *models.APIKey) error {
 	if err := r.decryptKey(key); err != nil {
 		return fmt.Errorf("failed to decrypt key %d for cache update: %w", key.ID, err)
@@ -128,78 +144,101 @@ func (r *gormKeyRepository) updateStoreCacheForKey(key *models.APIKey) error {
 	return r.store.Set(context.Background(), fmt.Sprintf(KeyDetails, key.ID), keyJSON, 0)
 }

+// removeStoreCacheForKey 从所有缓存结构中彻底移除一个APIKey。
 func (r *gormKeyRepository) removeStoreCacheForKey(ctx context.Context, key *models.APIKey) error {
 	groupIDs, err := r.GetGroupsForKey(ctx, key.ID)
 	if err != nil {
-		r.logger.Warnf("failed to get groups for key %d to clean up cache lists: %v", key.ID, err)
+		r.logger.WithError(err).Warnf("Failed to get groups for key %d during cache removal, cleanup may be partial.", key.ID)
 	}

 	pipe := r.store.Pipeline(ctx)
 	pipe.Del(fmt.Sprintf(KeyDetails, key.ID))

+	keyIDStr := strconv.FormatUint(uint64(key.ID), 10)
 	for _, groupID := range groupIDs {
 		pipe.Del(fmt.Sprintf(KeyMapping, groupID, key.ID))
-
-		keyIDStr := strconv.FormatUint(uint64(key.ID), 10)
 		pipe.SRem(fmt.Sprintf(KeyGroup, groupID), keyIDStr)
 		pipe.LRem(fmt.Sprintf(KeyGroupSequential, groupID), 0, keyIDStr)
 		pipe.SRem(fmt.Sprintf(KeyGroupRandomMain, groupID), keyIDStr)
 		pipe.SRem(fmt.Sprintf(KeyGroupRandomCooldown, groupID), keyIDStr)
-		go r.store.ZRem(context.Background(), fmt.Sprintf(KeyGroupLRU, groupID), keyIDStr)
+		pipe.ZRem(fmt.Sprintf(KeyGroupLRU, groupID), keyIDStr)
 	}
+
 	return pipe.Exec()
 }

+// updateStoreCacheForMapping 根据单个映射关系的状态，原子性地更新所有相关的缓存结构。
 func (r *gormKeyRepository) updateStoreCacheForMapping(mapping *models.GroupAPIKeyMapping) error {
-	pipe := r.store.Pipeline(context.Background())
-	activeKeyListKey := fmt.Sprintf("group:%d:keys:active", mapping.KeyGroupID)
-	pipe.LRem(activeKeyListKey, 0, mapping.APIKeyID)
+	keyIDStr := strconv.FormatUint(uint64(mapping.APIKeyID), 10)
+	groupID := mapping.KeyGroupID
+	ctx := context.Background()
+
+	pipe := r.store.Pipeline(ctx)
+
+	// 统一、无条件地从所有轮询结构中移除，确保状态清洁
+	pipe.SRem(fmt.Sprintf(KeyGroup, groupID), keyIDStr)
+	pipe.LRem(fmt.Sprintf(KeyGroupSequential, groupID), 0, keyIDStr)
+	pipe.SRem(fmt.Sprintf(KeyGroupRandomMain, groupID), keyIDStr)
+	pipe.SRem(fmt.Sprintf(KeyGroupRandomCooldown, groupID), keyIDStr)
+	pipe.ZRem(fmt.Sprintf(KeyGroupLRU, groupID), keyIDStr)
+
+	// 如果新状态是 Active，则重新添加到所有轮询结构中
 	if mapping.Status == models.StatusActive {
-		pipe.LPush(activeKeyListKey, mapping.APIKeyID)
+		pipe.SAdd(fmt.Sprintf(KeyGroup, groupID), keyIDStr)
+		pipe.LPush(fmt.Sprintf(KeyGroupSequential, groupID), keyIDStr)
+		pipe.SAdd(fmt.Sprintf(KeyGroupRandomMain, groupID), keyIDStr)
+
+		var score float64
+		if mapping.LastUsedAt != nil {
+			score = float64(mapping.LastUsedAt.UnixMilli())
+		}
+		pipe.ZAdd(fmt.Sprintf(KeyGroupLRU, groupID), map[string]float64{keyIDStr: score})
 	}
+
+	// 无论状态如何，都更新映射详情的 K-V 缓存
+	mappingJSON, err := json.Marshal(mapping)
+	if err != nil {
+		return fmt.Errorf("failed to marshal mapping: %w", err)
+	}
+	pipe.Set(fmt.Sprintf(KeyMapping, groupID, mapping.APIKeyID), mappingJSON, 0)
+
 	return pipe.Exec()
 }

+// HandleCacheUpdateEventBatch 批量、原子性地更新多个映射关系的缓存。
 func (r *gormKeyRepository) HandleCacheUpdateEventBatch(ctx context.Context, mappings []*models.GroupAPIKeyMapping) error {
 	if len(mappings) == 0 {
 		return nil
 	}
-	groupUpdates := make(map[uint]struct {
-		ToAdd    []interface{}
-		ToRemove []interface{}
-	})
+
+	pipe := r.store.Pipeline(ctx)
+
 	for _, mapping := range mappings {
 		keyIDStr := strconv.FormatUint(uint64(mapping.APIKeyID), 10)
-		update, ok := groupUpdates[mapping.KeyGroupID]
-		if !ok {
-			update = struct {
-				ToAdd    []interface{}
-				ToRemove []interface{}
-			}{}
-		}
+		groupID := mapping.KeyGroupID
+
+		// 对于批处理中的每一个mapping，都执行完整的、正确的“先删后增”逻辑
+		pipe.SRem(fmt.Sprintf(KeyGroup, groupID), keyIDStr)
+		pipe.LRem(fmt.Sprintf(KeyGroupSequential, groupID), 0, keyIDStr)
+		pipe.SRem(fmt.Sprintf(KeyGroupRandomMain, groupID), keyIDStr)
+		pipe.SRem(fmt.Sprintf(KeyGroupRandomCooldown, groupID), keyIDStr)
+		pipe.ZRem(fmt.Sprintf(KeyGroupLRU, groupID), keyIDStr)
+
 		if mapping.Status == models.StatusActive {
-			update.ToRemove = append(update.ToRemove, keyIDStr)
-			update.ToAdd = append(update.ToAdd, keyIDStr)
-		} else {
-			update.ToRemove = append(update.ToRemove, keyIDStr)
-		}
-		groupUpdates[mapping.KeyGroupID] = update
-	}
-	pipe := r.store.Pipeline(context.Background())
-	var pipelineError error
-	for groupID, updates := range groupUpdates {
-		activeKeyListKey := fmt.Sprintf("group:%d:keys:active", groupID)
-		if len(updates.ToRemove) > 0 {
-			for _, keyID := range updates.ToRemove {
-				pipe.LRem(activeKeyListKey, 0, keyID)
+			pipe.SAdd(fmt.Sprintf(KeyGroup, groupID), keyIDStr)
+			pipe.LPush(fmt.Sprintf(KeyGroupSequential, groupID), keyIDStr)
+			pipe.SAdd(fmt.Sprintf(KeyGroupRandomMain, groupID), keyIDStr)
+
+			var score float64
+			if mapping.LastUsedAt != nil {
+				score = float64(mapping.LastUsedAt.UnixMilli())
 			}
+			pipe.ZAdd(fmt.Sprintf(KeyGroupLRU, groupID), map[string]float64{keyIDStr: score})
 		}
-		if len(updates.ToAdd) > 0 {
-			pipe.LPush(activeKeyListKey, updates.ToAdd...)
-		}
+
+		mappingJSON, _ := json.Marshal(mapping) // 在批处理中忽略单个marshal错误，以保证大部分更新成功
+		pipe.Set(fmt.Sprintf(KeyMapping, groupID, mapping.APIKeyID), mappingJSON, 0)
 	}
-	if err := pipe.Exec(); err != nil {
-		pipelineError = fmt.Errorf("redis pipeline execution failed: %w", err)
-	}
-	return pipelineError
+
+	return pipe.Exec()
 }