fix(upload): 分片上传 - 先预留配额并原子提交容量 (#679)
Co-authored-by: ddcat666 <26564566+ddcat666@users.noreply.github.com>
This commit is contained in:
Catlaxy AI
141
backend/internal/repository/upload_reservations.go
Normal file
141
backend/internal/repository/upload_reservations.go
Normal file
@@ -0,0 +1,141 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"time"
|
||||
|
||||
"gorm.io/gorm"
|
||||
"infinite-canvas/backend/internal/model"
|
||||
)
|
||||
|
||||
var ErrUploadStorageLimit = errors.New("账号存储容量不足")
|
||||
var ErrUploadSessionLimit = errors.New("同时进行中的上传过多，请稍后重试")
|
||||
var ErrUploadReservationExpired = errors.New("上传会话不存在或已过期，请重新导入")
|
||||
|
||||
func lockUploadStorage(tx *gorm.DB, userID string) error {
|
||||
switch tx.Dialector.Name() {
|
||||
case "postgres":
|
||||
return tx.Exec("SELECT pg_advisory_xact_lock(hashtext(?)::bigint)", "upload-storage:"+userID).Error
|
||||
case "sqlite":
|
||||
return tx.Exec("UPDATE upload_reservations SET size = size WHERE 1 = 0").Error
|
||||
default:
|
||||
return gorm.ErrInvalidData
|
||||
}
|
||||
}
|
||||
|
||||
func (r *Repository) ReserveUploadSession(reservation *model.UploadReservation, dailyLimit, storedLimit int64, sessionLimit int) error {
|
||||
return r.db.Transaction(func(tx *gorm.DB) error {
|
||||
if err := lockUploadStorage(tx, reservation.UserID); err != nil {
|
||||
return err
|
||||
}
|
||||
repo := New(tx)
|
||||
now := time.Now()
|
||||
if !reservation.ExpiresAt.After(now) {
|
||||
return ErrUploadReservationExpired
|
||||
}
|
||||
reservation.Day = now.UTC().Format("2006-01-02")
|
||||
var expired []model.UploadReservation
|
||||
if err := tx.Where("user_id = ? AND expires_at <= ?", reservation.UserID, now).Find(&expired).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
for _, old := range expired {
|
||||
if err := repo.ReleaseDailyUpload(old.UserID, old.Day, old.Size); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Delete(&old).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
var active int64
|
||||
if err := tx.Model(&model.UploadReservation{}).Where("user_id = ?", reservation.UserID).Count(&active).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if active >= int64(sessionLimit) {
|
||||
return ErrUploadSessionLimit
|
||||
}
|
||||
stored, err := repo.UserStoredFileBytes(reservation.UserID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
var pending int64
|
||||
if err := tx.Model(&model.UploadReservation{}).Select("COALESCE(SUM(size),0)").Where("user_id = ?", reservation.UserID).Scan(&pending).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if reservation.Size <= 0 || stored >= storedLimit || pending >= storedLimit-stored || reservation.Size >= storedLimit-stored-pending {
|
||||
return ErrUploadStorageLimit
|
||||
}
|
||||
if err := repo.ReserveDailyUpload(reservation.UserID, reservation.Day, reservation.Size, dailyLimit); err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Create(reservation).Error
|
||||
})
|
||||
}
|
||||
|
||||
func (r *Repository) ReleaseUploadSession(userID, id string) error {
|
||||
return r.db.Transaction(func(tx *gorm.DB) error {
|
||||
if err := lockUploadStorage(tx, userID); err != nil {
|
||||
return err
|
||||
}
|
||||
var reservation model.UploadReservation
|
||||
if err := tx.Where("id = ? AND user_id = ?", id, userID).First(&reservation).Error; err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil
|
||||
}
|
||||
return err
|
||||
}
|
||||
if err := New(tx).ReleaseDailyUpload(userID, reservation.Day, reservation.Size); err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Delete(&reservation).Error
|
||||
})
|
||||
}
|
||||
|
||||
// SaveResourceWithinStorageLimit publishes ready bytes and consumes the session
|
||||
// reservation under the same database lock used by admission on every instance.
|
||||
func (r *Repository) SaveResourceWithinStorageLimit(resource *model.Resource, limit, dailyLimit int64, reservationID string) error {
|
||||
return r.db.Transaction(func(tx *gorm.DB) error {
|
||||
if err := lockUploadStorage(tx, resource.UserID); err != nil {
|
||||
return err
|
||||
}
|
||||
now := time.Now()
|
||||
if reservationID != "" {
|
||||
var reservation model.UploadReservation
|
||||
if err := tx.Where("id = ? AND user_id = ? AND size = ? AND expires_at > ?", reservationID, resource.UserID, resource.Size, now).First(&reservation).Error; err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return ErrUploadReservationExpired
|
||||
}
|
||||
return err
|
||||
}
|
||||
day := now.UTC().Format("2006-01-02")
|
||||
if reservation.Day != day {
|
||||
repo := New(tx)
|
||||
if err := repo.ReserveDailyUpload(resource.UserID, day, reservation.Size, dailyLimit); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := repo.ReleaseDailyUpload(resource.UserID, reservation.Day, reservation.Size); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
// Save inside the transaction so physical-object deduplication uses the
|
||||
// existing usage query; rollback leaves the resource non-ready on denial.
|
||||
if err := tx.Save(resource).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
stored, err := New(tx).UserStoredFileBytes(resource.UserID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
var pending int64
|
||||
if err := tx.Model(&model.UploadReservation{}).Select("COALESCE(SUM(size),0)").Where("user_id = ? AND id <> ? AND expires_at > ?", resource.UserID, reservationID, now).Scan(&pending).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
if stored >= limit || pending >= limit-stored {
|
||||
return ErrUploadStorageLimit
|
||||
}
|
||||
if reservationID != "" {
|
||||
return tx.Delete(&model.UploadReservation{}, "id = ? AND user_id = ?", reservationID, resource.UserID).Error
|
||||
}
|
||||
return nil
|
||||
})
|
||||
}
|
||||
195
backend/internal/repository/upload_reservations_test.go
Normal file
195
backend/internal/repository/upload_reservations_test.go
Normal file
@@ -0,0 +1,195 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"path/filepath"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"gorm.io/driver/sqlite"
|
||||
"gorm.io/gorm"
|
||||
"infinite-canvas/backend/internal/model"
|
||||
)
|
||||
|
||||
func uploadReservationTestDB(t *testing.T) (*gorm.DB, *Repository, *Repository) {
|
||||
t.Helper()
|
||||
dsn := filepath.Join(t.TempDir(), "uploads.db") + "?_busy_timeout=5000&_journal_mode=WAL"
|
||||
db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.AutoMigrate(&model.UploadReservation{}, &model.UserDailyUploadUsage{}, &model.Resource{}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
peer, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, conn := range []*gorm.DB{db, peer} {
|
||||
sqlDB, _ := conn.DB()
|
||||
t.Cleanup(func() { _ = sqlDB.Close() })
|
||||
}
|
||||
return db, New(db), New(peer)
|
||||
}
|
||||
|
||||
func TestUploadReservationConcurrentAdmission(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
size, daily, stored int64
|
||||
sessions, attempts, want int
|
||||
denied error
|
||||
}{
|
||||
{"sessions", 1, 1000, 1000, 32, 40, 32, ErrUploadSessionLimit},
|
||||
{"storage", 60, 1000, 100, 32, 2, 1, ErrUploadStorageLimit},
|
||||
{"daily", 60, 100, 1000, 32, 2, 1, ErrDailyUploadLimitExceeded},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
_, first, second := uploadReservationTestDB(t)
|
||||
start := make(chan struct{})
|
||||
results := make(chan error, tc.attempts)
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < tc.attempts; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
<-start
|
||||
repo := []*Repository{first, second}[i%2]
|
||||
results <- repo.ReserveUploadSession(&model.UploadReservation{ID: fmt.Sprint(i), UserID: "user", Size: tc.size, ExpiresAt: time.Now().Add(time.Hour)}, tc.daily, tc.stored, tc.sessions)
|
||||
}()
|
||||
}
|
||||
close(start)
|
||||
wg.Wait()
|
||||
close(results)
|
||||
accepted := 0
|
||||
for err := range results {
|
||||
if err == nil {
|
||||
accepted++
|
||||
} else if !errors.Is(err, tc.denied) {
|
||||
t.Fatalf("unexpected denial: %v", err)
|
||||
}
|
||||
}
|
||||
if accepted != tc.want {
|
||||
t.Fatalf("accepted=%d want=%d", accepted, tc.want)
|
||||
}
|
||||
if used, err := first.DailyUploadBytes("user", time.Now().UTC().Format("2006-01-02")); err != nil || used != int64(tc.want)*tc.size {
|
||||
t.Fatalf("reserved bytes=%d err=%v", used, err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestUploadReservationAtomicCapacityCommit(t *testing.T) {
|
||||
db, first, second := uploadReservationTestDB(t)
|
||||
start := make(chan struct{})
|
||||
results := make(chan error, 2)
|
||||
var wg sync.WaitGroup
|
||||
for i, repo := range []*Repository{first, second} {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
<-start
|
||||
results <- repo.SaveResourceWithinStorageLimit(&model.Resource{ID: fmt.Sprint(i), UserID: "user", ObjectKey: fmt.Sprint(i), Status: model.ResourceStatusReady, Size: 60}, 100, 1000, "")
|
||||
}()
|
||||
}
|
||||
close(start)
|
||||
wg.Wait()
|
||||
close(results)
|
||||
accepted := 0
|
||||
for err := range results {
|
||||
if err == nil {
|
||||
accepted++
|
||||
} else if !errors.Is(err, ErrUploadStorageLimit) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
if accepted != 1 {
|
||||
t.Fatalf("committed=%d want=1", accepted)
|
||||
}
|
||||
var count int64
|
||||
if err := db.Model(&model.Resource{}).Count(&count).Error; err != nil || count != 1 {
|
||||
t.Fatalf("resource rows=%d err=%v", count, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUploadReservationCommitAndExpiry(t *testing.T) {
|
||||
db, r, peer := uploadReservationTestDB(t)
|
||||
old := model.UploadReservation{ID: "old", UserID: "user", Size: 60, Day: "2026-10-04", ExpiresAt: time.Now().Add(-time.Hour)}
|
||||
if err := db.Create(&old).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := r.ReserveDailyUpload("user", old.Day, 60, 1000); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
lease := model.UploadReservation{ID: "new", UserID: "user", Size: 60, Day: "2026-10-05", ExpiresAt: time.Now().Add(time.Hour)}
|
||||
if err := peer.ReserveUploadSession(&lease, 1000, 100, 1); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if used, _ := r.DailyUploadBytes("user", old.Day); used != 0 {
|
||||
t.Fatalf("expired previous-day bytes=%d", used)
|
||||
}
|
||||
if err := r.ReleaseUploadSession("other", lease.ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
other := model.Resource{ID: "other", UserID: "user", Size: 40, Status: model.ResourceStatusReady, ObjectKey: "other"}
|
||||
if err := r.SaveResourceWithinStorageLimit(&other, 100, 1000, ""); !errors.Is(err, ErrUploadStorageLimit) {
|
||||
t.Fatalf("ordinary upload ignored reserved capacity: %v", err)
|
||||
}
|
||||
resource := model.Resource{ID: "resource", UserID: "user", Size: 59, Status: model.ResourceStatusReady, ObjectKey: "resource"}
|
||||
if err := r.SaveResourceWithinStorageLimit(&resource, 100, 1000, lease.ID); !errors.Is(err, ErrUploadReservationExpired) {
|
||||
t.Fatalf("mismatched lease: %v", err)
|
||||
}
|
||||
resource.Size = 60
|
||||
if err := r.SaveResourceWithinStorageLimit(&resource, 100, 1000, lease.ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := peer.ReleaseUploadSession("user", lease.ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if used, _ := r.DailyUploadBytes("user", lease.Day); used != 60 {
|
||||
t.Fatalf("committed daily bytes=%d", used)
|
||||
}
|
||||
if err := r.SaveResourceWithinStorageLimit(&resource, 100, 1000, lease.ID); !errors.Is(err, ErrUploadReservationExpired) {
|
||||
t.Fatalf("lease reused: %v", err)
|
||||
}
|
||||
resource.ID = "expired"
|
||||
if err := r.SaveResourceWithinStorageLimit(&resource, 100, 1000, old.ID); !errors.Is(err, ErrUploadReservationExpired) {
|
||||
t.Fatalf("expired lease accepted: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUploadReservationCommitAcrossUTCDay(t *testing.T) {
|
||||
db, r, _ := uploadReservationTestDB(t)
|
||||
today := time.Now().UTC().Format("2006-01-02")
|
||||
yesterday := time.Now().UTC().Add(-24 * time.Hour).Format("2006-01-02")
|
||||
lease := model.UploadReservation{ID: "midnight", UserID: "user", Size: 15, Day: yesterday, ExpiresAt: time.Now().Add(time.Hour)}
|
||||
if err := db.Create(&lease).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := r.ReserveDailyUpload("user", yesterday, 15, 100); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := r.ReserveDailyUpload("user", today, 90, 100); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
resource := model.Resource{ID: "midnight", UserID: "user", Size: 15, ObjectKey: "midnight", Status: model.ResourceStatusReady}
|
||||
if err := r.SaveResourceWithinStorageLimit(&resource, 1000, 100, lease.ID); !errors.Is(err, ErrDailyUploadLimitExceeded) {
|
||||
t.Fatalf("new-day quota bypassed: %v", err)
|
||||
}
|
||||
if used, _ := r.DailyUploadBytes("user", yesterday); used != 15 {
|
||||
t.Fatalf("failed commit refunded reservation: %d", used)
|
||||
}
|
||||
if err := r.ReleaseDailyUpload("user", today, 90); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := r.SaveResourceWithinStorageLimit(&resource, 1000, 100, lease.ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if used, _ := r.DailyUploadBytes("user", today); used != 15 {
|
||||
t.Fatalf("new-day usage=%d", used)
|
||||
}
|
||||
if used, _ := r.DailyUploadBytes("user", yesterday); used != 0 {
|
||||
t.Fatalf("old-day reservation=%d", used)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user