am-downloader/concurrency_test.go

package main

import (
	"os"
	"path/filepath"
	"strconv"
	"sync"
	"testing"

	"github.com/glebarez/sqlite"
	"gorm.io/gorm"
	"time"
)

func TestClaimTaskIsAtomic(t *testing.T) {
	a := newCreditTestApp(t)
	task := Task{UserID: 1, Source: "kw", RID: "atomic", Quality: "lossless", Status: "pending"}
	if err := a.db.Create(&task).Error; err != nil {
		t.Fatal(err)
	}
	if _, err := a.claimTask(task.ID); err != nil {
		t.Fatal(err)
	}
	if _, err := a.claimTask(task.ID); err != errTaskRunning {
		t.Fatalf("second claim error = %v, want task running", err)
	}
}

func TestRetryCommittedReservationDoesNotChargeAgain(t *testing.T) {
	a := newCreditTestApp(t)
	user := User{Username: "retry-committed", Role: RoleUser, Status: UserActive}
	if err := a.db.Create(&user).Error; err != nil {
		t.Fatal(err)
	}
	if err := a.db.Create(&CreditAccount{UserID: user.ID, Balance: 1}).Error; err != nil {
		t.Fatal(err)
	}
	task := Task{UserID: user.ID, Source: "kw", RID: "committed", Quality: "lossless", Status: "pending"}
	if err := a.db.Create(&task).Error; err != nil {
		t.Fatal(err)
	}
	r, err := a.reserveCredit(user.ID, task.ID, 1)
	if err != nil {
		t.Fatal(err)
	}
	task.CreditReservationID = r.ID
	if err := a.db.Model(&task).Update("credit_reservation_id", r.ID).Error; err != nil {
		t.Fatal(err)
	}
	if err := a.commitCredit(&task); err != nil {
		t.Fatal(err)
	}
	if err := a.db.Model(&task).Updates(map[string]any{"status": "error", "error": "legacy"}).Error; err != nil {
		t.Fatal(err)
	}
	var current Task
	if err := a.db.First(&current, task.ID).Error; err != nil {
		t.Fatal(err)
	}
	if err := a.prepareRetry(&current); err != nil {
		t.Fatal(err)
	}
	var account CreditAccount
	if err := a.db.Where("user_id = ?", user.ID).First(&account).Error; err != nil {
		t.Fatal(err)
	}
	if account.Balance != 0 || account.Reserved != 0 {
		t.Fatalf("credit changed on committed retry: %+v", account)
	}
	if current.Status != "upload_pending" {
		t.Fatalf("status = %q, want upload_pending", current.Status)
	}
}

func TestRetryRefundedReservationCreatesNewReservation(t *testing.T) {
	a := newCreditTestApp(t)
	user := User{Username: "retry-refunded", Role: RoleUser, Status: UserActive}
	if err := a.db.Create(&user).Error; err != nil {
		t.Fatal(err)
	}
	if err := a.db.Create(&CreditAccount{UserID: user.ID, Balance: 1}).Error; err != nil {
		t.Fatal(err)
	}
	task := Task{UserID: user.ID, Source: "kw", RID: "refunded", Quality: "lossless", Status: "error"}
	if err := a.db.Create(&task).Error; err != nil {
		t.Fatal(err)
	}
	r1, err := a.reserveCredit(user.ID, task.ID, 1)
	if err != nil {
		t.Fatal(err)
	}
	task.CreditReservationID = r1.ID
	if err := a.refundCredit(&task); err != nil {
		t.Fatal(err)
	}
	if err := a.prepareRetry(&task); err != nil {
		t.Fatal(err)
	}
	if task.CreditReservationID == r1.ID {
		t.Fatal("retry reused refunded reservation")
	}
	var account CreditAccount
	if err := a.db.Where("user_id = ?", user.ID).First(&account).Error; err != nil {
		t.Fatal(err)
	}
	if account.Balance != 1 || account.Reserved != 1 {
		t.Fatalf("unexpected account: %+v", account)
	}
}

