before update store

internal/repository/key_selector.go

@@ -11,7 +11,6 @@ import (
 	"io"
 	"sort"
 	"strconv"
-	"strings"
 	"time"
 
 	"gorm.io/gorm"
@@ -82,20 +81,13 @@ func (r *gormKeyRepository) SelectOneActiveKey(group *models.KeyGroup) (*models.
 
 // SelectOneActiveKeyFromBasePool 为智能聚合模式设计的全新轮询器。
 func (r *gormKeyRepository) SelectOneActiveKeyFromBasePool(pool *BasePool) (*models.APIKey, *models.KeyGroup, error) {
-	protocol := "default"
-	if pool.Protocol != "" {
-		protocol = string(pool.Protocol)
-	}
 	// 生成唯一的池ID，确保不同请求组合的轮询状态相互隔离
-	poolID := generatePoolID(pool.CandidateGroups, protocol)
-	log := r.logger.WithField("pool_id", poolID).WithField("protocol", protocol)
+	poolID := generatePoolID(pool.CandidateGroups)
+	log := r.logger.WithField("pool_id", poolID)
 
 	if err := r.ensureBasePoolCacheExists(pool, poolID); err != nil {
 		log.WithError(err).Error("Failed to ensure BasePool cache exists.")
-		if errors.Is(err, gorm.ErrRecordNotFound) {
-			return nil, nil, err
-		}
-		return nil, nil, fmt.Errorf("unexpected error while ensuring base pool cache: %w", err)
+		return nil, nil, err
 	}
 
 	var keyIDStr string
@@ -154,65 +146,78 @@ func (r *gormKeyRepository) SelectOneActiveKeyFromBasePool(pool *BasePool) (*mod
 // ensureBasePoolCacheExists 动态创建 BasePool 的 Redis 结构
 func (r *gormKeyRepository) ensureBasePoolCacheExists(pool *BasePool, poolID string) error {
 	listKey := fmt.Sprintf(BasePoolSequential, poolID)
+
+	// --- [逻辑优化] 提前处理“毒丸”，让逻辑更清晰 ---
 	exists, err := r.store.Exists(listKey)
 	if err != nil {
-		r.logger.WithError(err).Errorf("Failed to check existence of basepool key: %s", listKey)
-		return err
+		r.logger.WithError(err).Errorf("Failed to check existence for pool_id '%s'", poolID)
+		return err // 直接返回读取错误
 	}
 	if exists {
 		val, err := r.store.LIndex(listKey, 0)
 		if err != nil {
-			return err
+			// 如果连 LIndex 都失败，说明缓存可能已损坏，允许重建
+			r.logger.WithError(err).Warnf("Cache for pool_id '%s' exists but is unreadable. Forcing rebuild.", poolID)
+		} else {
+			if val == EmptyPoolPlaceholder {
+				return gorm.ErrRecordNotFound // 已知为空，直接返回
+			}
+			return nil // 缓存有效，直接返回
 		}
-		if val == EmptyPoolPlaceholder {
-			return gorm.ErrRecordNotFound
-		}
-		return nil
 	}
+	// --- [锁机制优化] 增加分布式锁，防止并发构建时的惊群效应 ---
 	lockKey := fmt.Sprintf("lock:basepool:%s", poolID)
-	acquired, err := r.store.SetNX(lockKey, []byte("1"), 10*time.Second)
+	acquired, err := r.store.SetNX(lockKey, []byte("1"), 10*time.Second) // 10秒锁超时
 	if err != nil {
-		r.logger.WithError(err).Errorf("Failed to acquire distributed lock for basepool build: %s", lockKey)
+		r.logger.WithError(err).Error("Failed to attempt acquiring distributed lock for basepool build.")
 		return err
 	}
 	if !acquired {
+		// 未获取到锁，等待一小段时间后重试，让持有锁的协程完成构建
 		time.Sleep(100 * time.Millisecond)
 		return r.ensureBasePoolCacheExists(pool, poolID)
 	}
-	defer r.store.Del(lockKey)
+	defer r.store.Del(lockKey) // 确保在函数退出时释放锁
+	// 双重检查，防止在获取锁的间隙，已有其他协程完成了构建
 	if exists, _ := r.store.Exists(listKey); exists {
 		return nil
 	}
-	r.logger.Infof("BasePool cache for pool_id '%s' not found. Building now...", poolID)
+	r.logger.Infof("BasePool cache for pool_id '%s' not found or is unreadable. Building now...", poolID)
 	var allActiveKeyIDs []string
 	lruMembers := make(map[string]float64)
-
 	for _, group := range pool.CandidateGroups {
 		activeKeySetKey := fmt.Sprintf(KeyGroup, group.ID)
 		groupKeyIDs, err := r.store.SMembers(activeKeySetKey)
+
+		// --- [核心修正] ---
+		// 这是整个问题的根源。我们绝不能在读取失败时，默默地`continue`。
+		// 任何读取源数据的失败，都必须被视为一次构建过程的彻底失败，并立即中止。
 		if err != nil {
-			r.logger.WithError(err).Warnf("Failed to get active keys for group %d during BasePool build", group.ID)
-			continue
+			r.logger.WithError(err).Errorf("FATAL: Failed to read active keys for group %d during BasePool build. Aborting build process for pool_id '%s'.", group.ID, poolID)
+			// 返回这个瞬时错误。这会导致本次请求失败，但绝不会写入“毒丸”，
+			// 从而给了下一次请求一个全新的、成功的机会。
+			return err
 		}
+		// 只有在 SMembers 成功时，才继续处理
+		allActiveKeyIDs = append(allActiveKeyIDs, groupKeyIDs...)
 		for _, keyIDStr := range groupKeyIDs {
 			keyID, _ := strconv.ParseUint(keyIDStr, 10, 64)
-
 			_, mapping, err := r.getKeyDetailsFromCache(uint(keyID), group.ID)
-			if err != nil {
-				if errors.Is(err, store.ErrNotFound) || strings.Contains(err.Error(), "failed to get") {
-					r.logger.WithError(err).Warnf("Cache inconsistency detected for KeyID %s in GroupID %d. Skipping.", keyIDStr, group.ID)
-					continue
-				} else {
-					return err
+			if err == nil && mapping != nil {
+				var score float64
+				if mapping.LastUsedAt != nil {
+					score = float64(mapping.LastUsedAt.UnixMilli())
 				}
-			}
-			allActiveKeyIDs = append(allActiveKeyIDs, keyIDStr)
-			if mapping != nil && mapping.LastUsedAt != nil {
-				lruMembers[keyIDStr] = float64(mapping.LastUsedAt.UnixMilli())
+				lruMembers[keyIDStr] = score
 			}
 		}
 	}
+
+	// --- [逻辑修正] ---
+	// 只有在“我们成功读取了所有数据，但发现数据本身是空的”这种情况下，
+	// 才允许写入“毒丸”。
 	if len(allActiveKeyIDs) == 0 {
+		r.logger.Warnf("No active keys found for any candidate groups for pool_id '%s'. Setting empty pool placeholder.", poolID)
 		pipe := r.store.Pipeline()
 		pipe.LPush(listKey, EmptyPoolPlaceholder)
 		pipe.Expire(listKey, EmptyCacheTTL)
@@ -221,16 +226,23 @@ func (r *gormKeyRepository) ensureBasePoolCacheExists(pool *BasePool, poolID str
 		}
 		return gorm.ErrRecordNotFound
 	}
+	// 使用管道填充所有轮询结构
 	pipe := r.store.Pipeline()
+	// 1. 顺序
 	pipe.LPush(fmt.Sprintf(BasePoolSequential, poolID), toInterfaceSlice(allActiveKeyIDs)...)
+	// 2. 随机
 	pipe.SAdd(fmt.Sprintf(BasePoolRandomMain, poolID), toInterfaceSlice(allActiveKeyIDs)...)
+
+	// 设置合理的过期时间，例如5分钟，以防止孤儿数据
 	pipe.Expire(fmt.Sprintf(BasePoolSequential, poolID), CacheTTL)
 	pipe.Expire(fmt.Sprintf(BasePoolRandomMain, poolID), CacheTTL)
 	pipe.Expire(fmt.Sprintf(BasePoolRandomCooldown, poolID), CacheTTL)
 	pipe.Expire(fmt.Sprintf(BasePoolLRU, poolID), CacheTTL)
+
 	if err := pipe.Exec(); err != nil {
 		return err
 	}
+
 	if len(lruMembers) > 0 {
 		r.store.ZAdd(fmt.Sprintf(BasePoolLRU, poolID), lruMembers)
 	}
@@ -246,7 +258,7 @@ func (r *gormKeyRepository) updateKeyUsageTimestampForPool(poolID string, keyID
 }
 
 // generatePoolID 根据候选组ID列表生成一个稳定的、唯一的字符串ID
-func generatePoolID(groups []*models.KeyGroup, protocol string) string {
+func generatePoolID(groups []*models.KeyGroup) string {
 	ids := make([]int, len(groups))
 	for i, g := range groups {
 		ids[i] = int(g.ID)
@@ -254,7 +266,7 @@ func generatePoolID(groups []*models.KeyGroup, protocol string) string {
 	sort.Ints(ids)
 
 	h := sha1.New()
-	io.WriteString(h, fmt.Sprintf("protocol:%s;groups:%v", protocol, ids))
+	io.WriteString(h, fmt.Sprintf("%v", ids))
 	return fmt.Sprintf("%x", h.Sum(nil))
 }