78 lines
3.0 KiB
Go
78 lines
3.0 KiB
Go
// Filename: internal/repository/key_writer.go
|
|
package repository
|
|
|
|
import (
|
|
"fmt"
|
|
"gemini-balancer/internal/errors"
|
|
"gemini-balancer/internal/models"
|
|
"strconv"
|
|
"time"
|
|
)
|
|
|
|
func (r *gormKeyRepository) UpdateKeyUsageTimestamp(groupID, keyID uint) {
|
|
lruKey := fmt.Sprintf(KeyGroupLRU, groupID)
|
|
timestamp := float64(time.Now().UnixMilli())
|
|
|
|
members := map[string]float64{
|
|
strconv.FormatUint(uint64(keyID), 10): timestamp,
|
|
}
|
|
|
|
if err := r.store.ZAdd(lruKey, members); err != nil {
|
|
r.logger.WithError(err).Warnf("Failed to update usage timestamp for key %d in group %d", keyID, groupID)
|
|
}
|
|
}
|
|
|
|
func (r *gormKeyRepository) SyncKeyStatusInPollingCaches(groupID, keyID uint, newStatus models.APIKeyStatus) {
|
|
r.logger.Infof("SYNC: Directly updating polling caches for G:%d K:%d -> %s", groupID, keyID, newStatus)
|
|
r.updatePollingCachesLogic(groupID, keyID, newStatus)
|
|
}
|
|
|
|
func (r *gormKeyRepository) HandleCacheUpdateEvent(groupID, keyID uint, newStatus models.APIKeyStatus) {
|
|
r.logger.Infof("EVENT: Updating polling caches for G:%d K:%d -> %s from an event", groupID, keyID, newStatus)
|
|
r.updatePollingCachesLogic(groupID, keyID, newStatus)
|
|
}
|
|
|
|
func (r *gormKeyRepository) updatePollingCachesLogic(groupID, keyID uint, newStatus models.APIKeyStatus) {
|
|
keyIDStr := strconv.FormatUint(uint64(keyID), 10)
|
|
sequentialKey := fmt.Sprintf(KeyGroupSequential, groupID)
|
|
lruKey := fmt.Sprintf(KeyGroupLRU, groupID)
|
|
mainPoolKey := fmt.Sprintf(KeyGroupRandomMain, groupID)
|
|
cooldownPoolKey := fmt.Sprintf(KeyGroupRandomCooldown, groupID)
|
|
|
|
_ = r.store.LRem(sequentialKey, 0, keyIDStr)
|
|
_ = r.store.ZRem(lruKey, keyIDStr)
|
|
_ = r.store.SRem(mainPoolKey, keyIDStr)
|
|
_ = r.store.SRem(cooldownPoolKey, keyIDStr)
|
|
|
|
if newStatus == models.StatusActive {
|
|
if err := r.store.LPush(sequentialKey, keyIDStr); err != nil {
|
|
r.logger.WithError(err).Warnf("Failed to add key %d to sequential list for group %d", keyID, groupID)
|
|
}
|
|
members := map[string]float64{keyIDStr: 0}
|
|
if err := r.store.ZAdd(lruKey, members); err != nil {
|
|
r.logger.WithError(err).Warnf("Failed to add key %d to LRU zset for group %d", keyID, groupID)
|
|
}
|
|
if err := r.store.SAdd(mainPoolKey, keyIDStr); err != nil {
|
|
r.logger.WithError(err).Warnf("Failed to add key %d to random main pool for group %d", keyID, groupID)
|
|
}
|
|
}
|
|
}
|
|
|
|
// UpdateKeyStatusAfterRequest is the new central hub for handling feedback.
|
|
func (r *gormKeyRepository) UpdateKeyStatusAfterRequest(group *models.KeyGroup, key *models.APIKey, success bool, apiErr *errors.APIError) {
|
|
if success {
|
|
if group.PollingStrategy == models.StrategyWeighted {
|
|
go r.UpdateKeyUsageTimestamp(group.ID, key.ID)
|
|
}
|
|
return
|
|
}
|
|
if apiErr == nil {
|
|
r.logger.Warnf("Request failed for KeyID %d in GroupID %d but no specific API error was provided.", key.ID, group.ID)
|
|
return
|
|
}
|
|
r.logger.Warnf("Request failed for KeyID %d in GroupID %d with error: %s. Temporarily removing from active polling caches.", key.ID, group.ID, apiErr.Message)
|
|
|
|
// This call is correct. It uses the synchronous, direct method.
|
|
r.SyncKeyStatusInPollingCaches(group.ID, key.ID, models.StatusCooldown)
|
|
}
|