256 lines
7.3 KiB
Go
256 lines
7.3 KiB
Go
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(¤t, 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(¤t).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
|
|
}
|