上传初始版，功能未经任何测试，仅为基本框架，目前为完全不可用状态

backend/services/quota_service.go

@@ -0,0 +1,162 @@
+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(&quota)
+	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(&quota); 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(&quota)
+	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(&quota)
+	}
+
+	return &quota, nil
+}
+
+// UpdateUserQuota 更新用户配额
+func (s *QuotaService) UpdateUserQuota(userID uint, totalTokens, tokenLimit int64) error {
+	var quota models.Quota
+	result := db.GetDB().Where("user_id = ?", userID).First(&quota)
+	if result.Error != nil {
+		return result.Error
+	}
+
+	// 更新配额
+	quota.TotalTokens = totalTokens
+	quota.TokenLimit = tokenLimit
+
+	return db.GetDB().Save(&quota).Error
+}
+
+// ResetUserQuota 重置单个用户配额
+func (s *QuotaService) ResetUserQuota(userID uint) error {
+	var quota models.Quota
+	result := db.GetDB().Where("user_id = ?", userID).First(&quota)
+	if result.Error != nil {
+		return result.Error
+	}
+
+	// 重置配额
+	quota.UsedTokens = 0
+	quota.ResetAt = time.Now().Add(s.cfg.Quota.ResetInterval)
+
+	return db.GetDB().Save(&quota).Error
+}