Adaptive auto-throttle (AIMD) — steer rate by target latency

internal/engine/adaptive.go

@@ -0,0 +1,55 @@
+package engine
+
+import (
+	"context"
+	"database/sql"
+	"time"
+
+	"marathon/internal/engine/throttle"
+)
+
+// adaptiveDefaultCeiling bounds an adaptive run whose task has no explicit rate
+// (rate_per_sec == 0). Without a ceiling the controller would ramp unbounded.
+const adaptiveDefaultCeiling = 50_000
+
+// probeInterval limits how often we sample target latency — at most once per
+// this window — so a fast batch loop doesn't flood the DB with probes.
+const probeInterval = 250 * time.Millisecond
+
+// adaptiveCeiling turns a task rate into the controller's upper bound: the
+// task's rate when set, otherwise a sane default.
+func adaptiveCeiling(taskRate int) float64 {
+	if taskRate > 0 {
+		return float64(taskRate)
+	}
+
+	return adaptiveDefaultCeiling
+}
+
+// newAdaptive builds an AIMD controller starting at the ceiling and ramping
+// within [min, ceiling] based on observed target latency.
+func newAdaptive(taskRate int) *throttle.AIMD {
+	ceiling := adaptiveCeiling(taskRate)
+
+	return throttle.NewAIMD(throttle.AIMDConfig{
+		Start:   ceiling, // begin optimistic; back off on the first painful sample
+		Ceiling: ceiling,
+	})
+}
+
+// probeQuery is the round-trip used to sample target latency. A test seam lets
+// integration tests substitute a deliberately slow query to exercise back-off.
+var probeQuery = "SELECT 1"
+
+// probeLatency times a trivial round-trip to the target. This is the health
+// signal the controller reacts to: when the DB is loaded, even `SELECT 1` slows.
+func probeLatency(ctx context.Context, db *sql.DB) time.Duration {
+	start := time.Now()
+
+	if err := db.QueryRowContext(ctx, probeQuery).Scan(new(int)); err != nil {
+		// A failed/slow probe is itself a strong "back off" signal.
+		return time.Second
+	}
+
+	return time.Since(start)
+}

internal/engine/adaptive_test.go

@@ -0,0 +1,96 @@
+package engine
+
+import (
+	"context"
+	"testing"
+
+	"marathon/internal/models"
+	"marathon/internal/store"
+)
+
+func TestAdaptiveCeiling(t *testing.T) {
+	if got := adaptiveCeiling(8000); got != 8000 {
+		t.Errorf("adaptiveCeiling(8000) = %v, want 8000", got)
+	}
+
+	if got := adaptiveCeiling(0); got != adaptiveDefaultCeiling {
+		t.Errorf("adaptiveCeiling(0) = %v, want default %d", got, adaptiveDefaultCeiling)
+	}
+}
+
+// An adaptive run completes correctly against a healthy DB — the feedback loop
+// must never break correctness (every row applied exactly once).
+func TestIntegration_AdaptiveRunCompletes(t *testing.T) {
+	control, target := openOrSkip(t)
+	controlSchema(t, control)
+
+	const n = 10000
+	seedTarget(t, target, "items", n, "")
+
+	task := &models.Task{
+		Name: "adaptive-ok", TargetDSN: dsn("MARATHON_TEST_TARGET_DSN", "postgres://demo:demo@localhost:5434/demo?sslmode=disable"),
+		SourceTable: "items", CursorColumn: "id", BatchSize: 1000,
+		RatePerSec: 20000, Adaptive: true,
+		OperationSQL: "UPDATE items SET applied = applied + 1 WHERE id >= $1 AND id <= $2",
+	}
+	newTask(t, control, task)
+
+	run, err := (store.RunStore{}).Create(context.Background(), control, task.ID)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	r := NewRunner(control, nopLogger{}, *task, run.ID)
+	if r.adaptive == nil {
+		t.Fatal("adaptive controller should be set for an adaptive task")
+	}
+
+	r.Run(context.Background())
+
+	assertState(t, control, run.ID, models.RunSucceeded)
+	assertCount(t, target, "SELECT count(*) FROM items WHERE applied <> 1", 0, "rows not applied exactly once")
+}
+
+// Under injected latency the controller backs the rate off well below its
+// ceiling — the point of the auto-throttle.
+func TestIntegration_AdaptiveBacksOffUnderLatency(t *testing.T) {
+	control, target := openOrSkip(t)
+	controlSchema(t, control)
+
+	const n = 4000
+	seedTarget(t, target, "items", n, "")
+
+	// Make each probe take ~80ms (> the 50ms high watermark) so every sample
+	// tells the controller to decrease.
+	orig := probeQuery
+	probeQuery = "SELECT 1 FROM pg_sleep(0.08)"
+	defer func() { probeQuery = orig }()
+
+	const ceiling = 20000
+	task := &models.Task{
+		Name: "adaptive-backoff", TargetDSN: dsn("MARATHON_TEST_TARGET_DSN", "postgres://demo:demo@localhost:5434/demo?sslmode=disable"),
+		SourceTable: "items", CursorColumn: "id", BatchSize: 500,
+		RatePerSec: ceiling, Adaptive: true,
+		OperationSQL: "UPDATE items SET applied = applied + 1 WHERE id >= $1 AND id <= $2",
+	}
+	newTask(t, control, task)
+
+	run, err := (store.RunStore{}).Create(context.Background(), control, task.ID)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	r := NewRunner(control, nopLogger{}, *task, run.ID)
+	r.Run(context.Background())
+
+	assertState(t, control, run.ID, models.RunSucceeded)
+
+	// The controller should have driven the rate far below the ceiling in
+	// response to the sustained high latency.
+	if got := r.adaptive.Rate(); got >= ceiling {
+		t.Errorf("adaptive rate = %.0f, expected well below ceiling %d after sustained latency", got, ceiling)
+	}
+
+	// Correctness is unaffected by throttling.
+	assertCount(t, target, "SELECT count(*) FROM items WHERE applied <> 1", 0, "rows not applied exactly once")
+}

