Fleet-mode live throttle: running workers pick up speed changes

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)
+	}
+}

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

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)