Merge pull request #5 from stackshy/hardening/prod-audit-fixes

Production hardening: fix duplicate-run race, scale cliffs, crash/leak/error-path bugs

Author: NitinKumar004

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"
 }
 

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

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}
 

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

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,18 +117,25 @@ 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))
 	}
 
 	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.

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)