package services import ( "log" "time" "d:\项目\dobaochet\backend\config" "d:\项目\dobaochet\backend\db" "d:\项目\dobaochet\backend\models" ) // QuotaService 配额管理服务 type QuotaService struct { cfg *config.Config } // NewQuotaService 创建配额管理服务实例 func NewQuotaService(cfg *config.Config) *QuotaService { return &QuotaService{cfg: cfg} } // StartQuotaResetJob 启动配额重置定时任务 func (s *QuotaService) StartQuotaResetJob() { // 立即执行一次重置检查 s.ResetExpiredQuotas() // 每天检查一次 ticker := time.NewTicker(24 * time.Hour) go func() { for range ticker.C { s.ResetExpiredQuotas() } }() } // ResetExpiredQuotas 重置过期的配额 func (s *QuotaService) ResetExpiredQuotas() { log.Println("开始检查并重置过期配额") // 查询所有需要重置配额的用户 var users []models.User result := db.GetDB().Preload("Quota").Where("quota.reset_at <= ?", time.Now()).Find(&users) if result.Error != nil { log.Printf("查询需要重置配额的用户失败: %v", result.Error) return } if len(users) == 0 { log.Println("没有需要重置配额的用户") return } log.Printf("发现 %d 个用户需要重置配额", len(users)) // 批量重置配额 tx := db.GetDB().Begin() defer func() { if r := recover(); r != nil { tx.Rollback() log.Printf("重置配额时发生 panic: %v", r) } }() for _, user := range users { // 重置配额 user.Quota.UsedTokens = 0 user.Quota.ResetAt = time.Now().Add(s.cfg.Quota.ResetInterval) if result := tx.Save(&user.Quota); result.Error != nil { tx.Rollback() log.Printf("重置用户 %s 配额失败: %v", user.Username, result.Error) return } log.Printf("已重置用户 %s 的配额", user.Username) } if err := tx.Commit().Error; err != nil { log.Printf("提交配额重置事务失败: %v", err) return } log.Printf("成功重置 %d 个用户的配额", len(users)) } // CheckAndUpdateQuota 检查并更新用户配额 func (s *QuotaService) CheckAndUpdateQuota(userID uint, usedTokens int64) (bool, error) { // 查询用户配额 var quota models.Quota result := db.GetDB().Where("user_id = ?", userID).First("a) if result.Error != nil { return false, result.Error } // 检查配额是否过期 if quota.ResetAt.Before(time.Now()) { // 重置配额 quota.UsedTokens = 0 quota.ResetAt = time.Now().Add(s.cfg.Quota.ResetInterval) } // 检查配额是否足够 if quota.UsedTokens+usedTokens > quota.TotalTokens { return false, nil } // 更新配额 quota.UsedTokens += usedTokens if result := db.GetDB().Save("a); result.Error != nil { return false, result.Error } return true, nil } // GetUserQuota 获取用户配额信息 func (s *QuotaService) GetUserQuota(userID uint) (*models.Quota, error) { var quota models.Quota result := db.GetDB().Where("user_id = ?", userID).First("a) if result.Error != nil { return nil, result.Error } // 检查并重置过期配额 if quota.ResetAt.Before(time.Now()) { quota.UsedTokens = 0 quota.ResetAt = time.Now().Add(s.cfg.Quota.ResetInterval) db.GetDB().Save("a) } return "a, nil } // UpdateUserQuota 更新用户配额 func (s *QuotaService) UpdateUserQuota(userID uint, totalTokens, tokenLimit int64) error { var quota models.Quota result := db.GetDB().Where("user_id = ?", userID).First("a) if result.Error != nil { return result.Error } // 更新配额 quota.TotalTokens = totalTokens quota.TokenLimit = tokenLimit return db.GetDB().Save("a).Error } // ResetUserQuota 重置单个用户配额 func (s *QuotaService) ResetUserQuota(userID uint) error { var quota models.Quota result := db.GetDB().Where("user_id = ?", userID).First("a) if result.Error != nil { return result.Error } // 重置配额 quota.UsedTokens = 0 quota.ResetAt = time.Now().Add(s.cfg.Quota.ResetInterval) return db.GetDB().Save("a).Error }