ttstd-infinite-canvas/backend/internal/app/upload_quota.go

package app

import (
	"errors"
	"fmt"
	"log"
	"time"

	"infinite-canvas/backend/internal/repository"
)

// 上传额度在写文件或 OSS 前原子预留，避免并发请求同时通过日限额检查。
func (s *Service) reserveUserUploadQuota(userID string, size int64) (string, error) {
	policy, err := s.RuntimePolicy()
	if err != nil {
		return "", err
	}
	return s.reserveUserStoredFileQuota(userID, size, megabytes(policy.Resource.ResourceUploadMB), megabytes(policy.Resource.DailyUploadMB), gigabytes(policy.Resource.StoredFileGB), fmt.Sprintf("单个上传文件必须小于 %dMB", policy.Resource.ResourceUploadMB))
}

// reserveChunkedUploadQuota 用于分片上传完成时预留额度：单文件上限对分片会话不适用（每片已独立校验），
// 仅受“今日上传”与“账号存储总量”约束；singleFileLimit 传 size+1 使单文件上限永不命中。
func (s *Service) reserveChunkedUploadQuota(userID string, size int64) (string, error) {
	policy, err := s.RuntimePolicy()
	if err != nil {
		return "", err
	}
	return s.reserveUserStoredFileQuota(userID, size, size+1, megabytes(policy.Resource.DailyUploadMB), gigabytes(policy.Resource.StoredFileGB), "")
}

func (s *Service) reserveGeneratedResourceQuota(userID string, size int64) (string, error) {
	policy, err := s.RuntimePolicy()
	if err != nil {
		return "", err
	}
	return s.reserveUserStoredFileQuota(userID, size, megabytes(policy.Resource.GeneratedFileMB)+1, megabytes(policy.Resource.DailyUploadMB), gigabytes(policy.Resource.StoredFileGB), fmt.Sprintf("单个生成文件不能超过 %dMB", policy.Resource.GeneratedFileMB))
}

// 失败资源的记录已经计入账号存储用量；重试只重新预留当日上传额度，避免重复计算存储容量。
func (s *Service) reserveRetryUploadQuota(userID string, size int64) (string, error) {
	policy, err := s.RuntimePolicy()
	if err != nil {
		return "", err
	}
	if size <= 0 {
		return "", BadAuthRequest("上传文件不能为空")
	}
	if size >= megabytes(policy.Resource.ResourceUploadMB) {
		return "", BadAuthRequest(fmt.Sprintf("单个上传文件必须小于 %dMB", policy.Resource.ResourceUploadMB))
	}
	day := time.Now().UTC().Format("2006-01-02")
	s.storageMu.Lock()
	defer s.storageMu.Unlock()
	if err := s.repo.ReserveDailyUpload(userID, day, size, megabytes(policy.Resource.DailyUploadMB)); err != nil {
		if errors.Is(err, repository.ErrDailyUploadLimitExceeded) {
			return "", QuotaExceeded(fmt.Sprintf("每个账号 UTC 自然日上传总量必须小于 %s", formatStorageLimit(megabytes(policy.Resource.DailyUploadMB))))
		}
		return "", err
	}
	return day, nil
}

func (s *Service) reserveUserStoredFileQuota(userID string, size int64, exclusiveSingleFileLimit int64, dailyLimit int64, storedLimit int64, singleFileMessage string) (string, error) {
	if size <= 0 {
		return "", BadAuthRequest("上传文件不能为空")
	}
	if size >= exclusiveSingleFileLimit {
		return "", BadAuthRequest(singleFileMessage)
	}
	day := time.Now().UTC().Format("2006-01-02")
	s.storageMu.Lock()
	defer s.storageMu.Unlock()
	storedBytes, err := s.repo.UserStoredFileBytes(userID)
	if err != nil {
		return "", err
	}
	if s.pendingStorage == nil {
		s.pendingStorage = map[string]int64{}
	}
	if storedBytes+s.pendingStorage[userID]+size >= storedLimit {
		return "", QuotaExceeded(fmt.Sprintf("账号资源和会话附件已达到 %s 上限，请联系管理员清理历史文件", formatStorageLimit(storedLimit)))
	}
	s.pendingStorage[userID] += size
	if err := s.repo.ReserveDailyUpload(userID, day, size, dailyLimit); err != nil {
		s.decreasePendingStorage(userID, size)
		if errors.Is(err, repository.ErrDailyUploadLimitExceeded) {
			return "", QuotaExceeded(fmt.Sprintf("每个账号 UTC 自然日上传总量必须小于 %s", formatStorageLimit(dailyLimit)))
		}
		return "", err
	}
	return day, nil
}

func formatStorageLimit(value int64) string {
	if value%(1<<30) == 0 {
		return fmt.Sprintf("%dGB", value>>30)
	}
	return fmt.Sprintf("%dMB", value>>20)
}

func (s *Service) releaseUserUploadQuota(userID string, day string, size int64) {
	if day == "" || size <= 0 {
		return
	}
	s.storageMu.Lock()
	defer s.storageMu.Unlock()
	s.decreasePendingStorage(userID, size)
	if err := s.repo.ReleaseDailyUpload(userID, day, size); err != nil {
		log.Printf("release upload quota failed: user=%s day=%s size=%d error=%v", userID, day, size, err)
	}
}

func (s *Service) releaseRetryUploadQuota(userID string, day string, size int64) {
	if day == "" || size <= 0 {
		return
	}
	s.storageMu.Lock()
	defer s.storageMu.Unlock()
	if err := s.repo.ReleaseDailyUpload(userID, day, size); err != nil {
		log.Printf("release retry upload quota failed: user=%s day=%s size=%d error=%v", userID, day, size, err)
	}
}

func (s *Service) commitUserUploadQuota(userID string, size int64) {
	if size <= 0 {
		return
	}
	s.storageMu.Lock()
	defer s.storageMu.Unlock()
	s.decreasePendingStorage(userID, size)
}

func (s *Service) decreasePendingStorage(userID string, size int64) {
	remaining := s.pendingStorage[userID] - size
	if remaining > 0 {
		s.pendingStorage[userID] = remaining
		return
	}
	delete(s.pendingStorage, userID)
}