feat: harden multi-user download concurrency

2026-09-18 09:06:59 +08:00
parent 0735cb2e09
commit b116303abe
11 changed files with 1003 additions and 202 deletions
@@ -6,6 +6,7 @@ import (
 	"strings"

 	"github.com/gin-gonic/gin"
+	"gorm.io/gorm"
 )

 func (a *App) search(c *gin.Context) {
@@ -60,46 +61,67 @@ func (a *App) addTask(c *gin.Context) {
 		in.Quality = a.cfg.DefaultQuality
 	}
 	user := c.MustGet("user").(*User)
+
 	var existing Task
-	err := a.db.Where("user_id = ? AND source = ? AND r_id = ? AND quality = ?", user.ID, in.Source, in.RID, in.Quality).Order("id desc").First(&existing).Error
+	err := a.db.Where("user_id = ? AND source = ? AND r_id = ? AND quality = ?",
+		user.ID, in.Source, in.RID, in.Quality).Order("id desc").First(&existing).Error
 	if err == nil {
-		if existing.Status != "error" {
+		existing.Title, existing.Artist, existing.Album, existing.Pic = in.Title, in.Artist, in.Album, in.Pic
+		switch existing.Status {
+		case "pending", "resolving", "downloading", "uploading", "upload_pending":
 			c.JSON(200, existing)
 			return
+		case "done":
+			c.JSON(200, existing)
+			return
+		case "error":
+			if err := a.prepareRetry(&existing); err != nil {
+				if err == errInsufficientCredits {
+					c.JSON(402, gin.H{"error": err.Error()})
+				} else {
+					c.JSON(409, gin.H{"error": err.Error()})
+				}
+				return
+			}
+			a.enqueueTask(existing.ID)
+			c.JSON(202, existing)
+			return
 		}
-		r, reserveErr := a.reserveCredit(user.ID, existing.ID, 1)
+	}
+
+	t := Task{
+		UserID: user.ID, Source: in.Source, RID: in.RID,
+		Title: in.Title, Artist: in.Artist, Album: in.Album,
+		Pic: in.Pic, Quality: in.Quality, Status: "pending",
+	}
+	var reservation *CreditReservation
+	err = a.transactionRetry(func(tx *gorm.DB) error {
+		if err := tx.Create(&t).Error; err != nil {
+			return err
+		}
+		var reserveErr error
+		reservation, reserveErr = reserveCreditTx(tx, user.ID, t.ID, 1)
 		if reserveErr != nil {
-			c.JSON(402, gin.H{"error": reserveErr.Error()})
-			return
+			return reserveErr
 		}
-		a.db.Model(&existing).Updates(map[string]any{"title": in.Title, "artist": in.Artist, "album": in.Album, "pic": in.Pic, "status": "pending", "error": "", "credit_reservation_id": r.ID})
-		existing.CreditReservationID = r.ID
-		a.queue <- existing.ID
-		c.JSON(202, existing)
-		return
-	}
-	t := Task{UserID: user.ID, Source: in.Source, RID: in.RID, Title: in.Title, Artist: in.Artist, Album: in.Album, Pic: in.Pic, Quality: in.Quality, Status: "pending"}
-	if err := a.db.Create(&t).Error; err != nil {
-		if a.db.Where("user_id = ? AND source = ? AND r_id = ? AND quality = ?", user.ID, in.Source, in.RID, in.Quality).Order("id desc").First(&existing).Error == nil {
+		t.CreditReservationID = reservation.ID
+		return tx.Model(&t).Update("credit_reservation_id", reservation.ID).Error
+	})
+	if err != nil {
+		if a.db.Where("user_id = ? AND source = ? AND r_id = ? AND quality = ?",
+			user.ID, in.Source, in.RID, in.Quality).First(&existing).Error == nil {
 			c.JSON(200, existing)
 			return
 		}
+		if err == errInsufficientCredits {
+			c.JSON(402, gin.H{"error": err.Error()})
+			return
+		}
 		c.JSON(500, gin.H{"error": err.Error()})
 		return
 	}
-	r, reserveErr := a.reserveCredit(user.ID, t.ID, 1)
-	if reserveErr != nil {
-		_ = a.db.Delete(&t).Error
-		c.JSON(402, gin.H{"error": reserveErr.Error()})
-		return
-	}
-	t.CreditReservationID = r.ID
-	if err := a.db.Model(&t).Update("credit_reservation_id", r.ID).Error; err != nil {
-		_ = a.refundCredit(&t)
-		c.JSON(500, gin.H{"error": err.Error()})
-		return
-	}
-	a.queue <- t.ID
+
+	a.enqueueTask(t.ID)
 	c.JSON(202, t)
 }

@@ -119,28 +141,35 @@ func (a *App) retryTask(c *gin.Context) {
 		c.Status(400)
 		return
 	}
-	var t Task
 	user := c.MustGet("user").(*User)
+	var t Task
 	if a.db.Where("id = ? AND user_id = ?", id, user.ID).First(&t).Error != nil {
 		c.Status(404)
 		return
 	}
-	if t.Status != "error" {
-		c.JSON(409, gin.H{"error": "only failed tasks can be retried"})
+	switch t.Status {
+	case "done":
+		c.JSON(409, gin.H{"error": "task is already completed"})
+		return
+	case "pending", "resolving", "downloading", "uploading":
+		c.JSON(409, gin.H{"error": "task is already running"})
+		return
+	case "upload_pending", "error":
+		if t.Status == "error" {
+			if err := a.prepareRetry(&t); err != nil {
+				if err == errInsufficientCredits {
+					c.JSON(402, gin.H{"error": err.Error()})
+				} else {
+					c.JSON(409, gin.H{"error": err.Error()})
+				}
+				return
+			}
+		}
+	default:
+		c.JSON(409, gin.H{"error": "task cannot be retried from current state"})
 		return
 	}
-	r, reserveErr := a.reserveCredit(user.ID, t.ID, 1)
-	if reserveErr != nil {
-		c.JSON(402, gin.H{"error": reserveErr.Error()})
-		return
-	}
-	if err := a.db.Model(&t).Updates(map[string]any{"status": "pending", "error": "", "progress": 0, "downloaded_bytes": 0, "credit_reservation_id": r.ID}).Error; err != nil {
-		_ = a.refundCredit(&Task{ID: t.ID, UserID: user.ID, CreditReservationID: r.ID})
-		c.JSON(500, gin.H{"error": err.Error()})
-		return
-	}
-	t.Status, t.Error, t.Progress, t.DownloadedBytes, t.CreditReservationID = "pending", "", 0, 0, r.ID
-	a.queue <- t.ID
+	a.enqueueTask(t.ID)
 	c.JSON(202, t)
 }