[prasadlohakpure]

AWOL-on-rollout plan

This PR is Phase 3. Phases 1 and 2 merged in #161.

• Phase 1: make SIGTERM reach Heimdall (exec in entrypoint; ECS stopTimeout 120s in deploy repo)
• Phase 2: drain on shutdown (stop claiming, HTTP shutdown, wait for in-flight jobs and log archiving)
• Phase 3 (this PR): reattach instead of fail (plugin.Resumer for sparkeks and ecs)
• Phase 4: AWOL only after 1 hour without activity (awol_timeout for stale heartbeats and never-claimed jobs)

Problem

Since #161, a rollout waits up to shutdown_timeout for async jobs to finish. Jobs that run longer, and every job on a task that crashes, still lose their heartbeat. After stale_job seconds the janitor fails them with awol: job is stale, marking failed and calls the plugin's Cleanup. For sparkeks and ecs that deletes a SparkApplication or stops ECS tasks that were usually still running fine.

Both plugins name their remote work after the job ID (spark-sql-job-<job_id>, and ECS startedBy heimdall-job-<job_id>-<n>), so another agent can find it and keep waiting on it.

Changes

• plugin.Resumer (pkg/plugin/plugin.go): an optional interface, like ResultStreamer. Resume reattaches to the work Execute started and waits for it as Execute would.

• Janitor (internal/pkg/janitor/job.go, queries/active_jobs_release.sql): a stale job with status running whose plugin implements Resumer is released instead of failed. Its active_jobs row gets agent_name and last_heartbeat set back to null, its status stays running, and Cleanup isn't called. Every other stale or canceling job is handled exactly as before.

• Claim and run (internal/pkg/heimdall/jobs_async.go, job.go): the claim query already takes unclaimed jobs regardless of status, so no SQL change is needed. A claimed job that is already running keeps that status instead of going to accepted, and runJob calls Resume instead of Execute. If its plugin no longer implements Resumer, it fails with the same awol: job is stale, marking failed error as before. When Resume fails for any reason, the worker calls the plugin's Cleanup, since the janitor skipped it on release and nothing else would stop the remote work.

• sparkeks: Resume calls Execute. Before submitting, Execute now looks up spark-sql-job-<job_id>: if it exists, it attaches to it; otherwise it submits it. So a resumed job waits on the application that is still running instead of starting a second one.

• ecs: Resume finds each task by its startedBy and polls it with the same code as Execute (waitForTasks, split out of Execute). It lists tasks with ListTasks by startedBy and desired status, running first, then stopped, and takes the newest. The task definition of that task is reused for restarts, so no new revision is registered. The number of tasks found seeds the retry count, so max_fail_count holds across agents. A task that can't be found fails the resume.

• ecs client token: RunTask now passes heimdall-job-<job_id>-<n>-<attempt> as ClientToken, so a retried call returns the task it already started. Because resume seeds the retry count from the tasks already started, a restart after resume uses a token the earlier agent hasn't used. Resume doesn't rely on the token to find tasks (tasks started before this release have none); token-based resume is a follow-up PR.

• stale_job default 45 s to 3600 s (cmd/heimdall/heimdall.go): a job is now treated as stale only after an hour without a heartbeat. heimdall-config doesn't set stale_job, so this applies in prod. Without it, a brief database or network hiccup longer than 45 s would fail every non-resumable job and release every resumable one.

No drain change is needed. A job still running at the drain deadline is cut off, its heartbeat goes stale, and the janitor releases it like a crashed one.

What happens to a job whose agent dies

Job	Before	After
sparkeks or ecs, status running	Failed as AWOL, remote work deleted or stopped	Released, then claimed by any agent and resumed; remote work untouched
sparkeks or ecs, status accepted (claimed, not started)	Failed as AWOL	Unchanged
Any other plugin	Failed as AWOL, Cleanup	Unchanged; its work died with the agent's connection
Released, then canceled	n/a	The janitor's canceling branch picks it up (no heartbeat), runs Cleanup, marks it canceled

Limits