func TestStaleAssetLeaseStealsSafely(t *testing.T) {
	a := newCreditTestApp(t)
	root := t.TempDir()
	a.cfg.DownloadDir = root
	first := Task{UserID: 1, Source: "kw", RID: "lease", Quality: "lossless", Format: "flac", Status: "downloading"}
	second := Task{UserID: 2, Source: "kw", RID: "lease", Quality: "lossless", Format: "flac", Status: "pending"}
	if err := a.db.Create(&first).Error; err != nil {
		t.Fatal(err)
	}
	if err := a.db.Create(&second).Error; err != nil {
		t.Fatal(err)
	}
	asset := MediaAsset{
		AssetKey: assetKey("kw", "lease", "lossless"),
		Source:   "kw", RID: "lease", Quality: "lossless",
		LocalPath: filepath.Join(root, "old.flac"), Format: "flac",
		Status: "downloading", OwnerTaskID: first.ID,
		LeaseUntil: time.Now().Add(-time.Minute), Generation: 1,
	}
	if err := a.db.Create(&asset).Error; err != nil {
		t.Fatal(err)
	}
	first.AssetID = asset.ID
	if err := a.db.Model(&first).Update("asset_id", asset.ID).Error; err != nil {
		t.Fatal(err)
	}
	got, reused, err := a.acquireAsset(&second)
	if err != nil || reused {
		t.Fatalf("acquire = asset=%+v reused=%v err=%v", got, reused, err)
	}
	if got.Generation != 2 {
		t.Fatalf("generation = %d, want 2", got.Generation)
	}
	if got.LocalPath == asset.LocalPath {
		t.Fatal("stolen generation reused old local path")
	}
	var old Task
	if err := a.db.First(&old, first.ID).Error; err != nil {
		t.Fatal(err)
	}
	if old.Status != "pending" {
		t.Fatalf("stale owner status = %q, want pending", old.Status)
	}
}

func TestAssetCachePathIsStablePerGeneration(t *testing.T) {
	root := t.TempDir()
	p1 := assetCachePath(root, "kw-rid-lossless", "flac", 1)
	p2 := assetCachePath(root, "kw-rid-lossless", "flac", 2)
	if p1 == p2 || filepath.Dir(p1) != filepath.Dir(p2) {
		t.Fatalf("unexpected cache paths: %q %q", p1, p2)
	}
	if err := os.MkdirAll(filepath.Dir(p1), 0755); err != nil {
		t.Fatal(err)
	}
}

func TestConcurrentCreditReservationsAreAtomic(t *testing.T) {
	dbPath := filepath.Join(t.TempDir(), "credits.db")
	db, err := gorm.Open(sqlite.Open(dbPath+"?_pragma=journal_mode(WAL)&_pragma=busy_timeout(10000)"), &gorm.Config{})
	if err != nil {
		t.Fatal(err)
	}
	if err := db.AutoMigrate(&Task{}, &MediaAsset{}, &User{}, &CreditAccount{}, &CreditTransaction{}, &CreditReservation{}); err != nil {
		t.Fatal(err)
	}
	app := &App{db: db}
	user := User{Username: "concurrent-credit", Role: RoleUser, Status: UserActive}
	if err := db.Create(&user).Error; err != nil {
		t.Fatal(err)
	}
	if err := db.Create(&CreditAccount{UserID: user.ID, Balance: 5}).Error; err != nil {
		t.Fatal(err)
	}
	for i := 0; i < 20; i++ {
		if err := db.Create(&Task{UserID: user.ID, Source: "kw", RID: "c-" + strconv.Itoa(i), Quality: "lossless", Status: "pending"}).Error; err != nil {
			t.Fatal(err)
		}
	}
	var wg sync.WaitGroup
	var mu sync.Mutex
	success, unexpected := 0, 0
	for i := 0; i < 20; i++ {
		var task Task
		if err := db.Where("user_id = ? AND r_id = ?", user.ID, "c-"+strconv.Itoa(i)).First(&task).Error; err != nil {
			t.Fatal(err)
		}
		wg.Add(1)
		go func(taskID uint) {
			defer wg.Done()
			if _, err := app.reserveCredit(user.ID, taskID, 1); err == nil {
				mu.Lock()
				success++
				mu.Unlock()
			} else if err != errInsufficientCredits {
				mu.Lock()
				unexpected++
				mu.Unlock()
			}
		}(task.ID)
	}
	wg.Wait()
	if success != 5 || unexpected != 0 {
		t.Fatalf("reservations: success=%d unexpected=%d", success, unexpected)
	}
	var account CreditAccount
	if err := db.Where("user_id = ?", user.ID).First(&account).Error; err != nil {
		t.Fatal(err)
	}
	if account.Balance != 5 || account.Reserved != 5 {
		t.Fatalf("account after concurrent reservations: %+v", account)
	}
}