fix(upload): 分片上传 - 先预留配额并原子提交容量 (#679)

Co-authored-by: ddcat666 <26564566+ddcat666@users.noreply.github.com>
2026-10-06 09:42:48 +08:00
1 parent 398f8e36f5
commit 2cc53fa25c
@@ -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
+	})
+}