• Release delay. A job is released stale_job seconds after its last heartbeat (now 1 hour by default), on the janitor's next pass (up to clean_interval, 60 s). The remote work keeps running meanwhile, but a job orphaned by a rollout finishes up to an hour later than its remote work does. Non-resumable jobs on a dead agent show running for up to an hour before they're failed, instead of 45 s.
• Logs. The resuming agent starts a new working directory, so the archived stdout and stderr only cover the resumed part. The old agent's local logs are lost, as before. sparkeks re-collects driver logs from the pods.
• Restarted counters. ecs restarts its timeout on resume. sparkeks restarts its 5-hour monitor timeout and re-uploads the same query.sql.
• ECS task gone. ECS keeps stopped tasks for about an hour. With stale_job at an hour, a task that stopped soon after its agent died is usually gone by the time the job is resumed, and the job fails. If some earlier failed tasks are gone but a later one isn't, the retry count comes out low and the job may restart a task more times than max_fail_count allows.
• Spark job can run twice. If the application is already gone when the job is resumed, Execute submits a new one and the Spark job runs again. That happens when the agent died after deleting a finished application but before writing the final status (sparkeks deletes before fetching results, as before), or when someone deleted it by hand. Jobs that write (INSERT, MERGE) would write twice.
• A released job nobody resumes stays running. Phase 4 adds the awol_timeout ceiling for that.
• Heartbeat outage on a live agent. If an agent is alive but can't write heartbeats for stale_job seconds, its job is released and another agent resumes it while the first is still waiting on it. Today the same outage fails the job and deletes its remote work, so this is no worse. The first agent's final status write can still race with the second's.

Testing

• go build, go vet and go test ./internal/... ./pkg/... pass.
• New TestSplitResumable: only running jobs whose plugin implements Resumer are released. accepted and canceling jobs, plugins without Resumer, and commands no longer loaded go down the old path.
• Locally against Postgres, with stale_job: 15 and a throwaway plugin whose Execute takes 120 s and whose Resume takes 5 s (not committed):
  1. Submitted one job on that plugin and one shell job running sleep 120. Both reached running.
  2. kill -9 the agent, then started a new one.
  3. About 15 s later, the shell job was failed with awol: job is stale, marking failed. In the same janitor pass the resumable job was released, claimed by the new agent under its own agent name, resumed, and finished succeeded with result resumed.
• Locally against Postgres and the real prod ECS cluster and spark-eks, with two agents and stale_job: 60, killing the agent that held the job with kill -9:
  • ECS, task still running: the other agent attached to the same task, then found it stopped after the sleep, fetched its CloudWatch logs and finished succeeded. No second task.
  • ECS, task already stopped: the resume found the stopped task by startedBy and handled its exit as Execute would.
  • ECS, failing task (exit 1, max_fail_count: 3): killed after restart 1. The other agent attached, counted two tasks, restarted with retry=2 (no ConflictException), and failed the job at "failed 3 times (max: 3)".
  • sparkeks: killed right after submit. The other agent logged Attached to existing Spark application and monitored it; no second submission. (The job itself failed on a Ranger permission for the local user, unrelated to resume.)
  • active_jobs: after stale_job the row moved to the other agent with a fresh heartbeat, and was removed when the job finished.

[Copilot]

🟡 Changes recommended

Critical recovery defects can block ECS reattachment, orphan running tasks after failed resume, and prevent recovery of completed ECS or Spark work when remote task records or applications are gone.

4 open issues

• Clean up or retry failed ECS reattachment · New
• Persist per-task completion outcomes for recovery · New
• Use supported ECS filters when reattaching tasks · New
• Preserve Spark results until terminal status is durable · New

What changed in this PR

Implements phase 3 of the rollout recovery plan: another agent can resume stale sparkeks and ecs jobs instead of failing them and stopping their remote work.

Changes:

• Adds an optional plugin.Resumer interface and dispatches reclaimed running jobs to it.
• Releases stale resumable jobs while preserving existing handling for other jobs.
• Adds SparkApplication and ECS task reattachment, shared monitoring paths, and janitor classification tests.

File	Description
pkg/​plugin/​plugin.go	Defines the optional resumption interface.
internal/​pkg/​object/​command/​sparkeks/​sparkeks.go	Reattaches to existing SparkApplications and shares monitoring logic.
internal/​pkg/​object/​command/​ecs/​ecs.go	Discovers existing tasks and adopts their task definition for resumed monitoring.
internal/​pkg/​janitor/​queries/​active_jobs_release.sql	Clears agent ownership and heartbeat for released jobs.
internal/​pkg/​janitor/​job.go	Separates resumable jobs from those requiring cleanup.
internal/​pkg/​janitor/​job_test.go	Tests which jobs qualify for release.
internal/​pkg/​heimdall/​jobs_async.go	Preserves running status and identifies resumed claims.
internal/​pkg/​heimdall/​job.go	Calls Resume and handles plugins that cannot resume.

_{Review effort: Balanced}