gpt-load/internal/services/key_cron_service.go

package services

import (
	"context"
	"gpt-load/internal/config"
	"gpt-load/internal/models"
	"sync"
	"time"

	"github.com/sirupsen/logrus"
	"gorm.io/gorm"
)

// KeyCronService is responsible for periodically validating all API keys.
type KeyCronService struct {
	DB              *gorm.DB
	Validator       *KeyValidatorService
	SettingsManager *config.SystemSettingsManager
	stopChan        chan struct{}
	wg              sync.WaitGroup
}

// NewKeyCronService creates a new KeyCronService.
func NewKeyCronService(db *gorm.DB, validator *KeyValidatorService, settingsManager *config.SystemSettingsManager) *KeyCronService {
	return &KeyCronService{
		DB:              db,
		Validator:       validator,
		SettingsManager: settingsManager,
		stopChan:        make(chan struct{}),
	}
}

// Start begins the cron job.
func (s *KeyCronService) Start() {
	logrus.Info("Starting KeyCronService...")
	s.wg.Add(1)
	go s.run()
}

// Stop stops the cron job.
func (s *KeyCronService) Stop() {
	logrus.Info("Stopping KeyCronService...")
	close(s.stopChan)
	s.wg.Wait()
	logrus.Info("KeyCronService stopped.")
}

func (s *KeyCronService) run() {
	defer s.wg.Done()
	ctx := context.Background()

	ticker := time.NewTicker(5 * time.Minute)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			s.validateAllGroups(ctx)
		case <-s.stopChan:
			return
		}
	}
}

func (s *KeyCronService) validateAllGroups(ctx context.Context) {
	logrus.Info("KeyCronService: Starting validation cycle for all groups.")
	var groups []models.Group
	if err := s.DB.Find(&groups).Error; err != nil {
		logrus.Errorf("KeyCronService: Failed to get groups: %v", err)
		return
	}

	validationStartTime := time.Now()
	var wg sync.WaitGroup
	for _, group := range groups {
		groupCopy := group
		wg.Add(1)
		go func(g models.Group) {
			defer wg.Done()
			defer func() {
				if r := recover(); r != nil {
					logrus.Errorf("KeyCronService: Panic recovered in group validation for %s: %v", g.Name, r)
				}
			}()

			effectiveSettings := s.SettingsManager.GetEffectiveConfig(g.Config)
			interval := time.Duration(effectiveSettings.KeyValidationIntervalMinutes) * time.Minute

			if g.LastValidatedAt == nil || validationStartTime.Sub(*g.LastValidatedAt) > interval {
				s.validateGroup(ctx, &g, validationStartTime)
			}
		}(groupCopy)
	}
	wg.Wait()
	logrus.Info("KeyCronService: Validation cycle finished.")
}

func (s *KeyCronService) validateGroup(ctx context.Context, group *models.Group, validationStartTime time.Time) {
	var keys []models.APIKey
	if err := s.DB.Where("group_id = ?", group.ID).Find(&keys).Error; err != nil {
		logrus.Errorf("KeyCronService: Failed to get keys for group %s: %v", group.Name, err)
		return
	}

	if len(keys) == 0 {
		if err := s.DB.Model(group).Update("last_validated_at", validationStartTime).Error; err != nil {
			logrus.Errorf("KeyCronService: Failed to update last_validated_at for empty group %s: %v", group.Name, err)
		}
		return
	}

	logrus.Infof("KeyCronService: Validating %d keys for group %s", len(keys), group.Name)

	jobs := make(chan models.APIKey, len(keys))
	results := make(chan models.APIKey, len(keys))

	concurrency := s.SettingsManager.GetInt("key_validation_concurrency", 10)
	if concurrency <= 0 {
		concurrency = 10
	}

	var workerWg sync.WaitGroup
	for range concurrency {
		workerWg.Add(1)
		go s.worker(ctx, &workerWg, group, jobs, results)
	}

	for _, key := range keys {
		jobs <- key
	}
	close(jobs)

	workerWg.Wait()
	close(results)

	var keysToUpdate []models.APIKey
	for key := range results {
		keysToUpdate = append(keysToUpdate, key)
	}

	if len(keysToUpdate) > 0 {
		s.batchUpdateKeyStatus(keysToUpdate)
	}

	if err := s.DB.Model(group).Update("last_validated_at", validationStartTime).Error; err != nil {
		logrus.Errorf("KeyCronService: Failed to update last_validated_at for group %s: %v", group.Name, err)
	}

	logrus.Infof("KeyCronService: Finished validating group %s. %d keys had their status changed.", group.Name, len(keysToUpdate))
}

func (s *KeyCronService) worker(ctx context.Context, wg *sync.WaitGroup, group *models.Group, jobs <-chan models.APIKey, results chan<- models.APIKey) {
	defer wg.Done()
	for key := range jobs {
		isValid, validationErr := s.Validator.ValidateSingleKey(ctx, &key, group)

		var newStatus string
		var newErrorReason string
		statusChanged := false

		if validationErr != nil {
			newStatus = "inactive"
			newErrorReason = validationErr.Error()
		} else {
			if isValid {
				newStatus = "active"
				newErrorReason = ""
			} else {
				newStatus = "inactive"
				newErrorReason = "Validation returned false without a specific error."
			}
		}

		if key.Status != newStatus || key.ErrorReason != newErrorReason {
			statusChanged = true
		}

		if statusChanged {
			key.Status = newStatus
			key.ErrorReason = newErrorReason
			results <- key
		}
	}
}

func (s *KeyCronService) batchUpdateKeyStatus(keys []models.APIKey) {
	if len(keys) == 0 {
		return
	}
	logrus.Infof("KeyCronService: Batch updating status/reason for %d keys.", len(keys))

	err := s.DB.Transaction(func(tx *gorm.DB) error {
		for _, key := range keys {
			updates := map[string]any{
				"status":       key.Status,
				"error_reason": key.ErrorReason,
			}
			if err := tx.Model(&models.APIKey{}).Where("id = ?", key.ID).Updates(updates).Error; err != nil {
				// Log the error for this specific key but continue the transaction
				logrus.Errorf("KeyCronService: Failed to update key ID %d: %v", key.ID, err)
			}
		}
		return nil // Commit the transaction even if some updates failed
	})

	if err != nil {
		// This error is for the transaction itself, not individual updates
		logrus.Errorf("KeyCronService: Transaction failed during batch update of key statuses: %v", err)
	}
}