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

package app

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

	"yingce/backend/internal/model"
	"yingce/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))
}

// 分片和普通上传使用相同的单文件上限。
func (s *Service) reserveChunkedUploadQuota(userID string, size int64) (string, error) {
	return s.reserveUserUploadQuota(userID, size)
}

func (s *Service) ReserveChunkUploadSession(userID, id string, size int64, expires time.Time, maxSessions int) error {
	policy, err := s.RuntimePolicy()
	if err != nil {
		return err
	}
	if size <= 0 || size >= megabytes(policy.Resource.ResourceUploadMB) {
		return BadAuthRequest(fmt.Sprintf("单个上传文件必须小于 %dMB", policy.Resource.ResourceUploadMB))
	}
	if !expires.After(time.Now()) || maxSessions <= 0 {
		return BadAuthRequest("上传会话参数无效")
	}
	err = s.repo.ReserveUploadSession(&model.UploadReservation{ID: id, UserID: userID, Size: size, Day: time.Now().UTC().Format("2006-01-02"), ExpiresAt: expires}, megabytes(policy.Resource.DailyUploadMB), gigabytes(policy.Resource.StoredFileGB), maxSessions)
	return uploadReservationError(err)
}

func (s *Service) ReleaseChunkUploadSession(userID, id string) error {
	return s.repo.ReleaseUploadSession(userID, id)
}

func uploadReservationError(err error) error {
	if errors.Is(err, repository.ErrUploadSessionLimit) {
		return NewAppError(429, err.Error())
	}
	if errors.Is(err, repository.ErrUploadStorageLimit) || errors.Is(err, repository.ErrDailyUploadLimitExceeded) {
		return QuotaExceeded(err.Error())
	}
	if errors.Is(err, repository.ErrUploadReservationExpired) {
		return BadAuthRequest(err.Error())
	}
	return err
}

func (s *Service) saveResourceWithinStorageLimit(resource *model.Resource, reservationID string) error {
	policy, err := s.RuntimePolicy()
	if err != nil {
		return err
	}
	return uploadReservationError(s.repo.SaveResourceWithinStorageLimit(resource, gigabytes(policy.Resource.StoredFileGB), megabytes(policy.Resource.DailyUploadMB), reservationID))
}

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)
}