internal/engine/integration_test.go

@@ -99,7 +99,8 @@ func controlSchema(t *testing.T, db *sql.DB) {
 		ADD COLUMN IF NOT EXISTS operation_url    TEXT NOT NULL DEFAULT '',
 		ADD COLUMN IF NOT EXISTS target_driver    TEXT NOT NULL DEFAULT 'postgres',
 		ADD COLUMN IF NOT EXISTS schedule_seconds INT  NOT NULL DEFAULT 0,
-		ADD COLUMN IF NOT EXISTS connection_id    BIGINT REFERENCES connections(id)`); err != nil {
+		ADD COLUMN IF NOT EXISTS connection_id    BIGINT REFERENCES connections(id),
+		ADD COLUMN IF NOT EXISTS adaptive         BOOLEAN NOT NULL DEFAULT false`); err != nil {
 		t.Fatalf("control alter: %v", err)
 	}
 

internal/engine/leaserunner.go

@@ -33,17 +33,19 @@ type LeaseWorker struct {
 	queue  *coord.Queue
 	task   models.Task
 	runID  int64
-	bucket   *throttle.Bucket
-	lastRate int
-	emit     ProgressFunc
+	bucket    *throttle.Bucket
+	adaptive  *throttle.AIMD // non-nil when the task opts into auto-throttle
+	lastRate  int
+	lastProbe time.Time
+	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{
+	w := &LeaseWorker{
 		id:       id,
 		db:       db,
 		log:      log,
@@ -53,19 +55,47 @@ func NewLeaseWorker(id string, db store.DB, log store.Logger, queue *coord.Queue
 		bucket:   throttle.New(task.RatePerSec),
 		lastRate: task.RatePerSec,
 	}
+
+	if task.Adaptive {
+		w.adaptive = newAdaptive(task.RatePerSec)
+	}
+
+	return w
 }
 
 // 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.
+// by re-reading the task's rate. Called once per lease, so already-running fleet
+// workers respond to speed changes. In adaptive mode the rate is the ceiling and
+// the controller still governs the actual rate; otherwise it sets the bucket.
 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
+
+	if w.adaptive != nil {
+		w.adaptive.SetCeiling(adaptiveCeiling(rate))
+		return
+	}
+
+	w.bucket.SetRate(rate)
+}
+
+// tune samples target latency (rate-limited) and steers the bucket rate via the
+// AIMD controller. No-op unless the task opts into adaptive throttling.
+func (w *LeaseWorker) tune(ctx context.Context, target *sql.DB) {
+	if w.adaptive == nil {
+		return
+	}
+
+	if !w.lastProbe.IsZero() && time.Since(w.lastProbe) < probeInterval {
+		return
+	}
+
+	w.lastProbe = time.Now()
+	w.bucket.SetRate(int(w.adaptive.Observe(probeLatency(ctx, target))))
 }
 
 func (w *LeaseWorker) WithProgress(fn ProgressFunc) *LeaseWorker {
@@ -170,6 +200,8 @@ func (w *LeaseWorker) processLease(ctx context.Context, target *sql.DB, lease co
 			return ctx.Err()
 		}
 
+		w.tune(ctx, target)
+
 		if err := w.bucket.Take(ctx, w.task.BatchSize); err != nil {
 			return err
 		}

internal/engine/runner.go

@@ -8,6 +8,7 @@ import (
 	"database/sql"
 	"errors"
 	"fmt"
+	"sync/atomic"
 	"time"
 
 	_ "github.com/jackc/pgx/v5/stdlib" // target-DB driver
@@ -34,8 +35,10 @@ type Runner struct {
 	task  models.Task
 	runID int64
 
-	bucket *throttle.Bucket
-	emit   ProgressFunc
+	bucket   *throttle.Bucket
+	adaptive *throttle.AIMD // non-nil when the task opts into auto-throttle
+	ceiling  atomic.Int64   // live ceiling (rows/sec) for adaptive mode
+	emit     ProgressFunc
 
 	batchIdx int
 	// crashHook, when set, is called after a batch's operation commits on the
@@ -66,13 +69,20 @@ type Progress struct {
 }
 
 func NewRunner(db store.DB, log store.Logger, task models.Task, runID int64) *Runner {
-	return &Runner{
+	r := &Runner{
 		db:     db,
 		log:    log,
 		task:   task,
 		runID:  runID,
 		bucket: throttle.New(task.RatePerSec),
 	}
+
+	if task.Adaptive {
+		r.adaptive = newAdaptive(task.RatePerSec)
+		r.ceiling.Store(int64(adaptiveCeiling(task.RatePerSec)))
+	}
+
+	return r
 }
 
 // WithProgress registers a live-progress emitter.
@@ -82,8 +92,16 @@ func (r *Runner) WithProgress(fn ProgressFunc) *Runner {
 }
 
 // SetRate changes the throttle rate (rows/sec, 0 = unlimited) of a running
-// runner, live. Used by the dashboard speed slider.
+// runner, live. Used by the dashboard speed slider. In adaptive mode the slider
+// sets the ceiling; the controller still governs the actual rate below it. The
+// ceiling is applied in the runner goroutine (via an atomic) so the AIMD is
+// never mutated concurrently.
 func (r *Runner) SetRate(rowsPerSec int) {
+	if r.adaptive != nil {
+		r.ceiling.Store(int64(adaptiveCeiling(rowsPerSec)))
+		return
+	}
+
 	r.bucket.SetRate(rowsPerSec)
 }
 
@@ -146,6 +164,8 @@ func (r *Runner) run(ctx context.Context) error {
 		}
 	}
 
+	var lastProbe time.Time
+
 	for {
 		state, err := r.runs.GetState(ctx, r.db, r.runID)
 		if err != nil {
@@ -169,6 +189,8 @@ func (r *Runner) run(ctx context.Context) error {
 			return nil // terminal state reached elsewhere; stop quietly
 		}
 
+		r.tune(ctx, target, &lastProbe)
+
 		if err := r.bucket.Take(ctx, r.task.BatchSize); err != nil {
 			return err
 		}
@@ -215,6 +237,25 @@ func (r *Runner) run(ctx context.Context) error {
 	}
 }
 
+// tune, in adaptive mode, samples target latency (rate-limited by probeInterval)
+// and steers the token-bucket rate via the AIMD controller. All AIMD access is
+// in this goroutine; SetRate only stores the ceiling atomically.
+func (r *Runner) tune(ctx context.Context, target *sql.DB, lastProbe *time.Time) {
+	if r.adaptive == nil {
+		return
+	}
+
+	if !lastProbe.IsZero() && time.Since(*lastProbe) < probeInterval {
+		return
+	}
+
+	*lastProbe = time.Now()
+
+	r.adaptive.SetCeiling(float64(r.ceiling.Load()))
+	rate := r.adaptive.Observe(probeLatency(ctx, target))
+	r.bucket.SetRate(int(rate))
+}
+
 // emitSnapshot publishes the run's current state to live subscribers.
 func (r *Runner) emitSnapshot(ctx context.Context, batchRows int) {
 	if r.emit == nil {

internal/models/task.go

@@ -35,6 +35,10 @@ type Task struct {
 	OperationSQL  string    `json:"operation_sql"`
 	OperationURL  string    `json:"operation_url"`
 	RatePerSec    int       `json:"rate_per_sec"`
+	// Adaptive enables the AIMD auto-throttle: the worker watches the target
+	// database's latency and adjusts the rate automatically, backing off when
+	// the DB is under load. rate_per_sec then acts as the ceiling.
+	Adaptive bool `json:"adaptive"`
 	// ScheduleSeconds > 0 makes the task recurring: the scheduler starts a new
 	// run every N seconds (skipping ticks while a run is still active). 0 = manual.
 	ScheduleSeconds int       `json:"schedule_seconds"`