428 lines
15 KiB
Go
428 lines
15 KiB
Go
// Filename: internal/repository/key_selector.go
|
||
package repository
|
||
|
||
import (
|
||
"context"
|
||
"crypto/sha1"
|
||
"encoding/json"
|
||
"errors"
|
||
"fmt"
|
||
"gemini-balancer/internal/models"
|
||
"gemini-balancer/internal/store"
|
||
"io"
|
||
"sort"
|
||
"strconv"
|
||
"time"
|
||
|
||
"gorm.io/gorm"
|
||
)
|
||
|
||
const (
|
||
CacheTTL = 5 * time.Minute
|
||
EmptyCacheTTL = 1 * time.Minute
|
||
)
|
||
|
||
// SelectOneActiveKey 根据指定的轮询策略，从单个密钥组缓存中选取一个可用的API密钥。
|
||
func (r *gormKeyRepository) SelectOneActiveKey(ctx context.Context, group *models.KeyGroup) (*models.APIKey, *models.GroupAPIKeyMapping, error) {
|
||
if group == nil {
|
||
return nil, nil, fmt.Errorf("group cannot be nil")
|
||
}
|
||
var keyIDStr string
|
||
var err error
|
||
switch group.PollingStrategy {
|
||
case models.StrategySequential:
|
||
sequentialKey := fmt.Sprintf(KeyGroupSequential, group.ID)
|
||
keyIDStr, err = r.store.Rotate(ctx, sequentialKey)
|
||
case models.StrategyWeighted:
|
||
lruKey := fmt.Sprintf(KeyGroupLRU, group.ID)
|
||
results, zerr := r.store.ZRange(ctx, lruKey, 0, 0)
|
||
if zerr == nil {
|
||
if len(results) > 0 {
|
||
keyIDStr = results[0]
|
||
} else {
|
||
zerr = gorm.ErrRecordNotFound
|
||
}
|
||
}
|
||
err = zerr
|
||
case models.StrategyRandom:
|
||
mainPoolKey := fmt.Sprintf(KeyGroupRandomMain, group.ID)
|
||
cooldownPoolKey := fmt.Sprintf(KeyGroupRandomCooldown, group.ID)
|
||
keyIDStr, err = r.store.PopAndCycleSetMember(ctx, mainPoolKey, cooldownPoolKey)
|
||
default:
|
||
activeKeySetKey := fmt.Sprintf(KeyGroup, group.ID)
|
||
keyIDStr, err = r.store.SRandMember(ctx, activeKeySetKey)
|
||
}
|
||
if err != nil {
|
||
if errors.Is(err, store.ErrNotFound) || errors.Is(err, gorm.ErrRecordNotFound) {
|
||
return nil, nil, gorm.ErrRecordNotFound
|
||
}
|
||
r.logger.WithError(err).Errorf("Failed to select key for group %d with strategy %s", group.ID, group.PollingStrategy)
|
||
return nil, nil, err
|
||
}
|
||
if keyIDStr == "" {
|
||
return nil, nil, gorm.ErrRecordNotFound
|
||
}
|
||
keyID, parseErr := strconv.ParseUint(keyIDStr, 10, 64)
|
||
if parseErr != nil {
|
||
r.logger.WithError(parseErr).Errorf("Invalid key ID format in group %d cache: %s", group.ID, keyIDStr)
|
||
return nil, nil, fmt.Errorf("invalid key ID in cache: %w", parseErr)
|
||
}
|
||
apiKey, mapping, err := r.getKeyDetailsFromCache(ctx, uint(keyID), group.ID)
|
||
if err != nil {
|
||
r.logger.WithError(err).Warnf("Cache inconsistency for key ID %d in group %d", keyID, group.ID)
|
||
return nil, nil, err
|
||
}
|
||
if group.PollingStrategy == models.StrategyWeighted {
|
||
go func() {
|
||
updateCtx, cancel := r.withTimeout(5 * time.Second)
|
||
defer cancel()
|
||
r.UpdateKeyUsageTimestamp(updateCtx, group.ID, uint(keyID))
|
||
}()
|
||
}
|
||
return apiKey, mapping, nil
|
||
}
|
||
|
||
// SelectOneActiveKeyFromBasePool 从智能聚合池中选取一个可用Key。
|
||
func (r *gormKeyRepository) SelectOneActiveKeyFromBasePool(ctx context.Context, pool *BasePool) (*models.APIKey, *models.KeyGroup, error) {
|
||
if pool == nil || len(pool.CandidateGroups) == 0 {
|
||
return nil, nil, fmt.Errorf("invalid or empty base pool configuration")
|
||
}
|
||
poolID := r.generatePoolID(pool.CandidateGroups)
|
||
log := r.logger.WithField("pool_id", poolID)
|
||
if err := r.ensureBasePoolCacheExists(ctx, pool, poolID); err != nil {
|
||
if !errors.Is(err, gorm.ErrRecordNotFound) {
|
||
log.WithError(err).Error("Failed to ensure BasePool cache exists")
|
||
}
|
||
return nil, nil, err
|
||
}
|
||
var keyIDStr string
|
||
var err error
|
||
switch pool.PollingStrategy {
|
||
case models.StrategySequential:
|
||
sequentialKey := fmt.Sprintf(BasePoolSequential, poolID)
|
||
keyIDStr, err = r.store.Rotate(ctx, sequentialKey)
|
||
case models.StrategyWeighted:
|
||
lruKey := fmt.Sprintf(BasePoolLRU, poolID)
|
||
results, zerr := r.store.ZRange(ctx, lruKey, 0, 0)
|
||
if zerr == nil {
|
||
if len(results) > 0 {
|
||
keyIDStr = results[0]
|
||
} else {
|
||
zerr = gorm.ErrRecordNotFound
|
||
}
|
||
}
|
||
err = zerr
|
||
case models.StrategyRandom:
|
||
mainPoolKey := fmt.Sprintf(BasePoolRandomMain, poolID)
|
||
cooldownPoolKey := fmt.Sprintf(BasePoolRandomCooldown, poolID)
|
||
keyIDStr, err = r.store.PopAndCycleSetMember(ctx, mainPoolKey, cooldownPoolKey)
|
||
default:
|
||
log.Warnf("Unknown polling strategy '%s'. Using sequential as fallback.", pool.PollingStrategy)
|
||
sequentialKey := fmt.Sprintf(BasePoolSequential, poolID)
|
||
keyIDStr, err = r.store.Rotate(ctx, sequentialKey)
|
||
}
|
||
if err != nil {
|
||
if errors.Is(err, store.ErrNotFound) || errors.Is(err, gorm.ErrRecordNotFound) {
|
||
return nil, nil, gorm.ErrRecordNotFound
|
||
}
|
||
log.WithError(err).Errorf("Failed to select key from BasePool with strategy %s", pool.PollingStrategy)
|
||
return nil, nil, err
|
||
}
|
||
if keyIDStr == "" {
|
||
return nil, nil, gorm.ErrRecordNotFound
|
||
}
|
||
go func() {
|
||
bgCtx, cancel := r.withTimeout(5 * time.Second)
|
||
defer cancel()
|
||
r.refreshBasePoolHeartbeat(bgCtx, poolID)
|
||
}()
|
||
keyID, parseErr := strconv.ParseUint(keyIDStr, 10, 64)
|
||
if parseErr != nil {
|
||
log.WithError(parseErr).Errorf("Invalid key ID format in BasePool cache: %s", keyIDStr)
|
||
return nil, nil, fmt.Errorf("invalid key ID in cache: %w", parseErr)
|
||
}
|
||
keyToGroupMapKey := fmt.Sprintf("basepool:%s:key_to_group", poolID)
|
||
groupIDStr, err := r.store.HGet(ctx, keyToGroupMapKey, keyIDStr)
|
||
if err != nil {
|
||
log.WithError(err).Errorf("Cache inconsistency: KeyID %d found in pool but not in key-to-group map", keyID)
|
||
return nil, nil, errors.New("cache inconsistency: key has no origin group mapping")
|
||
}
|
||
groupID, parseErr := strconv.ParseUint(groupIDStr, 10, 64)
|
||
if parseErr != nil {
|
||
log.WithError(parseErr).Errorf("Invalid group ID format in key-to-group map for key %d: %s", keyID, groupIDStr)
|
||
return nil, nil, errors.New("cache inconsistency: invalid group id in mapping")
|
||
}
|
||
apiKey, _, err := r.getKeyDetailsFromCache(ctx, uint(keyID), uint(groupID))
|
||
if err != nil {
|
||
log.WithError(err).Warnf("Cache inconsistency: Failed to get details for key %d in mapped group %d", keyID, groupID)
|
||
return nil, nil, err
|
||
}
|
||
var originGroup *models.KeyGroup
|
||
for _, g := range pool.CandidateGroups {
|
||
if g.ID == uint(groupID) {
|
||
originGroup = g
|
||
break
|
||
}
|
||
}
|
||
if originGroup == nil {
|
||
log.Errorf("Logic error: Mapped GroupID %d not found in pool's candidate groups list", groupID)
|
||
return nil, nil, errors.New("cache inconsistency: mapped group not in candidate list")
|
||
}
|
||
if pool.PollingStrategy == models.StrategyWeighted {
|
||
go func() {
|
||
bgCtx, cancel := r.withTimeout(5 * time.Second)
|
||
defer cancel()
|
||
r.updateKeyUsageTimestampForPool(bgCtx, poolID, uint(keyID))
|
||
}()
|
||
}
|
||
return apiKey, originGroup, nil
|
||
}
|
||
|
||
// ensureBasePoolCacheExists 动态创建或验证 BasePool 的 Redis 缓存结构。
|
||
func (r *gormKeyRepository) ensureBasePoolCacheExists(ctx context.Context, pool *BasePool, poolID string) error {
|
||
heartbeatKey := fmt.Sprintf("basepool:%s:heartbeat", poolID)
|
||
emptyMarkerKey := fmt.Sprintf("basepool:empty:%s", poolID)
|
||
// 预检查，快速失败
|
||
if exists, _ := r.store.Exists(ctx, emptyMarkerKey); exists {
|
||
return gorm.ErrRecordNotFound
|
||
}
|
||
if exists, _ := r.store.Exists(ctx, heartbeatKey); exists {
|
||
return nil
|
||
}
|
||
// 获取分布式锁
|
||
lockKey := fmt.Sprintf("lock:basepool:%s", poolID)
|
||
if err := r.acquireLock(ctx, lockKey); err != nil {
|
||
return err // acquireLock 内部已记录日志并返回明确错误
|
||
}
|
||
defer r.releaseLock(context.Background(), lockKey)
|
||
// 双重检查锁定
|
||
if exists, _ := r.store.Exists(ctx, emptyMarkerKey); exists {
|
||
return gorm.ErrRecordNotFound
|
||
}
|
||
if exists, _ := r.store.Exists(ctx, heartbeatKey); exists {
|
||
return nil
|
||
}
|
||
// 在执行重度操作前，最后检查一次上下文是否已取消
|
||
select {
|
||
case <-ctx.Done():
|
||
return ctx.Err()
|
||
default:
|
||
}
|
||
r.logger.Infof("Building BasePool cache for pool_id '%s'", poolID)
|
||
// 手动聚合所有 Keys 并同时构建 key-to-group 映射
|
||
keyToGroupMap := make(map[string]any)
|
||
allKeyIDsSet := make(map[string]struct{})
|
||
for _, group := range pool.CandidateGroups {
|
||
groupKeySet := fmt.Sprintf(KeyGroup, group.ID)
|
||
groupKeyIDs, err := r.store.SMembers(ctx, groupKeySet)
|
||
if err != nil {
|
||
r.logger.WithError(err).Warnf("Failed to get members for group %d during pool build", group.ID)
|
||
continue
|
||
}
|
||
groupIDStr := strconv.FormatUint(uint64(group.ID), 10)
|
||
for _, keyID := range groupKeyIDs {
|
||
if _, exists := allKeyIDsSet[keyID]; !exists {
|
||
allKeyIDsSet[keyID] = struct{}{}
|
||
keyToGroupMap[keyID] = groupIDStr
|
||
}
|
||
}
|
||
}
|
||
// 处理空池情况
|
||
if len(allKeyIDsSet) == 0 {
|
||
emptyCacheTTL := time.Duration(r.config.Repository.BasePoolTTIMinutes) * time.Minute / 2
|
||
if emptyCacheTTL < time.Minute {
|
||
emptyCacheTTL = time.Minute
|
||
}
|
||
r.logger.Warnf("No active keys found for pool_id '%s', setting empty marker.", poolID)
|
||
if err := r.store.Set(ctx, emptyMarkerKey, []byte("1"), emptyCacheTTL); err != nil {
|
||
r.logger.WithError(err).Warnf("Failed to set empty marker for pool_id '%s'", poolID)
|
||
}
|
||
return gorm.ErrRecordNotFound
|
||
}
|
||
allActiveKeyIDs := make([]string, 0, len(allKeyIDsSet))
|
||
for keyID := range allKeyIDsSet {
|
||
allActiveKeyIDs = append(allActiveKeyIDs, keyID)
|
||
}
|
||
// 使用 Pipeline 原子化构建所有缓存结构
|
||
basePoolTTL := time.Duration(r.config.Repository.BasePoolTTLMinutes) * time.Minute
|
||
basePoolTTI := time.Duration(r.config.Repository.BasePoolTTIMinutes) * time.Minute
|
||
mainPoolKey := fmt.Sprintf(BasePoolRandomMain, poolID)
|
||
sequentialKey := fmt.Sprintf(BasePoolSequential, poolID)
|
||
cooldownKey := fmt.Sprintf(BasePoolRandomCooldown, poolID)
|
||
lruKey := fmt.Sprintf(BasePoolLRU, poolID)
|
||
keyToGroupMapKey := fmt.Sprintf("basepool:%s:key_to_group", poolID)
|
||
pipe := r.store.Pipeline(ctx)
|
||
pipe.Del(mainPoolKey, sequentialKey, cooldownKey, lruKey, emptyMarkerKey, keyToGroupMapKey)
|
||
pipe.SAdd(mainPoolKey, r.toInterfaceSlice(allActiveKeyIDs)...)
|
||
pipe.LPush(sequentialKey, r.toInterfaceSlice(allActiveKeyIDs)...)
|
||
if len(keyToGroupMap) > 0 {
|
||
pipe.HSet(keyToGroupMapKey, keyToGroupMap)
|
||
pipe.Expire(keyToGroupMapKey, basePoolTTL)
|
||
}
|
||
pipe.Expire(mainPoolKey, basePoolTTL)
|
||
pipe.Expire(sequentialKey, basePoolTTL)
|
||
pipe.Expire(cooldownKey, basePoolTTL)
|
||
pipe.Expire(lruKey, basePoolTTL)
|
||
pipe.Set(heartbeatKey, []byte("1"), basePoolTTI)
|
||
if err := pipe.Exec(); err != nil {
|
||
r.logger.WithError(err).Errorf("Failed to populate polling structures for pool_id '%s'", poolID)
|
||
cleanupCtx, cancel := r.withTimeout(5 * time.Second)
|
||
defer cancel()
|
||
r.store.Del(cleanupCtx, mainPoolKey, sequentialKey, cooldownKey, lruKey, heartbeatKey, emptyMarkerKey, keyToGroupMapKey)
|
||
return err
|
||
}
|
||
// 异步填充 LRU 缓存，并传入已构建好的映射
|
||
go r.populateBasePoolLRUCache(context.Background(), poolID, allActiveKeyIDs, keyToGroupMap)
|
||
return nil
|
||
}
|
||
|
||
// --- 辅助方法 ---
|
||
// acquireLock 封装了带重试和指数退避的分布式锁获取逻辑。
|
||
func (r *gormKeyRepository) acquireLock(ctx context.Context, lockKey string) error {
|
||
const (
|
||
lockTTL = 30 * time.Second
|
||
lockMaxRetries = 5
|
||
lockBaseBackoff = 50 * time.Millisecond
|
||
)
|
||
for i := 0; i < lockMaxRetries; i++ {
|
||
acquired, err := r.store.SetNX(ctx, lockKey, []byte("1"), lockTTL)
|
||
if err != nil {
|
||
r.logger.WithError(err).Error("Failed to attempt acquiring distributed lock")
|
||
return err
|
||
}
|
||
if acquired {
|
||
return nil
|
||
}
|
||
time.Sleep(lockBaseBackoff * (1 << i))
|
||
}
|
||
return fmt.Errorf("failed to acquire lock for key '%s' after %d retries", lockKey, lockMaxRetries)
|
||
}
|
||
|
||
// releaseLock 封装了分布式锁的释放逻辑。
|
||
func (r *gormKeyRepository) releaseLock(ctx context.Context, lockKey string) {
|
||
if err := r.store.Del(ctx, lockKey); err != nil {
|
||
r.logger.WithError(err).Errorf("Failed to release distributed lock for key '%s'", lockKey)
|
||
}
|
||
}
|
||
|
||
// withTimeout 是 context.WithTimeout 的一个简单包装，便于测试和模拟。
|
||
func (r *gormKeyRepository) withTimeout(duration time.Duration) (context.Context, context.CancelFunc) {
|
||
return context.WithTimeout(context.Background(), duration)
|
||
}
|
||
|
||
// refreshBasePoolHeartbeat 异步刷新心跳Key的TTI
|
||
func (r *gormKeyRepository) refreshBasePoolHeartbeat(ctx context.Context, poolID string) {
|
||
basePoolTTI := time.Duration(r.config.Repository.BasePoolTTIMinutes) * time.Minute
|
||
heartbeatKey := fmt.Sprintf("basepool:%s:heartbeat", poolID)
|
||
// 使用 EXPIRE 命令来刷新，如果Key不存在，它什么也不做，是安全的
|
||
if err := r.store.Expire(ctx, heartbeatKey, basePoolTTI); err != nil {
|
||
if ctx.Err() == nil { // 避免在context取消后打印不必要的错误
|
||
r.logger.WithError(err).Warnf("Failed to refresh heartbeat for pool_id '%s'", poolID)
|
||
}
|
||
}
|
||
}
|
||
|
||
// populateBasePoolLRUCache 异步填充 BasePool 的 LRU 缓存结构
|
||
func (r *gormKeyRepository) populateBasePoolLRUCache(
|
||
parentCtx context.Context,
|
||
currentPoolID string,
|
||
keys []string,
|
||
keyToGroupMap map[string]any,
|
||
) {
|
||
lruMembers := make(map[string]float64, len(keys))
|
||
for _, keyIDStr := range keys {
|
||
select {
|
||
case <-parentCtx.Done():
|
||
return
|
||
default:
|
||
}
|
||
groupIDStr, ok := keyToGroupMap[keyIDStr].(string)
|
||
if !ok {
|
||
continue
|
||
}
|
||
keyID, _ := strconv.ParseUint(keyIDStr, 10, 64)
|
||
groupID, _ := strconv.ParseUint(groupIDStr, 10, 64)
|
||
mappingKey := fmt.Sprintf(KeyMapping, groupID, keyID)
|
||
data, err := r.store.Get(parentCtx, mappingKey)
|
||
if err != nil {
|
||
continue
|
||
}
|
||
var mapping models.GroupAPIKeyMapping
|
||
if json.Unmarshal(data, &mapping) == nil {
|
||
var score float64
|
||
if mapping.LastUsedAt != nil {
|
||
score = float64(mapping.LastUsedAt.UnixMilli())
|
||
}
|
||
lruMembers[keyIDStr] = score
|
||
}
|
||
}
|
||
if len(lruMembers) > 0 {
|
||
lruKey := fmt.Sprintf(BasePoolLRU, currentPoolID)
|
||
if err := r.store.ZAdd(parentCtx, lruKey, lruMembers); err != nil {
|
||
if parentCtx.Err() == nil {
|
||
r.logger.WithError(err).Warnf("Failed to populate LRU cache for pool '%s'", currentPoolID)
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// updateKeyUsageTimestampForPool 更新 BasePool 的 LUR ZSET
|
||
func (r *gormKeyRepository) updateKeyUsageTimestampForPool(ctx context.Context, poolID string, keyID uint) {
|
||
lruKey := fmt.Sprintf(BasePoolLRU, poolID)
|
||
err := r.store.ZAdd(ctx, lruKey, map[string]float64{
|
||
strconv.FormatUint(uint64(keyID), 10): r.nowMilli(),
|
||
})
|
||
if err != nil {
|
||
r.logger.WithError(err).Warnf("Failed to update key usage for pool %s", poolID)
|
||
}
|
||
}
|
||
|
||
// generatePoolID 根据候选组ID列表生成一个稳定的、唯一的字符串ID
|
||
func (r *gormKeyRepository) generatePoolID(groups []*models.KeyGroup) string {
|
||
ids := make([]int, len(groups))
|
||
for i, g := range groups {
|
||
ids[i] = int(g.ID)
|
||
}
|
||
sort.Ints(ids)
|
||
h := sha1.New()
|
||
io.WriteString(h, fmt.Sprintf("%v", ids))
|
||
return fmt.Sprintf("%x", h.Sum(nil))
|
||
}
|
||
|
||
// toInterfaceSlice 类型转换辅助函数
|
||
func (r *gormKeyRepository) toInterfaceSlice(slice []string) []interface{} {
|
||
result := make([]interface{}, len(slice))
|
||
for i, v := range slice {
|
||
result[i] = v
|
||
}
|
||
return result
|
||
}
|
||
|
||
// nowMilli 返回当前的Unix毫秒时间戳，用于LRU/Weighted策略
|
||
func (r *gormKeyRepository) nowMilli() float64 {
|
||
return float64(time.Now().UnixMilli())
|
||
}
|
||
|
||
// getKeyDetailsFromCache 从缓存中获取Key和Mapping的JSON数据。
|
||
func (r *gormKeyRepository) getKeyDetailsFromCache(ctx context.Context, keyID, groupID uint) (*models.APIKey, *models.GroupAPIKeyMapping, error) {
|
||
apiKeyJSON, err := r.store.Get(ctx, fmt.Sprintf(KeyDetails, keyID))
|
||
if err != nil {
|
||
return nil, nil, fmt.Errorf("failed to get key details for key %d: %w", keyID, err)
|
||
}
|
||
var apiKey models.APIKey
|
||
if err := json.Unmarshal(apiKeyJSON, &apiKey); err != nil {
|
||
return nil, nil, fmt.Errorf("failed to unmarshal api key %d: %w", keyID, err)
|
||
}
|
||
|
||
mappingJSON, err := r.store.Get(ctx, fmt.Sprintf(KeyMapping, groupID, keyID))
|
||
if err != nil {
|
||
return nil, nil, fmt.Errorf("failed to get mapping details for key %d in group %d: %w", keyID, groupID, err)
|
||
}
|
||
var mapping models.GroupAPIKeyMapping
|
||
if err := json.Unmarshal(mappingJSON, &mapping); err != nil {
|
||
return nil, nil, fmt.Errorf("failed to unmarshal mapping for key %d in group %d: %w", keyID, groupID, err)
|
||
}
|
||
|
||
return &apiKey, &mapping, nil
|
||
}
|