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
@@ -0,0 +1,255 @@
+package main
+
+import (
+	"crypto/sha256"
+	"encoding/hex"
+	"errors"
+	"fmt"
+	"os"
+	"path/filepath"
+	"time"
+
+	"gorm.io/gorm"
+)
+
+const assetLease = 15 * time.Minute
+
+var (
+	errTaskRunning     = errors.New("task is already running")
+	errTaskAlreadyDone = errors.New("task is already completed")
+)
+
+func (a *App) claimTask(id uint) (*Task, error) {
+	result := a.db.Model(&Task{}).Where("id = ? AND status = ?", id, "pending").
+		Updates(map[string]any{"status": "resolving", "error": "", "progress": 0})
+	if result.Error != nil {
+		return nil, result.Error
+	}
+	if result.RowsAffected != 1 {
+		var t Task
+		if err := a.db.First(&t, id).Error; err != nil {
+			return nil, err
+		}
+		if t.Status == "upload_pending" {
+			return nil, fmt.Errorf("task %d is waiting for upload", id)
+		}
+		if t.Status == "done" {
+			return nil, errTaskAlreadyDone
+		}
+		return nil, errTaskRunning
+	}
+	var t Task
+	if err := a.db.First(&t, id).Error; err != nil {
+		return nil, err
+	}
+	return &t, nil
+}
+
+func (a *App) claimUpload(t *Task) bool {
+	result := a.db.Model(&Task{}).
+		Where("id = ? AND (status = ? OR (status = ? AND (remote_path = '' OR remote_path IS NULL)))",
+			t.ID, "upload_pending", "done").
+		Updates(map[string]any{"status": "uploading", "error": ""})
+	if result.Error != nil || result.RowsAffected != 1 {
+		return false
+	}
+	t.Status, t.Error = "uploading", ""
+	return true
+}
+
+func (a *App) taskReservationStatus(taskID, reservationID uint) (string, error) {
+	if reservationID == 0 {
+		return "", nil
+	}
+	var r CreditReservation
+	if err := a.db.First(&r, reservationID).Error; err != nil {
+		return "", err
+	}
+	if r.TaskID != taskID {
+		return "", fmt.Errorf("credit reservation does not belong to task")
+	}
+	return r.Status, nil
+}
+
+func (a *App) repairCommittedTask(t *Task) error {
+	status, err := a.taskReservationStatus(t.ID, t.CreditReservationID)
+	if err != nil {
+		return err
+	}
+	if status == "committed" {
+		if t.Status == "error" {
+			t.Status, t.Error = "upload_pending", "下载已完成，正在恢复 WebDAV 上传"
+			return a.db.Model(t).Updates(map[string]any{
+				"status": t.Status, "error": t.Error,
+			}).Error
+		}
+		return nil
+	}
+	return nil
+}
+
+func (a *App) markAssetHeartbeat(assetID, taskID uint) {
+	if assetID == 0 || taskID == 0 {
+		return
+	}
+	now := time.Now()
+	lease := now.Add(assetLease)
+	_ = a.db.Model(&MediaAsset{}).
+		Where("id = ? AND owner_task_id = ? AND status = ?", assetID, taskID, "downloading").
+		Updates(map[string]any{"lease_until": lease, "updated_at": now}).Error
+}
+
+func (a *App) assetLeaseExpired(asset *MediaAsset) bool {
+	return asset.Status == "downloading" && !asset.LeaseUntil.IsZero() && time.Now().After(asset.LeaseUntil)
+}
+func (a *App) claimAsset(t *Task) (*MediaAsset, bool, error) {
+	for i := 0; i < 180; i++ {
+		asset, err := a.ensureAsset(t)
+		if err != nil {
+			return nil, false, err
+		}
+		if asset.Status == "done" && asset.LocalPath != "" {
+			if fi, statErr := os.Stat(asset.LocalPath); statErr == nil && fi.Size() > 0 {
+				t.LocalPath, t.Format, t.ActualQuality = asset.LocalPath, asset.Format, asset.ActualQuality
+				t.TotalBytes, t.DownloadedBytes, t.Progress, t.Status = fi.Size(), fi.Size(), 100, "done"
+				if err := a.db.Save(t).Error; err != nil {
+					return nil, false, err
+				}
+				return asset, true, nil
+			}
+		}
+
+		claimed := false
+		now := time.Now()
+		lease := now.Add(assetLease)
+		err = a.transactionRetry(func(tx *gorm.DB) error {
+			var current MediaAsset
+			if err := tx.First(&current, asset.ID).Error; err != nil {
+				return err
+			}
+			if current.Status == "downloading" && current.OwnerTaskID != t.ID {
+				if !current.LeaseUntil.IsZero() && now.Before(current.LeaseUntil) {
+					return nil
+				}
+				if current.OwnerTaskID != 0 {
+					_ = tx.Model(&Task{}).Where("id = ? AND status IN ?", current.OwnerTaskID, []string{"resolving", "downloading"}).Updates(map[string]any{"status": "pending", "error": "共享资源下载锁超时，已自动恢复", "progress": 0})
+				}
+			}
+			wasEvicted := current.Status == "evicted"
+			wasStolen := current.Status == "downloading" && current.OwnerTaskID != 0 && current.OwnerTaskID != t.ID && (current.LeaseUntil.IsZero() || !now.Before(current.LeaseUntil))
+			current.Status = "downloading"
+			current.OwnerTaskID = t.ID
+			current.LeaseUntil = lease
+			current.Generation++
+			if current.LocalPath == "" || wasEvicted || wasStolen {
+				current.LocalPath = assetCachePath(a.cfg.DownloadDir, current.AssetKey, t.Format, current.Generation)
+			}
+			if err := tx.Save(&current).Error; err != nil {
+				return err
+			}
+			*asset = current
+			claimed = true
+			return nil
+		})
+		if err != nil {
+			return nil, false, err
+		}
+		if claimed {
+			t.LocalPath = asset.LocalPath
+			return asset, false, nil
+		}
+		time.Sleep(500 * time.Millisecond)
+	}
+	return nil, false, fmt.Errorf("timed out waiting for shared media asset")
+}
+func assetCachePath(root, key, format string, generation uint64) string {
+	sum := sha256.Sum256([]byte(key))
+	return filepath.Join(root, ".assets", hex.EncodeToString(sum[:]), fmt.Sprintf("asset-%d.%s", generation, normalizedFormat(format)))
+}
+
+func (a *App) assertAssetOwner(t *Task) bool {
+	if t.AssetID == 0 {
+		return true
+	}
+	var asset MediaAsset
+	if a.db.First(&asset, t.AssetID).Error != nil {
+		return false
+	}
+	return asset.OwnerTaskID == t.ID && asset.Status == "downloading"
+}
+
+func (a *App) finishAsset(t *Task) bool {
+	if t.AssetID == 0 {
+		return true
+	}
+	result := a.db.Model(&MediaAsset{}).
+		Where("id = ? AND owner_task_id = ? AND status = ?", t.AssetID, t.ID, "downloading").
+		Updates(map[string]any{
+			"status": "done", "owner_task_id": 0, "lease_until": nil,
+			"local_path": t.LocalPath, "format": t.Format,
+			"actual_quality": t.ActualQuality, "size": t.DownloadedBytes,
+		})
+	return result.Error == nil && result.RowsAffected == 1
+}
+
+func (a *App) invalidateAsset(t *Task) {
+	if t.AssetID == 0 {
+		return
+	}
+	_ = a.db.Model(&MediaAsset{}).
+		Where("id = ? AND owner_task_id = ?", t.AssetID, t.ID).
+		Updates(map[string]any{"status": "error", "owner_task_id": 0, "lease_until": nil})
+}
+
+func (a *App) prepareRetry(t *Task) error {
+	status, err := a.taskReservationStatus(t.ID, t.CreditReservationID)
+	if err != nil {
+		return err
+	}
+	switch status {
+	case "committed":
+		t.Status = "upload_pending"
+		t.Error = "下载已完成，恢复 WebDAV 上传"
+		return a.db.Model(t).Updates(map[string]any{
+			"status": "upload_pending", "error": t.Error,
+		}).Error
+	case "reserved":
+		t.Status, t.Error = "pending", ""
+		return a.db.Model(t).Updates(map[string]any{
+			"status": "pending", "error": "", "progress": 0,
+		}).Error
+	case "refunded", "":
+		r, reserveErr := a.reserveCredit(t.UserID, t.ID, 1)
+		if reserveErr != nil {
+			return reserveErr
+		}
+		t.CreditReservationID = r.ID
+		t.Status, t.Error, t.Progress, t.DownloadedBytes = "pending", "", 0, 0
+		return a.db.Model(t).Updates(map[string]any{
+			"status": "pending", "error": "", "progress": 0,
+			"downloaded_bytes": 0, "credit_reservation_id": r.ID,
+		}).Error
+	default:
+		return errReservationState
+	}
+}
+
+func (a *App) enqueueTask(id uint) {
+	select {
+	case a.queue <- id:
+	default:
+		go func() { a.queue <- id }()
+	}
+}
+
+func (a *App) transactionRetry(fn func(*gorm.DB) error) error {
+	var err error
+	for attempt := 0; attempt < 8; attempt++ {
+		err = a.db.Transaction(fn)
+		if !isSQLiteBusy(err) {
+			return err
+		}
+		time.Sleep(time.Duration(25*(1<<attempt)) * time.Millisecond)
+	}
+	return err
+}