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
@@ -7,6 +7,7 @@ import (
 	"os"
 	"path/filepath"
 	"strings"
+	"time"
 )

 func (a *App) startWorkers() {
@@ -15,21 +16,36 @@ func (a *App) startWorkers() {
 		go func() {
 			defer a.wg.Done()
 			for id := range a.queue {
-				a.download(id)
+				a.processTask(id)
 			}
 		}()
 	}
 }

-func (a *App) download(id uint) {
+func (a *App) processTask(id uint) {
 	var t Task
-	if a.db.First(&t, id).Error != nil || t.Status != "pending" {
+	if a.db.First(&t, id).Error != nil {
 		return
 	}
-	a.db.Model(&t).Updates(map[string]any{"status": "resolving", "error": "", "progress": 0})
+	if t.Status == "upload_pending" {
+		a.upload(&t)
+		return
+	}
+	if t.Status != "pending" {
+		return
+	}
+	a.download(id)
+}
+
+func (a *App) download(id uint) {
+	t, err := a.claimTask(id)
+	if err != nil {
+		return
+	}
+
 	info, err := a.am.Resolve(t.Source, t.RID, t.Quality)
 	if err != nil {
-		a.fail(&t, err)
+		a.fail(t, err)
 		return
 	}

@@ -41,17 +57,30 @@ func (a *App) download(id uint) {
 		ext = "mp3"
 	}
 	t.Format, t.ResolvedQuality, t.ActualQuality = ext, info.Level, ""
-	asset, reused, err := a.acquireAsset(&t)
+	if err := a.db.Model(t).Updates(map[string]any{
+		"format": t.Format, "resolved_quality": t.ResolvedQuality,
+	}).Error; err != nil {
+		a.fail(t, err)
+		return
+	}
+
+	asset, reused, err := a.acquireAsset(t)
 	if err != nil {
-		a.fail(&t, err)
+		a.fail(t, err)
 		return
 	}
 	if reused {
-		a.finishExisting(&t, asset.Size)
+		a.finishExisting(t, asset.Size)
 		return
 	}
-	a.db.Model(&t).Updates(map[string]any{"format": t.Format, "local_path": t.LocalPath})
-	a.downloadFile(&t, info)
+	if err := a.db.Model(t).Updates(map[string]any{
+		"local_path": t.LocalPath, "status": "downloading",
+	}).Error; err != nil {
+		a.invalidateAsset(t)
+		a.fail(t, err)
+		return
+	}
+	a.downloadFile(t, info)
 }

 func (a *App) finishExisting(t *Task, size int64) {
@@ -70,8 +99,9 @@ func (a *App) finishExisting(t *Task, size int64) {
 		t.ActualQuality = actualQuality(t.Quality, t.ResolvedQuality, actual)
 	}
 	t.TotalBytes, t.DownloadedBytes, t.Progress, t.Status = size, size, 100, "done"
-	a.db.Save(t)
-	// Existing local media is reused and therefore does not consume a credit.
+	if err := a.db.Save(t).Error; err != nil {
+		return
+	}
 	if err := a.refundCredit(t); err != nil {
 		a.fail(t, err)
 		return
@@ -80,7 +110,6 @@ func (a *App) finishExisting(t *Task, size int64) {
 }

 func (a *App) downloadFile(t *Task, info playInfo) {
-	a.db.Model(t).Updates(map[string]any{"status": "downloading", "format": t.Format, "local_path": t.LocalPath})
 	req, err := http.NewRequest("GET", info.URL, nil)
 	if err != nil {
 		a.fail(t, err)
@@ -99,27 +128,62 @@ func (a *App) downloadFile(t *Task, info playInfo) {
 		a.fail(t, fmt.Errorf("media http %d", resp.StatusCode))
 		return
 	}
-	t.Status, t.TotalBytes, t.DownloadedBytes, t.Progress = "downloading", resp.ContentLength, 0, 0
-	a.db.Save(t)

-	tmp := t.LocalPath + ".part"
+	t.Status = "downloading"
+	t.TotalBytes = resp.ContentLength
+	t.DownloadedBytes = 0
+	t.Progress = 0
+	_ = a.db.Model(t).Updates(map[string]any{
+		"status": t.Status, "total_bytes": t.TotalBytes,
+		"downloaded_bytes": 0, "progress": 0,
+	}).Error
+
+	if err := os.MkdirAll(filepath.Dir(t.LocalPath), 0755); err != nil {
+		a.fail(t, err)
+		return
+	}
+	tmp := fmt.Sprintf("%s.task-%d.part", t.LocalPath, t.ID)
+	_ = os.Remove(tmp)
 	f, err := os.Create(tmp)
 	if err != nil {
 		a.fail(t, err)
 		return
 	}
-	n, err := a.copyWithProgress(f, resp.Body, t)
+
+	stop := make(chan struct{})
+	ticker := time.NewTicker(30 * time.Second)
+	go func() {
+		for {
+			select {
+			case <-ticker.C:
+				a.markAssetHeartbeat(t.AssetID, t.ID)
+			case <-stop:
+				return
+			}
+		}
+	}()
+	n, copyErr := a.copyWithProgress(f, resp.Body, t)
 	closeErr := f.Close()
-	if err != nil {
-		a.fail(t, err)
+	ticker.Stop()
+	close(stop)
+
+	if copyErr != nil {
+		a.fail(t, copyErr)
+		_ = os.Remove(tmp)
 		return
 	}
 	if closeErr != nil {
 		a.fail(t, closeErr)
+		_ = os.Remove(tmp)
 		return
 	}
 	if t.TotalBytes > 0 && n != t.TotalBytes {
 		a.fail(t, fmt.Errorf("incomplete download: got %d of %d bytes", n, t.TotalBytes))
+		_ = os.Remove(tmp)
+		return
+	}
+	if !a.assertAssetOwner(t) {
+		_ = os.Remove(tmp)
 		return
 	}

@@ -131,26 +195,40 @@ func (a *App) downloadFile(t *Task, info playInfo) {
 	}
 	actual = normalizedFormat(actual)
 	finalPath := strings.TrimSuffix(t.LocalPath, filepath.Ext(t.LocalPath)) + "." + actual
-	if err = os.Rename(tmp, finalPath); err != nil {
+	if err := os.Rename(tmp, finalPath); err != nil {
 		a.fail(t, err)
 		_ = os.Remove(tmp)
 		return
 	}
-	t.Format, t.LocalPath, t.ActualQuality, t.DownloadedBytes = actual, finalPath, actualQuality(t.Quality, t.ResolvedQuality, actual), n
-	if t.TotalBytes > 0 {
-		t.Progress = 100
+	t.LocalPath = finalPath
+	t.Format = actual
+	t.ActualQuality = actualQuality(t.Quality, t.ResolvedQuality, actual)
+	t.DownloadedBytes = n
+	t.Progress = 100
+	if err := a.db.Model(t).Updates(map[string]any{
+		"format": t.Format, "local_path": t.LocalPath,
+		"actual_quality":   t.ActualQuality,
+		"downloaded_bytes": n, "progress": 100,
+	}).Error; err != nil {
+		a.invalidateAsset(t)
+		return
+	}
+	if !a.finishAsset(t) {
+		a.fail(t, fmt.Errorf("shared asset ownership lost before completion"))
+		return
 	}
-	a.db.Save(t)
-	a.markAssetDone(t)
-	// The credit pays for the successful media download, not the later WebDAV sync.
-	// Commit before upload so a storage outage cannot turn a successful download into a free retry.
 	if err := a.commitCredit(t); err != nil {
 		a.fail(t, err)
 		return
 	}
+	t.Status, t.Error = "upload_pending", ""
+	if err := a.db.Model(t).Updates(map[string]any{
+		"status": "upload_pending", "error": "",
+	}).Error; err != nil {
+		return
+	}
 	a.upload(t)
 }
-
 func (a *App) copyWithProgress(dst *os.File, src io.Reader, t *Task) (int64, error) {
 	var n, last int64
 	buf := make([]byte, 256*1024)
@@ -168,7 +246,12 @@ func (a *App) copyWithProgress(dst *os.File, src io.Reader, t *Task) (int64, err
 					progress = float64(n) * 100 / float64(t.TotalBytes)
 				}
 				t.DownloadedBytes, t.Progress = n, progress
-				a.db.Model(t).Updates(map[string]any{"downloaded_bytes": n, "progress": progress})
+				if err := a.db.Model(t).Updates(map[string]any{
+					"downloaded_bytes": n, "progress": progress,
+				}).Error; err != nil {
+					return n, err
+				}
+				a.markAssetHeartbeat(t.AssetID, t.ID)
 				last = n
 			}
 		}