feat: harden multi-user download concurrency
This commit is contained in:
OpenClaw Agent
@@ -0,0 +1,223 @@
|
||||
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(¤t, task.ID).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := a.prepareRetry(¤t); 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)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user