feat: 优化事务

internal/keypool/provider.go

@@ -82,94 +82,112 @@ func (p *KeyProvider) UpdateStatus(keyID uint, groupID uint, isSuccess bool) {
 		activeKeysListKey := fmt.Sprintf("group:%d:active_keys", groupID)
 
 		if isSuccess {
-			p.handleSuccess(keyID, keyHashKey, activeKeysListKey)
+			if err := p.handleSuccess(keyID, keyHashKey, activeKeysListKey); err != nil {
+				logrus.WithFields(logrus.Fields{"keyID": keyID, "error": err}).Error("Failed to handle key success")
+			}
 		} else {
-			p.handleFailure(keyID, keyHashKey, activeKeysListKey)
+			if err := p.handleFailure(keyID, keyHashKey, activeKeysListKey); err != nil {
+				logrus.WithFields(logrus.Fields{"keyID": keyID, "error": err}).Error("Failed to handle key failure")
+			}
 		}
 	}()
 }
 
-func (p *KeyProvider) handleSuccess(keyID uint, keyHashKey, activeKeysListKey string) {
+func (p *KeyProvider) handleSuccess(keyID uint, keyHashKey, activeKeysListKey string) error {
 	keyDetails, err := p.store.HGetAll(keyHashKey)
 	if err != nil {
-		logrus.WithFields(logrus.Fields{"keyID": keyID, "error": err}).Error("Failed to get key details on success")
+		return fmt.Errorf("failed to get key details from store: %w", err)
-		return
 	}
 
 	failureCount, _ := strconv.ParseInt(keyDetails["failure_count"], 10, 64)
 	isActive := keyDetails["status"] == models.KeyStatusActive
 
 	if failureCount == 0 && isActive {
-		return
+		return nil
 	}
 
-	if err := p.store.HSet(keyHashKey, map[string]any{"failure_count": 0}); err != nil {
+	return p.db.Transaction(func(tx *gorm.DB) error {
-		logrus.WithFields(logrus.Fields{"keyID": keyID, "error": err}).Error("Failed to reset failure count in store, aborting DB update.")
+		var key models.APIKey
-		return
+		if err := tx.Set("gorm:query_option", "FOR UPDATE").First(&key, keyID).Error; err != nil {
-	}
+			return fmt.Errorf("failed to lock key %d for update: %w", keyID, err)
-
-	updates := map[string]any{"failure_count": 0}
-
-	if !isActive {
-		logrus.WithField("keyID", keyID).Info("Key has recovered and is being restored to active pool.")
-		if err := p.store.HSet(keyHashKey, map[string]any{"status": models.KeyStatusActive}); err != nil {
-			logrus.WithFields(logrus.Fields{"keyID": keyID, "error": err}).Error("Failed to update key status to active in store, aborting DB update.")
-			return
 		}
-		// To prevent duplicates, first remove any existing instance of the key from the list.
-		// This makes the recovery operation idempotent.
-		if err := p.store.LRem(activeKeysListKey, 0, keyID); err != nil {
-			logrus.WithFields(logrus.Fields{"keyID": keyID, "error": err}).Error("Failed to LRem key before LPush on recovery, aborting DB update.")
-			return
-		}
-		if err := p.store.LPush(activeKeysListKey, keyID); err != nil {
-			logrus.WithFields(logrus.Fields{"keyID": keyID, "error": err}).Error("Failed to LPush key back to active list, aborting DB update.")
-			return
-		}
-		updates["status"] = models.KeyStatusActive
-	}
 
-	if err := p.db.Model(&models.APIKey{}).Where("id = ?", keyID).Updates(updates).Error; err != nil {
+		updates := map[string]any{"failure_count": 0}
-		logrus.WithFields(logrus.Fields{"keyID": keyID, "error": err}).Error("Failed to update key status in DB on success")
+		if !isActive {
-	}
+			updates["status"] = models.KeyStatusActive
+		}
+
+		if err := tx.Model(&key).Updates(updates).Error; err != nil {
+			return fmt.Errorf("failed to update key in DB: %w", err)
+		}
+
+		if err := p.store.HSet(keyHashKey, updates); err != nil {
+			return fmt.Errorf("failed to update key details in store: %w", err)
+		}
+
+		if !isActive {
+			logrus.WithField("keyID", keyID).Info("Key has recovered and is being restored to active pool.")
+			if err := p.store.LRem(activeKeysListKey, 0, keyID); err != nil {
+				return fmt.Errorf("failed to LRem key before LPush on recovery: %w", err)
+			}
+			if err := p.store.LPush(activeKeysListKey, keyID); err != nil {
+				return fmt.Errorf("failed to LPush key back to active list: %w", err)
+			}
+		}
+
+		return nil
+	})
 }
 
-func (p *KeyProvider) handleFailure(keyID uint, keyHashKey, activeKeysListKey string) {
+func (p *KeyProvider) handleFailure(keyID uint, keyHashKey, activeKeysListKey string) error {
 	keyDetails, err := p.store.HGetAll(keyHashKey)
 	if err != nil {
-		logrus.WithFields(logrus.Fields{"keyID": keyID, "error": err}).Error("Failed to get key details on failure")
+		return fmt.Errorf("failed to get key details from store: %w", err)
-		return
-	}
-	if keyDetails["status"] == models.KeyStatusInvalid {
-		return
 	}
 
-	newFailureCount, err := p.store.HIncrBy(keyHashKey, "failure_count", 1)
+	failureCount, _ := strconv.ParseInt(keyDetails["failure_count"], 10, 64)
-	if err != nil {
+
-		logrus.WithFields(logrus.Fields{"keyID": keyID, "error": err}).Error("Failed to increment failure count")
+	if keyDetails["status"] == models.KeyStatusInvalid {
-		return
+		return nil
 	}
 
 	settings := p.settingsManager.GetSettings()
 	blacklistThreshold := settings.BlacklistThreshold
-	updates := map[string]any{"failure_count": newFailureCount}
 
-	if newFailureCount >= int64(blacklistThreshold) {
+	return p.db.Transaction(func(tx *gorm.DB) error {
-		logrus.WithFields(logrus.Fields{"keyID": keyID, "threshold": blacklistThreshold}).Warn("Key has reached blacklist threshold, disabling.")
+		var key models.APIKey
-		if err := p.store.LRem(activeKeysListKey, 0, keyID); err != nil {
+		if err := tx.Set("gorm:query_option", "FOR UPDATE").First(&key, keyID).Error; err != nil {
-			logrus.WithFields(logrus.Fields{"keyID": keyID, "error": err}).Error("Failed to LRem key from active list, aborting DB update.")
+			return fmt.Errorf("failed to lock key %d for update: %w", keyID, err)
-			return
 		}
-		if err := p.store.HSet(keyHashKey, map[string]any{"status": models.KeyStatusInvalid}); err != nil {
-			logrus.WithFields(logrus.Fields{"keyID": keyID, "error": err}).Error("Failed to update key status to invalid in store, aborting DB update.")
-			return
-		}
-		updates["status"] = models.KeyStatusInvalid
-	}
 
-	if err := p.db.Model(&models.APIKey{}).Where("id = ?", keyID).Updates(updates).Error; err != nil {
+		newFailureCount := failureCount + 1
-		logrus.WithFields(logrus.Fields{"keyID": keyID, "error": err}).Error("Failed to update key stats in DB on failure")
+
-	}
+		updates := map[string]any{"failure_count": newFailureCount}
+		shouldBlacklist := newFailureCount >= int64(blacklistThreshold)
+		if shouldBlacklist {
+			updates["status"] = models.KeyStatusInvalid
+		}
+
+		if err := tx.Model(&key).Updates(updates).Error; err != nil {
+			return fmt.Errorf("failed to update key stats in DB: %w", err)
+		}
+
+		if _, err := p.store.HIncrBy(keyHashKey, "failure_count", 1); err != nil {
+			return fmt.Errorf("failed to increment failure count in store: %w", err)
+		}
+
+		if shouldBlacklist {
+			logrus.WithFields(logrus.Fields{"keyID": keyID, "threshold": blacklistThreshold}).Warn("Key has reached blacklist threshold, disabling.")
+			if err := p.store.LRem(activeKeysListKey, 0, keyID); err != nil {
+				return fmt.Errorf("failed to LRem key from active list: %w", err)
+			}
+			if err := p.store.HSet(keyHashKey, map[string]any{"status": models.KeyStatusInvalid}); err != nil {
+				return fmt.Errorf("failed to update key status to invalid in store: %w", err)
+			}
+		}
+
+		return nil
+	})
 }
 
 // LoadKeysFromDB 从数据库加载所有分组和密钥，并填充到 Store 中。
@@ -278,7 +296,8 @@ func (p *KeyProvider) AddKeys(groupID uint, keys []models.APIKey) error {
 
 		for _, key := range keys {
 			if err := p.addKeyToStore(&key); err != nil {
-				logrus.WithFields(logrus.Fields{"keyID": key.ID, "error": err}).Error("Failed to add key to store after DB creation")
+				logrus.WithFields(logrus.Fields{"keyID": key.ID, "error": err}).Error("Failed to add key to store after DB creation, rolling back transaction")
+				return err
 			}
 		}
 		return nil
@@ -315,7 +334,8 @@ func (p *KeyProvider) RemoveKeys(groupID uint, keyValues []string) (int64, error
 
 		for _, key := range keysToDelete {
 			if err := p.removeKeyFromStore(key.ID, key.GroupID); err != nil {
-				logrus.WithFields(logrus.Fields{"keyID": key.ID, "error": err}).Error("Failed to remove key from store after DB deletion")
+				logrus.WithFields(logrus.Fields{"keyID": key.ID, "error": err}).Error("Failed to remove key from store after DB deletion, rolling back transaction")
+				return err
 			}
 		}
 
@@ -353,7 +373,8 @@ func (p *KeyProvider) RestoreKeys(groupID uint) (int64, error) {
 			key.Status = models.KeyStatusActive
 			key.FailureCount = 0
 			if err := p.addKeyToStore(&key); err != nil {
-				logrus.WithFields(logrus.Fields{"keyID": key.ID, "error": err}).Error("Failed to restore key in store after DB update")
+				logrus.WithFields(logrus.Fields{"keyID": key.ID, "error": err}).Error("Failed to restore key in store after DB update, rolling back transaction")
+				return err
 			}
 		}
 		return nil
@@ -432,9 +453,9 @@ func (p *KeyProvider) RemoveInvalidKeys(groupID uint) (int64, error) {
 		removedCount = result.RowsAffected
 
 		for _, key := range invalidKeys {
-			keyHashKey := fmt.Sprintf("key:%d", key.ID)
+			if err := p.removeKeyFromStore(key.ID, key.GroupID); err != nil {
-			if err := p.store.Delete(keyHashKey); err != nil {
+				logrus.WithFields(logrus.Fields{"keyID": key.ID, "error": err}).Error("Failed to remove invalid key from store after DB deletion, rolling back transaction")
-				logrus.WithFields(logrus.Fields{"keyID": key.ID, "error": err}).Error("Failed to remove invalid key HASH from store after DB deletion")
+				return err
 			}
 		}
 		return nil