diff --git a/internal/auth/rbac.go b/internal/auth/rbac.go index 1bb2cad..15babea 100644 --- a/internal/auth/rbac.go +++ b/internal/auth/rbac.go @@ -81,10 +81,12 @@ func requiredRole(method string) Role { } // openPaths bypass auth entirely: health/observability probes and the static -// dashboard assets (the dashboard's API calls still enforce RBAC). +// dashboard assets (the dashboard's API calls still enforce RBAC). The /app +// match is boundary-anchored so unrelated paths like "/apple" or a traversal +// like "/app/../tasks" don't slip past auth. func openPath(path string) bool { return strings.HasPrefix(path, "/.well-known/") || - strings.HasPrefix(path, "/app") || + path == "/app" || strings.HasPrefix(path, "/app/") || path == "/" || path == "/health" || path == "/metrics" || path == "/favicon.ico" } diff --git a/internal/coord/lease.go b/internal/coord/lease.go index 7fb6d80..fcd2606 100644 --- a/internal/coord/lease.go +++ b/internal/coord/lease.go @@ -60,19 +60,34 @@ func (q *Queue) key(suffix string) string { return fmt.Sprintf("mara:run:%d:%s", q.runID, suffix) } +// enqueueChunk bounds how many RPUSH commands go in one pipeline, so publishing +// a large lease set doesn't build one huge Redis transaction. +const enqueueChunk = 1000 + // Enqueue publishes all leases for the run and seals it (no more will be added). // Sealing lets workers distinguish "queue momentarily empty" from "run done". +// Leases are pushed in chunks to keep each Redis round-trip bounded. func (q *Queue) Enqueue(ctx context.Context, leases []Lease) error { - pipe := q.rdb.TxPipeline() + for start := 0; start < len(leases); start += enqueueChunk { + end := start + enqueueChunk + if end > len(leases) { + end = len(leases) + } - for _, l := range leases { - pipe.RPush(ctx, q.key("pending"), l.encode()) - } + pipe := q.rdb.Pipeline() + for _, l := range leases[start:end] { + pipe.RPush(ctx, q.key("pending"), l.encode()) + } - pipe.Set(ctx, q.key("sealed"), "1", 0) + if _, err := pipe.Exec(ctx); err != nil { + return fmt.Errorf("enqueue leases: %w", err) + } + } - if _, err := pipe.Exec(ctx); err != nil { - return fmt.Errorf("enqueue leases: %w", err) + // Seal only after every lease is enqueued, so a worker never sees an empty + // pending list as "drained" mid-publish. + if err := q.rdb.Set(ctx, q.key("sealed"), "1", 0).Err(); err != nil { + return fmt.Errorf("seal queue: %w", err) } return nil @@ -119,6 +134,34 @@ func (q *Queue) Complete(ctx context.Context, l Lease) error { return nil } +// Requeue returns a lease to pending after the owning worker failed to process +// it (e.g. target unreachable), and bumps the failure counter. Another worker +// will retry it; the monitor uses the failure counter to detect a run that is +// failing with no forward progress. +func (q *Queue) Requeue(ctx context.Context, l Lease) error { + pipe := q.rdb.TxPipeline() + pipe.LRem(ctx, q.key("processing"), 1, l.encode()) + pipe.RPush(ctx, q.key("pending"), l.encode()) + pipe.ZRem(ctx, q.key("hb"), l.encode()) + pipe.Incr(ctx, q.key("fail")) + + if _, err := pipe.Exec(ctx); err != nil { + return fmt.Errorf("requeue failed lease: %w", err) + } + + return nil +} + +// Fails returns the cumulative lease-failure count for the run. +func (q *Queue) Fails(ctx context.Context) (int64, error) { + n, err := q.rdb.Get(ctx, q.key("fail")).Int64() + if err == redis.Nil { + return 0, nil + } + + return n, err +} + // Reap requeues leases whose heartbeat is older than the cutoff — the // signature move of dead-worker recovery. Returns how many were requeued. func (q *Queue) Reap(ctx context.Context, cutoffMillis int64) (int, error) { @@ -176,25 +219,50 @@ func (q *Queue) Drained(ctx context.Context) (bool, error) { func (q *Queue) Cleanup(ctx context.Context) error { return q.rdb.Del(ctx, q.key("pending"), q.key("processing"), q.key("hb"), - q.key("sealed"), q.key("done"), + q.key("sealed"), q.key("done"), q.key("fail"), ).Err() } -// PlanLeases divides [minID, maxID] into leases of the given stride. +// MaxLeases caps how many leases a single run is split into, regardless of how +// large or sparse the key space is. This bounds memory and the Redis enqueue +// fan-out; each lease just covers a wider id range (workers sub-batch within). +const MaxLeases = 100_000 + +// PlanLeases divides [minID, maxID] into at most MaxLeases leases. It widens the +// stride when the key space is huge or sparse (e.g. snowflake ids up to 10^18) +// and is overflow-safe near math.MaxInt64, so it can never build billions of +// leases or spin forever. func PlanLeases(minID, maxID, stride int64) []Lease { + if maxID < minID { + return nil + } + if stride <= 0 { stride = 1 } - var leases []Lease + // Widen the stride so the lease count can't exceed MaxLeases. float64 avoids + // int64 overflow on very wide spans; the hard loop cap below is the backstop. + spanF := float64(maxID) - float64(minID) + 1 + if spanF/float64(stride) > float64(MaxLeases) { + stride = int64(spanF/float64(MaxLeases)) + 1 + } + + leases := make([]Lease, 0, MaxLeases+1) - for lo := minID; lo <= maxID; lo += stride { + for lo := minID; len(leases) <= MaxLeases; { hi := lo + stride - 1 - if hi > maxID { + if hi < lo || hi > maxID { // overflow (wrapped) or past the end hi = maxID } leases = append(leases, Lease{Lo: lo, Hi: hi}) + + if hi >= maxID { + break // full coverage reached + } + + lo = hi + 1 } return leases diff --git a/internal/coord/lease_test.go b/internal/coord/lease_test.go index 1fe4480..2436ec8 100644 --- a/internal/coord/lease_test.go +++ b/internal/coord/lease_test.go @@ -29,6 +29,43 @@ func TestPlanLeases(t *testing.T) { } } +func TestPlanLeasesBounded(t *testing.T) { + // A huge/sparse key space with a tiny stride must NOT produce billions of + // leases or hang — the count is capped and coverage is complete. + got := PlanLeases(1, 1_000_000_000_000, 1000) + if len(got) > MaxLeases+1 { + t.Fatalf("lease count = %d, want <= %d", len(got), MaxLeases+1) + } + + if got[0].Lo != 1 { + t.Errorf("first lease Lo = %d, want 1", got[0].Lo) + } + + if last := got[len(got)-1]; last.Hi != 1_000_000_000_000 { + t.Errorf("last lease Hi = %d, want full coverage to 10^12", last.Hi) + } + + // No gaps or overlaps across the whole plan. + for i := 1; i < len(got); i++ { + if got[i].Lo != got[i-1].Hi+1 { + t.Fatalf("gap/overlap at %d: prev.Hi=%d cur.Lo=%d", i, got[i-1].Hi, got[i].Lo) + } + } +} + +func TestPlanLeasesOverflowSafe(t *testing.T) { + // Near math.MaxInt64 the stride addition would overflow; must terminate. + const maxInt64 = int64(9223372036854775807) + got := PlanLeases(maxInt64-10, maxInt64, 4) + if len(got) == 0 || len(got) > MaxLeases+1 { + t.Fatalf("overflow plan len = %d", len(got)) + } + + if got[len(got)-1].Hi != maxInt64 { + t.Errorf("last Hi = %d, want %d", got[len(got)-1].Hi, maxInt64) + } +} + func TestLeaseRoundTrip(t *testing.T) { l := Lease{Lo: 42, Hi: 1337} diff --git a/internal/engine/hardening_test.go b/internal/engine/hardening_test.go new file mode 100644 index 0000000..b4bc83f --- /dev/null +++ b/internal/engine/hardening_test.go @@ -0,0 +1,103 @@ +package engine + +import ( + "context" + "errors" + "sync" + "testing" + + "marathon/internal/models" + "marathon/internal/store" +) + +// The partial unique index must allow only one active run per task, even under +// a burst of concurrent Create calls (the duplicate-run race). +func TestIntegration_OneActiveRunPerTask(t *testing.T) { + control, _ := openOrSkip(t) + controlSchema(t, control) + + task := &models.Task{ + Name: "dup-guard", TargetDSN: dsn("MARATHON_TEST_TARGET_DSN", "postgres://demo:demo@localhost:5434/demo?sslmode=disable"), + SourceTable: "items", CursorColumn: "id", BatchSize: 1000, + OperationSQL: "UPDATE items SET applied = applied + 1 WHERE id >= $1 AND id <= $2", + } + newTask(t, control, task) + + const racers = 12 + var ( + wg sync.WaitGroup + mu sync.Mutex + created int + rejected int + ) + + for range racers { + wg.Add(1) + go func() { + defer wg.Done() + _, err := (store.RunStore{}).Create(context.Background(), control, task.ID) + + mu.Lock() + defer mu.Unlock() + switch { + case err == nil: + created++ + case errors.Is(err, store.ErrActiveRunExists): + rejected++ + default: + t.Errorf("unexpected error: %v", err) + } + }() + } + wg.Wait() + + if created != 1 { + t.Errorf("created %d active runs, want exactly 1", created) + } + + if rejected != racers-1 { + t.Errorf("rejected %d, want %d", rejected, racers-1) + } + + // Confirm the DB agrees: exactly one active run. + var active int + control.QueryRow(`SELECT count(*) FROM runs WHERE task_id=$1 AND state IN ('queued','running','paused')`, task.ID).Scan(&active) + if active != 1 { + t.Errorf("db shows %d active runs, want 1", active) + } +} + +// A run's operation is available for a new run only once the previous one is +// terminal — after which Create succeeds again. +func TestIntegration_ActiveRunFreesOnTerminal(t *testing.T) { + control, _ := openOrSkip(t) + controlSchema(t, control) + + task := &models.Task{ + Name: "free-on-terminal", TargetDSN: dsn("MARATHON_TEST_TARGET_DSN", "postgres://demo:demo@localhost:5434/demo?sslmode=disable"), + SourceTable: "items", CursorColumn: "id", BatchSize: 1000, + OperationSQL: "UPDATE items SET applied = applied + 1 WHERE id >= $1 AND id <= $2", + } + newTask(t, control, task) + + runs := store.RunStore{} + ctx := context.Background() + + r1, err := runs.Create(ctx, control, task.ID) + if err != nil { + t.Fatal(err) + } + + if _, err := runs.Create(ctx, control, task.ID); !errors.Is(err, store.ErrActiveRunExists) { + t.Fatalf("second Create = %v, want ErrActiveRunExists", err) + } + + // Finish the first run; a new one may now start. + if err := runs.SetState(ctx, control, r1.ID, models.RunSucceeded, ""); err != nil { + t.Fatal(err) + } + + if _, err := runs.Create(ctx, control, task.ID); err != nil { + t.Fatalf("Create after terminal = %v, want success", err) + } +} diff --git a/internal/engine/httpop.go b/internal/engine/httpop.go index 86f9318..02d996b 100644 --- a/internal/engine/httpop.go +++ b/internal/engine/httpop.go @@ -86,8 +86,21 @@ func applyHTTPBatch( } var cr callbackResponse - if err := json.Unmarshal(payload, &cr); err != nil { - return 0, fmt.Errorf("decode callback response: %w", err) + + // A 2xx with a body we can't parse into per-row results is a contract + // violation. We must NOT silently claim all rows succeeded (that would mask + // data never being applied), and we should NOT fail the whole run over one + // bad response. Quarantine the batch's rows so the operator sees it and can + // retry once the callback is fixed. Same for an empty results list. + if err := json.Unmarshal(payload, &cr); err != nil || len(cr.Results) == 0 { + reason := "callback returned no per-row results (expected {\"results\":[{\"row_key\":...,\"ok\":...}]})" + if err != nil { + reason = "callback response was not valid JSON: " + err.Error() + } + + quarantineBatch(ctx, quarantine, control, runID, b, reason, log) + + return 0, nil } var affected, failed int64 @@ -104,11 +117,6 @@ func applyHTTPBatch( } } - // Rows the service never mentioned count as applied (fire-and-forget rows). - if len(cr.Results) == 0 { - affected = int64(len(rows)) - } - if failed > 0 { log.Infof("run %d: batch [%v, %v] callback quarantined %d/%d rows", runID, b.Lo, b.Hi, failed, len(rows)) } @@ -116,6 +124,18 @@ func applyHTTPBatch( return affected, nil } +// quarantineBatch quarantines every row in a batch with the same reason. +func quarantineBatch(ctx context.Context, quarantine store.QuarantineStore, control store.DB, runID int64, b batch.Bounds, reason string, log store.Logger) { + for _, key := range b.Keys { + enc, _ := batch.Encode(key) + if qErr := quarantine.Add(ctx, control, runID, enc, reason); qErr != nil { + log.Errorf("run %d: quarantine write failed: %v", runID, qErr) + } + } + + log.Infof("run %d: batch [%v, %v] callback response unusable, quarantined %d rows: %s", runID, b.Lo, b.Hi, len(b.Keys), reason) +} + // fetchRows loads full rows for the given cursor keys as JSON-ready maps. Each // row carries a "_row_key" (the encoded cursor value) so the callback service // can report per-row success/failure back to us. diff --git a/internal/engine/integration_test.go b/internal/engine/integration_test.go index 3b0f18c..275bfea 100644 --- a/internal/engine/integration_test.go +++ b/internal/engine/integration_test.go @@ -104,6 +104,12 @@ func controlSchema(t *testing.T, db *sql.DB) { t.Fatalf("control alter: %v", err) } + // Partial unique index: at most one active run per task (matches migration 012). + if _, err := db.Exec(`CREATE UNIQUE INDEX IF NOT EXISTS runs_one_active_per_task + ON runs (task_id) WHERE state IN ('queued','running','paused')`); err != nil { + t.Fatalf("control unique index: %v", err) + } + // Clean slate. TRUNCATE ... CASCADE resets identities too. if _, err := db.Exec(`TRUNCATE run_events, quarantined_rows, checkpoints, runs, tasks, connections RESTART IDENTITY CASCADE`); err != nil { t.Fatalf("truncate control: %v", err) diff --git a/internal/engine/leaserunner.go b/internal/engine/leaserunner.go index 2e3060b..c55a82e 100644 --- a/internal/engine/leaserunner.go +++ b/internal/engine/leaserunner.go @@ -83,6 +83,25 @@ func (w *LeaseWorker) refreshRate(ctx context.Context) { w.bucket.SetRate(rate) } +// heartbeatLoop refreshes the lease's liveness on a fixed interval until ctx is +// canceled (lease done, worker stopping, or run killed). Runs in its own +// goroutine so batch duration can never starve heartbeats. +func (w *LeaseWorker) heartbeatLoop(ctx context.Context, lease coord.Lease) { + t := time.NewTicker(leaseHeartbeatEvery) + defer t.Stop() + + for { + select { + case <-ctx.Done(): + return + case <-t.C: + if err := w.queue.Heartbeat(context.Background(), lease, nowMillis()); err != nil { + w.log.Errorf("run %d: heartbeat failed: %v", w.runID, err) + } + } + } +} + // 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) { @@ -109,15 +128,18 @@ func (w *LeaseWorker) WithProgress(fn ProgressFunc) *LeaseWorker { func (w *LeaseWorker) Run(ctx context.Context) error { dialect := batch.DialectFor(w.task.TargetDriver) - target, err := sql.Open(dialect.DriverName(), w.task.TargetDSN) + target, err := openTarget(dialect.DriverName(), w.task.TargetDSN) if err != nil { return fmt.Errorf("open target: %w", err) } defer target.Close() - if err := target.PingContext(ctx); err != nil { + pingCtx, cancelPing := withTimeout(ctx, 10*time.Second) + if err := target.PingContext(pingCtx); err != nil { + cancelPing() return fmt.Errorf("connect target: %w", err) } + cancelPing() for { if ctx.Err() != nil { @@ -169,7 +191,21 @@ func (w *LeaseWorker) Run(ctx context.Context) error { 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. + if ctx.Err() != nil { + return ctx.Err() // killed/shutdown — leave lease for the reaper + } + + // The lease failed (e.g. target unreachable). Requeue it (bumping the + // failure counter) and record the error so the monitor can fail the + // run if failures pile up with no progress. + if rqErr := w.queue.Requeue(ctx, lease); rqErr != nil { + w.log.Errorf("run %d: requeue failed lease: %v", w.runID, rqErr) + } + + if seErr := w.runs.SetError(ctx, w.db, w.runID, err.Error()); seErr != nil { + w.log.Errorf("run %d: record lease error: %v", w.runID, seErr) + } + return err } @@ -191,9 +227,17 @@ func (w *LeaseWorker) processLease(ctx context.Context, target *sql.DB, lease co src := batch.Source{Table: w.task.SourceTable, CursorColumn: w.task.CursorColumn, Filter: rangeFilter, Dialect: dialect} - var after any + // Heartbeat on a background ticker, independent of batch duration. Without + // this, a batch slower than the reap cutoff (throttled rate, slow target) + // would let the monitor falsely reap this still-alive lease and hand it to + // another worker → duplicate processing. The ticker keeps the lease alive + // for as long as this goroutine runs. + beatCtx, stopBeat := context.WithCancel(ctx) + defer stopBeat() - lastBeat := time.Now() + go w.heartbeatLoop(beatCtx, lease) + + var after any for { if ctx.Err() != nil { @@ -206,16 +250,21 @@ func (w *LeaseWorker) processLease(ctx context.Context, target *sql.DB, lease co return err } - b, err := batch.NextBounds(ctx, target, src, after, w.task.BatchSize) + batchCtx, cancelBatch := withTimeout(ctx, batchQueryTimeout) + + b, err := batch.NextBounds(batchCtx, target, src, after, w.task.BatchSize) if err != nil { + cancelBatch() return err } if b.Count == 0 { + cancelBatch() return nil } - affected, err := applyBatch(ctx, target, w.task, b, w.quarantine, w.db, w.runID, w.log) + affected, err := applyBatch(batchCtx, target, w.task, b, w.quarantine, w.db, w.runID, w.log) + cancelBatch() if err != nil { return err } @@ -225,14 +274,6 @@ func (w *LeaseWorker) processLease(ctx context.Context, target *sql.DB, lease co return err } - if time.Since(lastBeat) >= leaseHeartbeatEvery { - if err := w.queue.Heartbeat(ctx, lease, nowMillis()); err != nil { - w.log.Errorf("run %d: heartbeat failed: %v", w.runID, err) - } - - lastBeat = time.Now() - } - w.emitSnapshot(ctx, b.Count) after = b.Hi diff --git a/internal/engine/runner.go b/internal/engine/runner.go index a17b2d7..5941513 100644 --- a/internal/engine/runner.go +++ b/internal/engine/runner.go @@ -40,6 +40,13 @@ type Runner struct { ceiling atomic.Int64 // live ceiling (rows/sec) for adaptive mode emit ProgressFunc + // Running counters mirror the control store so per-batch progress emits + // don't need an extra SELECT on every batch (millions of batches at scale). + total int64 + processed int64 + affected int64 + cursor string + batchIdx int // crashHook, when set, is called after a batch's operation commits on the // target but before its checkpoint is written — a test-only injection point @@ -122,15 +129,18 @@ func (r *Runner) Run(ctx context.Context) { func (r *Runner) run(ctx context.Context) error { dialect := batch.DialectFor(r.task.TargetDriver) - target, err := sql.Open(dialect.DriverName(), r.task.TargetDSN) + target, err := openTarget(dialect.DriverName(), r.task.TargetDSN) if err != nil { return fmt.Errorf("open target database: %w", err) } defer target.Close() - if err := target.PingContext(ctx); err != nil { + pingCtx, cancelPing := withTimeout(ctx, 10*time.Second) + if err := target.PingContext(pingCtx); err != nil { + cancelPing() return fmt.Errorf("connect to target database: %w", err) } + cancelPing() src := batch.Source{ Table: r.task.SourceTable, @@ -146,15 +156,26 @@ func (r *Runner) run(ctx context.Context) error { if run.RowsTotal < 0 { var total int64 - if err := target.QueryRowContext(ctx, batch.CountQuery(src)).Scan(&total); err != nil { + countCtx, cancelCount := withTimeout(ctx, countQueryTimeout) + err := target.QueryRowContext(countCtx, batch.CountQuery(src)).Scan(&total) + cancelCount() + if err != nil { return fmt.Errorf("count total rows: %w", err) } if err := r.runs.SetTotal(ctx, r.db, r.runID, total); err != nil { return err } + + run.RowsTotal = total } + // Seed in-memory counters from the run (handles resume where progress != 0). + r.total = run.RowsTotal + r.processed = run.RowsProcessed + r.affected = run.RowsAffected + r.cursor = run.LastCursor + // Resume point: the run row's cursor (kept in lockstep with checkpoints). var after any if run.LastCursor != "" { @@ -195,16 +216,21 @@ func (r *Runner) run(ctx context.Context) error { return err } - b, err := batch.NextBounds(ctx, target, src, after, r.task.BatchSize) + batchCtx, cancelBatch := withTimeout(ctx, batchQueryTimeout) + + b, err := batch.NextBounds(batchCtx, target, src, after, r.task.BatchSize) if err != nil { + cancelBatch() return err } if b.Count == 0 { + cancelBatch() return r.succeed() } - affected, err := applyBatch(ctx, target, r.task, b, r.quarantine, r.db, r.runID, r.log) + affected, err := applyBatch(batchCtx, target, r.task, b, r.quarantine, r.db, r.runID, r.log) + cancelBatch() if err != nil { return err } @@ -233,10 +259,33 @@ func (r *Runner) run(ctx context.Context) error { after = b.Hi - r.emitSnapshot(ctx, b.Count) + // Update in-memory counters and emit from them — no extra SELECT. + r.processed += int64(b.Count) + r.affected += affected + r.cursor = cursor + r.emitProgress(string(models.RunRunning), b.Count) } } +// emitProgress publishes a live update from the in-memory counters (hot path, +// no DB read). +func (r *Runner) emitProgress(state string, batchRows int) { + if r.emit == nil { + return + } + + r.emit(Progress{ + RunID: r.runID, + State: state, + Processed: r.processed, + Affected: r.affected, + Total: r.total, + Percent: percent(r.processed, r.total), + BatchRows: batchRows, + Cursor: r.cursor, + }) +} + // 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. @@ -303,6 +352,19 @@ func (r *Runner) succeed() error { r.log.Errorf("run %d: audit write failed: %v", r.runID, err) } + // Surface a coverage gap: keyset pagination skips rows whose cursor is NULL + // (they never satisfy cursor > x). If we processed materially fewer rows than + // the initial count, warn — the usual cause is a NULLable cursor column. + if run, err := r.runs.Get(ctx, r.db, r.runID); err == nil && run.RowsTotal > 0 { + if gap := run.RowsTotal - run.RowsProcessed; gap > 0 { + detail := fmt.Sprintf("%d of %d rows were not walked — the cursor column %q may contain NULLs (keyset pagination skips NULL cursors)", + gap, run.RowsTotal, r.task.CursorColumn) + if err := r.events.Append(ctx, r.db, r.runID, "runner", "run.coverage_gap", detail); err != nil { + r.log.Errorf("run %d: audit write failed: %v", r.runID, err) + } + } + } + r.emitSnapshot(ctx, 0) return nil diff --git a/internal/engine/target.go b/internal/engine/target.go new file mode 100644 index 0000000..7d1d85f --- /dev/null +++ b/internal/engine/target.go @@ -0,0 +1,50 @@ +package engine + +import ( + "context" + "database/sql" + "time" +) + +// Target-connection safety limits. A single runner/worker processes batches +// sequentially, so a small pool is plenty; capping it prevents a fleet of +// workers from exhausting the customer database's max_connections, and a +// lifetime bound recycles connections that go stale on multi-hour runs. +const ( + targetMaxOpenConns = 4 + targetMaxIdleConns = 2 + targetConnMaxLifetime = 30 * time.Minute + + // batchQueryTimeout bounds any single fetch/operation/probe against the + // target so a hung query can't pin a run in "running" forever. Generous + // enough for a large batch UPDATE, short enough to fail a wedged DB. + batchQueryTimeout = 10 * time.Minute + + // countQueryTimeout bounds the one-time SELECT count(*) at run start, which + // on a huge table can be a long full scan. + countQueryTimeout = 5 * time.Minute +) + +// OpenTarget opens the customer database with conservative pool limits, for +// callers outside this package (orchestrator, dry-run). +func OpenTarget(driverName, dsn string) (*sql.DB, error) { return openTarget(driverName, dsn) } + +// openTarget opens the customer database with conservative pool limits. +func openTarget(driverName, dsn string) (*sql.DB, error) { + db, err := sql.Open(driverName, dsn) + if err != nil { + return nil, err + } + + db.SetMaxOpenConns(targetMaxOpenConns) + db.SetMaxIdleConns(targetMaxIdleConns) + db.SetConnMaxLifetime(targetConnMaxLifetime) + + return db, nil +} + +// withTimeout derives a child context bounded by d, unless the parent is already +// sooner. Callers must call the returned cancel. +func withTimeout(ctx context.Context, d time.Duration) (context.Context, context.CancelFunc) { + return context.WithTimeout(ctx, d) +} diff --git a/internal/fleet/worker.go b/internal/fleet/worker.go index ed37dee..089011a 100644 --- a/internal/fleet/worker.go +++ b/internal/fleet/worker.go @@ -6,6 +6,7 @@ package fleet import ( "context" "database/sql" + "fmt" "sync" "time" @@ -81,6 +82,15 @@ func (m *Manager) tick(ctx context.Context) error { func (m *Manager) work(ctx context.Context, run models.Run) { defer m.release(run.ID) + // A panic in one lease must not crash the worker process (and abandon every + // other run this process is serving). Recover and record the error. + defer func() { + if p := recover(); p != nil { + m.log.Errorf("fleet worker %s: run %d panic recovered: %v", m.id, run.ID, p) + _ = m.runs.SetError(ctx, m.db, run.ID, fmt.Sprintf("worker panic: %v", p)) + } + }() + task, err := m.tasks.Get(ctx, m.db, run.TaskID) if err != nil { m.log.Errorf("fleet worker %s: load task for run %d: %v", m.id, run.ID, err) diff --git a/internal/handler/apierr/apierr.go b/internal/handler/apierr/apierr.go new file mode 100644 index 0000000..66d7d0c --- /dev/null +++ b/internal/handler/apierr/apierr.go @@ -0,0 +1,49 @@ +// Package apierr maps internal store/service errors to GoFr HTTP errors so the +// API returns correct status codes (404/409/400) instead of a blanket 500 that +// leaks database internals. +package apierr + +import ( + "database/sql" + "errors" + "strings" + + gofrhttp "gofr.dev/pkg/gofr/http" + + "marathon/internal/models" + "marathon/internal/store" +) + +// Map converts err to the appropriate HTTP error. It returns err unchanged when +// no mapping applies (GoFr then treats it as 500). +func Map(err error, entity, id string) error { + if err == nil { + return nil + } + + switch { + case errors.Is(err, sql.ErrNoRows), + errors.Is(err, store.ErrConnectionNotFound): + return gofrhttp.ErrorEntityNotFound{Name: entity, Value: id} + + case errors.Is(err, store.ErrActiveRunExists), + errors.Is(err, store.ErrConnectionInUse), + isUniqueViolation(err): + return gofrhttp.ErrorEntityAlreadyExist{} + + case errors.Is(err, models.ErrInvalidTransition): + return gofrhttp.ErrorInvalidParam{Params: []string{err.Error()}} + + default: + return err + } +} + +// isUniqueViolation reports whether err is a Postgres/MySQL duplicate-key +// error, matched by SQLSTATE / message shape (portable across drivers). +func isUniqueViolation(err error) bool { + s := err.Error() + return strings.Contains(s, "23505") || // Postgres unique_violation + strings.Contains(s, "duplicate key value") || + strings.Contains(s, "Duplicate entry") // MySQL +} diff --git a/internal/handler/connection/handler.go b/internal/handler/connection/handler.go index 07e9efa..8bf2475 100644 --- a/internal/handler/connection/handler.go +++ b/internal/handler/connection/handler.go @@ -13,6 +13,7 @@ import ( gofrhttp "gofr.dev/pkg/gofr/http" "marathon/internal/engine/batch" + "marathon/internal/handler/apierr" "marathon/internal/store" ) @@ -41,7 +42,12 @@ func (h *Handler) Create(ctx *gofr.Context) (any, error) { return nil, gofrhttp.ErrorInvalidParam{Params: []string{"name, driver (postgres|mysql), dsn"}} } - return h.conns.Create(ctx, ctx.SQL, req.Name, req.Driver, req.DSN) + conn, err := h.conns.Create(ctx, ctx.SQL, req.Name, req.Driver, req.DSN) + if err != nil { + return nil, apierr.Map(err, "connection", req.Name) + } + + return conn, nil } // List handles GET /connections. @@ -57,7 +63,7 @@ func (h *Handler) Delete(ctx *gofr.Context) (any, error) { } if err := h.conns.Delete(ctx, ctx.SQL, id); err != nil { - return nil, err + return nil, apierr.Map(err, "connection", ctx.PathParam("id")) } return map[string]bool{"deleted": true}, nil @@ -72,7 +78,7 @@ func (h *Handler) Test(ctx *gofr.Context) (any, error) { driver, dsn, err := h.conns.Resolve(ctx, ctx.SQL, id) if err != nil { - return nil, err + return nil, apierr.Map(err, "connection", ctx.PathParam("id")) } return probe(batch.DialectFor(driver).DriverName(), dsn), nil diff --git a/internal/handler/run/handler.go b/internal/handler/run/handler.go index 35cc836..cd53c41 100644 --- a/internal/handler/run/handler.go +++ b/internal/handler/run/handler.go @@ -2,11 +2,13 @@ package run import ( + "errors" "strconv" "gofr.dev/pkg/gofr" gofrhttp "gofr.dev/pkg/gofr/http" + "marathon/internal/handler/apierr" "marathon/internal/models" "marathon/internal/service/orchestrator" "marathon/internal/store" @@ -71,7 +73,11 @@ func (h *Handler) Start(ctx *gofr.Context) (any, error) { run, err := h.orch.StartRun(ctx, taskID) if err != nil { - return nil, err + if errors.Is(err, orchestrator.ErrTaskBusy) { + return nil, gofrhttp.ErrorEntityAlreadyExist{} + } + + return nil, apierr.Map(err, "task", ctx.PathParam("id")) } return toStatus(run), nil @@ -86,7 +92,7 @@ func (h *Handler) Get(ctx *gofr.Context) (any, error) { run, err := h.runs.Get(ctx, ctx.SQL, id) if err != nil { - return nil, err + return nil, apierr.Map(err, "run", ctx.PathParam("id")) } return toStatus(run), nil @@ -117,7 +123,7 @@ func (h *Handler) Control(ctx *gofr.Context) (any, error) { run, err := h.orch.SetSpeed(ctx, id, *req.Rate) if err != nil { - return nil, err + return nil, apierr.Map(err, "run", ctx.PathParam("id")) } return toStatus(run), nil @@ -125,7 +131,7 @@ func (h *Handler) Control(ctx *gofr.Context) (any, error) { run, err := h.orch.Control(ctx, id, req.Action) if err != nil { - return nil, err + return nil, apierr.Map(err, "run", ctx.PathParam("id")) } return toStatus(run), nil @@ -159,5 +165,10 @@ func (h *Handler) RetryQuarantine(ctx *gofr.Context) (any, error) { return nil, gofrhttp.ErrorInvalidParam{Params: []string{"id"}} } - return h.orch.RetryQuarantine(ctx, id) + res, err := h.orch.RetryQuarantine(ctx, id) + if err != nil { + return nil, apierr.Map(err, "run", ctx.PathParam("id")) + } + + return res, nil } diff --git a/internal/handler/task/handler.go b/internal/handler/task/handler.go index 29a0560..0e8dea5 100644 --- a/internal/handler/task/handler.go +++ b/internal/handler/task/handler.go @@ -8,6 +8,7 @@ import ( "gofr.dev/pkg/gofr" gofrhttp "gofr.dev/pkg/gofr/http" + "marathon/internal/handler/apierr" "marathon/internal/models" "marathon/internal/service/dryrun" "marathon/internal/store" @@ -33,7 +34,7 @@ func (h *Handler) Create(ctx *gofr.Context) (any, error) { } if err := h.tasks.Create(ctx, ctx.SQL, &t); err != nil { - return nil, err + return nil, apierr.Map(err, "task", t.Name) } return t, nil @@ -49,7 +50,12 @@ func (h *Handler) Get(ctx *gofr.Context) (any, error) { return nil, gofrhttp.ErrorInvalidParam{Params: []string{"id"}} } - return h.tasks.Get(ctx, ctx.SQL, id) + task, err := h.tasks.Get(ctx, ctx.SQL, id) + if err != nil { + return nil, apierr.Map(err, "task", ctx.PathParam("id")) + } + + return task, nil } // DryRun analyzes a task without writing to the target (POST /tasks/{id}/dry-run). @@ -61,11 +67,11 @@ func (h *Handler) DryRun(ctx *gofr.Context) (any, error) { task, err := h.tasks.Get(ctx, ctx.SQL, id) if err != nil { - return nil, err + return nil, apierr.Map(err, "task", ctx.PathParam("id")) } if err := (store.ConnectionStore{}).ResolveInto(ctx, ctx.SQL, &task); err != nil { - return nil, err + return nil, apierr.Map(err, "connection", "") } return dryrun.Analyze(ctx, task) diff --git a/internal/service/dryrun/service.go b/internal/service/dryrun/service.go index f2c71e9..abc0804 100644 --- a/internal/service/dryrun/service.go +++ b/internal/service/dryrun/service.go @@ -8,9 +8,9 @@ import ( "context" "database/sql" "fmt" + "time" - _ "github.com/jackc/pgx/v5/stdlib" - + "marathon/internal/engine" "marathon/internal/engine/batch" "marathon/internal/models" ) @@ -36,20 +36,28 @@ type Sample struct { func Analyze(ctx context.Context, task models.Task) (Report, error) { dialect := batch.DialectFor(task.TargetDriver) - db, err := sql.Open(dialect.DriverName(), task.TargetDSN) + db, err := engine.OpenTarget(dialect.DriverName(), task.TargetDSN) if err != nil { return Report{}, fmt.Errorf("open target: %w", err) } defer db.Close() - if err := db.PingContext(ctx); err != nil { + pingCtx, cancelPing := context.WithTimeout(ctx, 10*time.Second) + err = db.PingContext(pingCtx) + cancelPing() + if err != nil { return Report{}, fmt.Errorf("connect to target: %w", err) } src := batch.Source{Table: task.SourceTable, CursorColumn: task.CursorColumn, Filter: task.RowFilter, Dialect: dialect} var rep Report - if err := db.QueryRowContext(ctx, batch.CountQuery(src)).Scan(&rep.PlannedRows); err != nil { + + // Bound the count so dry-run on a huge table can't hang the request. + countCtx, cancelCount := context.WithTimeout(ctx, 5*time.Minute) + err = db.QueryRowContext(countCtx, batch.CountQuery(src)).Scan(&rep.PlannedRows) + cancelCount() + if err != nil { return Report{}, fmt.Errorf("count planned rows: %w", err) } diff --git a/internal/service/orchestrator/service.go b/internal/service/orchestrator/service.go index 7f6c593..9934aed 100644 --- a/internal/service/orchestrator/service.go +++ b/internal/service/orchestrator/service.go @@ -5,13 +5,11 @@ package orchestrator import ( "context" - "database/sql" "errors" "fmt" "sync" "time" - _ "github.com/jackc/pgx/v5/stdlib" "github.com/redis/go-redis/v9" "gofr.dev/pkg/gofr" "gofr.dev/pkg/gofr/container" @@ -33,6 +31,11 @@ const leaseStrideBatches = 5 // heartbeat interval by a comfortable margin to avoid reaping live workers. const deadWorkerCutoff = 10 * time.Second +// stallStrikes is how many consecutive monitor ticks a fleet run may make zero +// progress while leases keep failing before the run is marked failed. At the +// 5s fleet-monitor cadence this is ~15s of "failing with no progress". +const stallStrikes = 3 + var ( ErrTaskBusy = errors.New("task already has an active run") ErrUnknownAction = errors.New("unknown action (want pause, resume, or kill)") @@ -45,6 +48,14 @@ type activeRun struct { runner *engine.Runner } +// stallState tracks a fleet run's progress across monitor ticks for stall +// detection (see detectStall). +type stallState struct { + processed int64 + fails int64 + strikes int +} + type Orchestrator struct { mu sync.Mutex active map[int64]*activeRun // runID → live runner @@ -52,6 +63,8 @@ type Orchestrator struct { hub *stream.Hub rdb *redis.Client // non-nil enables fleet (lease) mode + stall map[int64]stallState // runID → stall tracking (fleet monitor only) + tasks store.TaskStore runs store.RunStore events store.EventStore @@ -62,6 +75,7 @@ type Orchestrator struct { func New(hub *stream.Hub) *Orchestrator { return &Orchestrator{ active: make(map[int64]*activeRun), + stall: make(map[int64]stallState), hub: hub, } } @@ -90,6 +104,9 @@ func (o *Orchestrator) StartRun(ctx *gofr.Context, taskID int64) (models.Run, er return models.Run{}, err } + // Fast, friendly pre-check. The DB partial unique index is the real + // backstop: concurrent starts that both pass this check still can't both + // insert an active run. busy, err := o.runs.ActiveForTask(ctx, ctx.SQL, taskID) if err != nil { return models.Run{}, err @@ -101,6 +118,10 @@ func (o *Orchestrator) StartRun(ctx *gofr.Context, taskID int64) (models.Run, er run, err := o.runs.Create(ctx, ctx.SQL, taskID) if err != nil { + if errors.Is(err, store.ErrActiveRunExists) { + return models.Run{}, fmt.Errorf("%w: task %d", ErrTaskBusy, taskID) + } + return models.Run{}, err } @@ -127,7 +148,7 @@ func (o *Orchestrator) StartRun(ctx *gofr.Context, taskID int64) (models.Run, er func (o *Orchestrator) planFleetRun(ctx *gofr.Context, task models.Task, runID int64) error { dialect := batch.DialectFor(task.TargetDriver) - target, err := sql.Open(dialect.DriverName(), task.TargetDSN) + target, err := engine.OpenTarget(dialect.DriverName(), task.TargetDSN) if err != nil { return fmt.Errorf("open target: %w", err) } @@ -135,8 +156,12 @@ func (o *Orchestrator) planFleetRun(ctx *gofr.Context, task models.Task, runID i src := batch.Source{Table: task.SourceTable, CursorColumn: task.CursorColumn, Filter: task.RowFilter, Dialect: dialect} + // Bound the count and min/max scans so a huge table can't hang StartRun. + planCtx, cancelPlan := context.WithTimeout(ctx, 5*time.Minute) + defer cancelPlan() + var total int64 - if err := target.QueryRowContext(ctx, batch.CountQuery(src)).Scan(&total); err != nil { + if err := target.QueryRowContext(planCtx, batch.CountQuery(src)).Scan(&total); err != nil { return fmt.Errorf("count rows: %w", err) } @@ -151,7 +176,7 @@ func (o *Orchestrator) planFleetRun(ctx *gofr.Context, task models.Task, runID i var minID, maxID int64 - if err := target.QueryRowContext(ctx, batch.MinMaxQuery(src)).Scan(&minID, &maxID); err != nil { + if err := target.QueryRowContext(planCtx, batch.MinMaxQuery(src)).Scan(&minID, &maxID); err != nil { return fmt.Errorf("compute key bounds (fleet mode needs an integer cursor): %w", err) } @@ -211,8 +236,58 @@ func (o *Orchestrator) MonitorFleet(ctx *gofr.Context) { if err := o.runs.SetState(ctx, ctx.SQL, run.ID, models.RunSucceeded, ""); err == nil { _ = o.events.Append(ctx, ctx.SQL, run.ID, "monitor", "run.succeeded", "") _ = queue.Cleanup(ctx) + delete(o.stall, run.ID) ctx.Logger.Infof("fleet monitor: run %d succeeded", run.ID) } + + continue + } + + o.detectStall(ctx, run, queue) + } +} + +// detectStall fails a fleet run whose leases are erroring with no forward +// progress — the signature of a persistently unreachable/broken target. It +// compares rows_processed and the lease-failure counter across monitor ticks: +// a healthy run (even a heavily throttled one) advances rows_processed, so only +// a run that is failing leases while making zero progress trips this. +func (o *Orchestrator) detectStall(ctx *gofr.Context, run models.Run, queue *coord.Queue) { + fails, err := queue.Fails(ctx) + if err != nil { + return + } + + prev, seen := o.stall[run.ID] + cur := stallState{processed: run.RowsProcessed, fails: fails} + + // Progress resets the stall clock. + if !seen || run.RowsProcessed != prev.processed { + cur.strikes = 0 + o.stall[run.ID] = cur + return + } + + // No progress since last tick. If failures are also climbing, count a strike. + if fails > prev.fails { + cur.strikes = prev.strikes + 1 + } else { + cur.strikes = prev.strikes + } + + o.stall[run.ID] = cur + + if cur.strikes >= stallStrikes { + msg := run.Error + if msg == "" { + msg = "run stalled: leases failing with no progress" + } + + if err := o.runs.SetState(ctx, ctx.SQL, run.ID, models.RunFailed, msg); err == nil { + _ = o.events.Append(ctx, ctx.SQL, run.ID, "monitor", "run.failed", "stalled: "+msg) + _ = queue.Cleanup(ctx) + delete(o.stall, run.ID) + ctx.Logger.Errorf("fleet monitor: run %d failed (stalled): %s", run.ID, msg) } } } @@ -325,7 +400,7 @@ func (o *Orchestrator) RetryQuarantine(ctx *gofr.Context, runID int64) (RetryRes return RetryResult{}, err } - target, err := sql.Open("pgx", task.TargetDSN) + target, err := engine.OpenTarget(batch.DialectFor(task.TargetDriver).DriverName(), task.TargetDSN) if err != nil { return RetryResult{}, fmt.Errorf("open target: %w", err) } @@ -440,12 +515,23 @@ func (o *Orchestrator) spawn(c *container.Container, task models.Task, runID int o.active[runID] = &activeRun{cancel: cancel, runner: runner} + log := c.Logger + runs := o.runs + db := c.SQL + go func() { defer func() { o.mu.Lock() delete(o.active, runID) o.mu.Unlock() cancel() + + // A panic in a driver/scan must not crash the process and take down + // every other live run. Recover, mark this run failed, keep serving. + if p := recover(); p != nil { + log.Errorf("run %d: panic recovered: %v", runID, p) + _ = runs.SetState(context.Background(), db, runID, models.RunFailed, fmt.Sprintf("panic: %v", p)) + } }() runner.Run(runCtx) diff --git a/internal/store/run.go b/internal/store/run.go index adeac93..33e10a0 100644 --- a/internal/store/run.go +++ b/internal/store/run.go @@ -3,12 +3,18 @@ package store import ( "context" "database/sql" + "errors" "fmt" "strings" "marathon/internal/models" ) +// ErrActiveRunExists is returned by Create when the task already has an active +// run — the partial unique index runs_one_active_per_task rejected the insert. +// This is the atomic backstop for the check-then-insert race. +var ErrActiveRunExists = errors.New("task already has an active run") + type RunStore struct{} const runColumns = `id, task_id, state, rows_total, rows_processed, rows_affected, @@ -46,6 +52,12 @@ func (RunStore) Create(ctx context.Context, db DB, taskID int64) (models.Run, er r, err := scanRun(row) if err != nil { + // The partial unique index fires when another active run exists. Detect + // it by the index name (portable across pq/pgx error shapes). + if strings.Contains(err.Error(), "runs_one_active_per_task") { + return models.Run{}, ErrActiveRunExists + } + return models.Run{}, fmt.Errorf("create run: %w", err) } @@ -122,6 +134,17 @@ func (RunStore) SetState(ctx context.Context, db DB, id int64, to models.RunStat return nil } +// SetError records the latest error on a run without changing its state (used +// by fleet workers so the monitor can surface why a run is failing). +func (RunStore) SetError(ctx context.Context, db DB, id int64, msg string) error { + _, err := db.ExecContext(ctx, `UPDATE runs SET error = $1 WHERE id = $2`, msg, id) + if err != nil { + return fmt.Errorf("set run %d error: %w", id, err) + } + + return nil +} + func (RunStore) SetTotal(ctx context.Context, db DB, id, total int64) error { _, err := db.ExecContext(ctx, `UPDATE runs SET rows_total = $1 WHERE id = $2`, total, id) if err != nil { diff --git a/internal/stream/hub.go b/internal/stream/hub.go index d2143f3..629f367 100644 --- a/internal/stream/hub.go +++ b/internal/stream/hub.go @@ -9,54 +9,79 @@ import ( "marathon/internal/engine" ) +// deadSubscriberMisses is how many consecutive full-buffer publishes mark a +// subscriber as dead. GoFr never tells us when a browser tab closes, so a +// subscriber whose buffer stays full (nobody draining) is our signal that the +// client is gone; we then drop it to prevent an unbounded subscription leak. +// A live client drains within milliseconds, resetting its miss counter. +const deadSubscriberMisses = 64 + +type subscriber struct { + ch chan engine.Progress + misses int +} + // Hub is an in-process publish/subscribe broker keyed by run ID. type Hub struct { - mu sync.RWMutex - subs map[int64]map[int]chan engine.Progress + mu sync.Mutex + subs map[int64]map[int]*subscriber next int } func NewHub() *Hub { - return &Hub{subs: make(map[int64]map[int]chan engine.Progress)} + return &Hub{subs: make(map[int64]map[int]*subscriber)} } -// Publish delivers a progress update to every subscriber of its run. -// It never blocks: a slow subscriber drops the update rather than stalling -// the runner (the next update supersedes it anyway). +// Publish delivers a progress update to every subscriber of its run. It never +// blocks: a slow subscriber drops the update (the next one supersedes it). A +// subscriber that keeps missing is assumed disconnected and is reaped. func (h *Hub) Publish(p engine.Progress) { - h.mu.RLock() - defer h.mu.RUnlock() + h.mu.Lock() + defer h.mu.Unlock() - for _, ch := range h.subs[p.RunID] { + subs := h.subs[p.RunID] + for id, s := range subs { select { - case ch <- p: + case s.ch <- p: + s.misses = 0 default: + s.misses++ + if s.misses >= deadSubscriberMisses { + close(s.ch) + delete(subs, id) + } } } + + if len(subs) == 0 { + delete(h.subs, p.RunID) + } } // Subscribe returns a channel of updates for a run and an unsubscribe func. +// The unsubscribe func is idempotent and safe to call after the hub has already +// reaped a dead subscriber. func (h *Hub) Subscribe(runID int64) (<-chan engine.Progress, func()) { h.mu.Lock() defer h.mu.Unlock() if h.subs[runID] == nil { - h.subs[runID] = make(map[int]chan engine.Progress) + h.subs[runID] = make(map[int]*subscriber) } id := h.next h.next++ ch := make(chan engine.Progress, 16) - h.subs[runID][id] = ch + h.subs[runID][id] = &subscriber{ch: ch} return ch, func() { h.mu.Lock() defer h.mu.Unlock() if subs := h.subs[runID]; subs != nil { - if c, ok := subs[id]; ok { - close(c) + if s, ok := subs[id]; ok { + close(s.ch) delete(subs, id) } diff --git a/migrations/012_one_active_run.go b/migrations/012_one_active_run.go new file mode 100644 index 0000000..c3efce8 --- /dev/null +++ b/migrations/012_one_active_run.go @@ -0,0 +1,31 @@ +package migrations + +import "gofr.dev/pkg/gofr/migration" + +// Enforces at most one active (queued/running/paused) run per task at the DB +// level. This closes the check-then-insert race in StartRun: two concurrent +// starts (manual + scheduler, or double-click) can no longer both create a run. +func oneActiveRunPerTask() migration.Migrate { + return migration.Migrate{ + UP: func(d migration.Datasource) error { + // Defensive: collapse any pre-existing duplicates before adding the + // unique index, keeping only the newest active run per task. + if _, err := d.SQL.Exec(` + UPDATE runs SET state = 'killed', finished_at = now() + WHERE state IN ('queued','running','paused') + AND id NOT IN ( + SELECT max(id) FROM runs + WHERE state IN ('queued','running','paused') + GROUP BY task_id)`); err != nil { + return err + } + + _, err := d.SQL.Exec(` + CREATE UNIQUE INDEX IF NOT EXISTS runs_one_active_per_task + ON runs (task_id) + WHERE state IN ('queued','running','paused')`) + + return err + }, + } +} diff --git a/migrations/all.go b/migrations/all.go index 0de43ef..eb6402f 100644 --- a/migrations/all.go +++ b/migrations/all.go @@ -17,5 +17,6 @@ func All() map[int64]migration.Migrate { 9: addSchedule(), 10: addTaskConnection(), 11: addAdaptive(), + 12: oneActiveRunPerTask(), } } diff --git a/web/src/api.ts b/web/src/api.ts index f9ce113..67a30c7 100644 --- a/web/src/api.ts +++ b/web/src/api.ts @@ -67,25 +67,59 @@ export const api = { }), }; -// liveSocket opens the progress websocket and performs the {run_id} handshake. +export type LiveStatus = "connecting" | "live" | "reconnecting"; + +// liveSocket opens the progress websocket, performs the {run_id} handshake, and +// transparently reconnects if the connection drops (server restart, network +// blip) with a capped backoff. Returns a close() to stop for good. onStatus +// lets the UI show a "reconnecting" indicator instead of silently freezing. export function liveSocket( runId: number, - onMessage: (p: import("./types").Progress) => void -): WebSocket { + onMessage: (p: import("./types").Progress) => void, + onStatus?: (s: LiveStatus) => void +): () => void { const proto = location.protocol === "https:" ? "wss" : "ws"; - const ws = new WebSocket(`${proto}://${location.host}/runs/${runId}/live`); - ws.onopen = () => ws.send(JSON.stringify({ run_id: runId })); - ws.onmessage = (ev) => { - try { - const p = JSON.parse(ev.data); - // The server can emit a `null` frame (e.g. an empty handshake result); - // only forward well-formed progress objects. - if (p && typeof p === "object" && "run_id" in p) { - onMessage(p as import("./types").Progress); + let ws: WebSocket | null = null; + let closed = false; + let attempt = 0; + let retryTimer: ReturnType | undefined; + + const connect = () => { + if (closed) return; + onStatus?.(attempt === 0 ? "connecting" : "reconnecting"); + ws = new WebSocket(`${proto}://${location.host}/runs/${runId}/live`); + + ws.onopen = () => { + attempt = 0; + onStatus?.("live"); + ws?.send(JSON.stringify({ run_id: runId })); + }; + ws.onmessage = (ev) => { + try { + const p = JSON.parse(ev.data); + if (p && typeof p === "object" && "run_id" in p) { + onMessage(p as import("./types").Progress); + } + } catch { + /* ignore malformed frame */ } - } catch { - /* ignore malformed frame */ - } + }; + ws.onclose = () => { + if (closed) return; + // Reconnect with backoff: 0.5s, 1s, 2s, … capped at 10s. + const delay = Math.min(10000, 500 * 2 ** attempt); + attempt++; + onStatus?.("reconnecting"); + retryTimer = setTimeout(connect, delay); + }; + ws.onerror = () => ws?.close(); + }; + + connect(); + + return () => { + closed = true; + if (retryTimer) clearTimeout(retryTimer); + ws?.close(); }; - return ws; } diff --git a/web/src/pages/RunLive.tsx b/web/src/pages/RunLive.tsx index 3b8b6ac..48f754b 100644 --- a/web/src/pages/RunLive.tsx +++ b/web/src/pages/RunLive.tsx @@ -22,9 +22,13 @@ export function RunLive({ runId, onBack }: { runId: number; onBack: () => void } const [speed, setSpeed] = useState(null); const [speedMax, setSpeedMax] = useState(100000); const [adaptive, setAdaptive] = useState(false); + const [conn, setConn] = useState<"connecting" | "live" | "reconnecting">("connecting"); const last = useRef<{ t: number; processed: number } | null>(null); const speedTimer = useRef | undefined>(undefined); + // Clear a pending slider-debounce timer on unmount (avoid setState-after-unmount). + useEffect(() => () => { if (speedTimer.current) clearTimeout(speedTimer.current); }, []); + // Load the task's current throttle to initialize the slider. useEffect(() => { if (!run) return; @@ -48,24 +52,28 @@ export function RunLive({ runId, onBack }: { runId: number; onBack: () => void } useEffect(() => { api.getRun(runId).then(setRun).catch(() => {}); - const ws = liveSocket(runId, (p: Progress) => { - if (!p) return; - setRun((prev) => ({ ...(prev as Run), ...p, id: p.run_id })); - const now = Date.now(); - if (last.current) { - const dt = (now - last.current.t) / 1000; - if (dt > 0.05) { - const rps = Math.max(0, (p.rows_processed - last.current.processed) / dt); - setRate(rps); - setHistory((h) => [...h.slice(-79), rps]); + const close = liveSocket( + runId, + (p: Progress) => { + if (!p) return; + setRun((prev) => ({ ...(prev as Run), ...p, id: p.run_id })); + const now = Date.now(); + if (last.current) { + const dt = (now - last.current.t) / 1000; + if (dt > 0.05) { + const rps = Math.max(0, (p.rows_processed - last.current.processed) / dt); + setRate(rps); + setHistory((h) => [...h.slice(-79), rps]); + last.current = { t: now, processed: p.rows_processed }; + } + } else { last.current = { t: now, processed: p.rows_processed }; } - } else { - last.current = { t: now, processed: p.rows_processed }; - } - if (["succeeded", "failed", "killed"].includes(p.state)) loadSecondary(); - }); - return () => ws.close(); + if (["succeeded", "failed", "killed"].includes(p.state)) loadSecondary(); + }, + setConn + ); + return () => close(); // eslint-disable-next-line react-hooks/exhaustive-deps }, [runId]); @@ -98,6 +106,14 @@ export function RunLive({ runId, onBack }: { runId: number; onBack: () => void } } }; + if (!Number.isFinite(runId)) { + return ( +
+
Invalid run id.
+ +
+ ); + } if (!run) return
Loading run {runId}…
; const q = Array.isArray(quarantine) ? quarantine : []; const ev = Array.isArray(events) ? events : []; @@ -117,6 +133,9 @@ export function RunLive({ runId, onBack }: { runId: number; onBack: () => void }
Live execution console
+ {active && conn !== "live" && ( + {conn === "connecting" ? "Connecting…" : "Reconnecting…"} + )} diff --git a/web/src/styles.css b/web/src/styles.css index 76f1332..af7aeaf 100644 --- a/web/src/styles.css +++ b/web/src/styles.css @@ -302,6 +302,7 @@ a { color: var(--accent); text-decoration: none; } } .speed input[type="range"]:focus-visible { box-shadow: 0 0 0 3px var(--accent-ring); } .auto-chip { font-size: 9px; font-weight: 800; letter-spacing: 0.6px; padding: 1px 6px; border-radius: 5px; background: var(--good-weak); color: var(--good); margin-left: 6px; vertical-align: 1px; } +.reconnect-chip { font-size: 11px; font-weight: 600; padding: 4px 10px; border-radius: 999px; background: var(--warn-weak); color: var(--warn); border: 1px solid var(--warn-border); } .controls { display: flex; gap: 10px; margin-top: 20px; flex-wrap: wrap; align-items: center; } .done-note { color: var(--good); font-weight: 550; font-size: 13.5px; }