104 lines
3.6 KiB
Go
104 lines
3.6 KiB
Go
package service
|
|
|
|
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))
|
|
}
|
|
|
|
func (s *Service) reserveSessionUploadQuota(userID string, size int64) (string, error) {
|
|
policy, err := s.RuntimePolicy()
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
return s.reserveUserStoredFileQuota(userID, size, megabytes(policy.Resource.SessionUploadMB)+1, megabytes(policy.Resource.DailyUploadMB), gigabytes(policy.Resource.StoredFileGB), fmt.Sprintf("会话文件不能超过 %dMB", policy.Resource.SessionUploadMB))
|
|
}
|
|
|
|
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) 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 "", BadAuthRequest(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 "", BadAuthRequest(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) 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)
|
|
}
|