diff --git a/internal/engine/fleet_speed_test.go b/internal/engine/fleet_speed_test.go new file mode 100644 index 0000000..631e6e8 --- /dev/null +++ b/internal/engine/fleet_speed_test.go @@ -0,0 +1,50 @@ +package engine + +import ( + "context" + "testing" + + "marathon/internal/models" + "marathon/internal/store" +) + +// Proves a fleet worker adopts a live rate change: refreshRate re-reads the +// task's rate and updates the worker between leases, so the dashboard speed +// slider reaches already-running workers. +func TestIntegration_FleetLiveSpeed(t *testing.T) { + control, _ := openOrSkip(t) + controlSchema(t, control) + + task := &models.Task{ + Name: "fleet-speed", TargetDSN: dsn("MARATHON_TEST_TARGET_DSN", "postgres://demo:demo@localhost:5434/demo?sslmode=disable"), + SourceTable: "items", CursorColumn: "id", BatchSize: 1000, RatePerSec: 1000, + OperationSQL: "UPDATE items SET applied = applied + 1 WHERE id >= $1 AND id <= $2", + } + newTask(t, control, task) + + w := NewLeaseWorker("w0", control, nopLogger{}, nil, *task, 1) + if w.lastRate != 1000 { + t.Fatalf("initial lastRate = %d, want 1000", w.lastRate) + } + + // Operator drags the slider → the task rate changes in the control store. + if err := (store.TaskStore{}).SetRate(context.Background(), control, task.ID, 5000); err != nil { + t.Fatal(err) + } + + // Between leases the worker refreshes and adopts the new rate. + w.refreshRate(context.Background()) + if w.lastRate != 5000 { + t.Errorf("after refresh lastRate = %d, want 5000", w.lastRate) + } + + // A further change to unthrottled (0) is also picked up. + if err := (store.TaskStore{}).SetRate(context.Background(), control, task.ID, 0); err != nil { + t.Fatal(err) + } + + w.refreshRate(context.Background()) + if w.lastRate != 0 { + t.Errorf("after refresh lastRate = %d, want 0 (unthrottled)", w.lastRate) + } +} diff --git a/internal/engine/leaserunner.go b/internal/engine/leaserunner.go index 653e040..e0ced7f 100644 --- a/internal/engine/leaserunner.go +++ b/internal/engine/leaserunner.go @@ -33,25 +33,41 @@ type LeaseWorker struct { queue *coord.Queue task models.Task runID int64 - bucket *throttle.Bucket - emit ProgressFunc + bucket *throttle.Bucket + lastRate int + emit ProgressFunc + tasks store.TaskStore runs store.RunStore quarantine store.QuarantineStore } func NewLeaseWorker(id string, db store.DB, log store.Logger, queue *coord.Queue, task models.Task, runID int64) *LeaseWorker { return &LeaseWorker{ - id: id, - db: db, - log: log, - queue: queue, - task: task, - runID: runID, - bucket: throttle.New(task.RatePerSec), + id: id, + db: db, + log: log, + queue: queue, + task: task, + runID: runID, + bucket: throttle.New(task.RatePerSec), + lastRate: task.RatePerSec, } } +// refreshRate picks up a live throttle change (from the dashboard speed slider) +// by re-reading the task's rate and adjusting this worker's bucket. Called once +// per lease, so already-running fleet workers respond to speed changes. +func (w *LeaseWorker) refreshRate(ctx context.Context) { + rate, err := w.tasks.GetRate(ctx, w.db, w.task.ID) + if err != nil || rate == w.lastRate { + return + } + + w.bucket.SetRate(rate) + w.lastRate = rate +} + func (w *LeaseWorker) WithProgress(fn ProgressFunc) *LeaseWorker { w.emit = fn return w @@ -120,6 +136,8 @@ func (w *LeaseWorker) Run(ctx context.Context) error { continue } + w.refreshRate(ctx) // pick up any live speed change before this lease + if err := w.processLease(ctx, target, lease); err != nil { // Leave the lease in processing; the reaper will requeue it. return err diff --git a/internal/store/task.go b/internal/store/task.go index 5c8671d..d0fd15d 100644 --- a/internal/store/task.go +++ b/internal/store/task.go @@ -28,6 +28,18 @@ func (TaskStore) Create(ctx context.Context, db DB, t *models.Task) error { return nil } +// GetRate returns a task's current throttle rate (rows/sec, 0 = unlimited). +func (TaskStore) GetRate(ctx context.Context, db DB, id int64) (int, error) { + var rate int + + err := db.QueryRowContext(ctx, `SELECT rate_per_sec FROM tasks WHERE id = $1`, id).Scan(&rate) + if err != nil { + return 0, fmt.Errorf("get task %d rate: %w", id, err) + } + + return rate, nil +} + // SetRate updates a task's throttle rate (rows/sec, 0 = unlimited). func (TaskStore) SetRate(ctx context.Context, db DB, id int64, rate int) error { _, err := db.ExecContext(ctx, `UPDATE tasks SET rate_per_sec = $1 WHERE id = $2`, rate, id)