Fix task executor worker identity reconciliation races

[Andyz26]

Problem

Task ownership spans submission, preparation, execution, cancellation, and cleanup. Previously, the agent did not expose ownership until preparation completed, allowing the control plane and agent to disagree about which worker occupied an executor.

Ambiguous assignment failures could then discard ownership, cancel the wrong worker, or release an executor prematurely. Two additional lifecycle races could leave executors permanently unschedulable:

• Disconnecting while Assigned(A) preserved that state across reconnect, but an Available heartbeat could not transition it back to Pending.
• Stale disconnect, timeout, and status messages could revive an archived executor as active but unregistered. A late assignment failure could then mutate it, throw, and be dropped by the actor.

Summary

• Reserve accepted worker identity before acknowledging submission.
• Keep ambiguous executors fenced until ownership is reconciled.
• Reconcile Assigned disconnects before returning an executor to scheduling.
• Prevent stale lifecycle messages from reviving archived executors.
• Make cancellation idempotent, prompt to acknowledge, and bounded during shutdown.
• Dispose prepared tasks and classloader leases when cancellation races with preparation.

State Transition Analysis

Agent task ownership

flowchart LR
    I[Idle] -->|submit A: reserve before Ack| P[Preparing A]
    P -->|prepared| R[Running A]
    P -->|cancel accepted| C[Cancelling A]
    R -->|cancel accepted| C
    C -->|cleanup settles| I
    P -->|submit B| X[Reject with worker A]
    R -->|submit B| X

Loading

The reserved slot is the source of worker identity for submission checks, reports, cancellation, and cleanup. Cancellation acknowledgement confirms acceptance rather than completion of teardown.

Assignment failure direction

flowchart TB
    F[Assignment attempt] --> G{Did submit begin?}
    G -->|No| N[NotSent]
    N -->|retry remains| R[Retry through fresh gateway]
    N -->|expired| B[Block late submission]
    G -->|Yes or unknown| M[MayHaveRun]
    M --> Q[Fence executor and cancel expected worker]
    F -->|AlreadyRunning B| C[Quarantine and cancel B]
    F -->|stale assignment epoch| S[Ignore failure]

Loading

Issue 1: reconnect while assigned

flowchart LR
    subgraph Before
        A1[Assigned A] --> B1[Disconnect and archive]
        B1 --> C1[Reconnect as Assigned A]
        C1 -->|Available heartbeat| D1[Assigned A forever]
    end

    subgraph After
        A2[Assigned A] --> B2[Disconnect]
        B2 --> C2[Cancelling A and mark A Lost]
        C2 --> D2[Archive and reconnect]
        D2 -->|Occupied A| E2[Quarantine and cancel A]
        D2 -->|Available heartbeat 1| F2[Verifying]
        F2 -->|Available heartbeat 2| G2[Pending]
    end

Loading

Only reservation-capable agents use the two-heartbeat recovery path. An unsequenced status-change RPC cannot bypass reconciliation.

Issue 3: stale lifecycle messages

flowchart LR
    subgraph Before
        A1[Archived executor] --> B1[Stale disconnect or status]
        B1 --> C1[Revived active but unregistered]
        C1 --> D1[Assignment gate passes]
        D1 --> E1[Mutation throws and message is dropped]
    end

    subgraph After
        A2[Archived executor] -->|registration or heartbeat| B2[Revive and reconcile]
        A2 -->|duplicate disconnect or timeout| C2[Ignore]
        A2 -->|late status| D2[Return NotFound]
        A2 -->|late assignment failure| E2[Ignore]
    end

Loading

Archived state is now revived only by messages that establish renewed executor activity. Assignment failures additionally require a registered, current, unreconciled assignment.

Tests

./gradlew :mantis-control-plane:mantis-control-plane-server:test \
  :mantis-server:mantis-server-agent:test

Coverage includes preparation ownership, cancellation cleanup, stop failures, bounded shutdown, conflicting assignments, delayed failures, reconnect reconciliation, duplicate disconnects, and late status messages.

[github-actions]

Test Results

845 tests  +22   834 ✅ +22   10m 11s ⏱️ -7s
167 suites ± 0    11 💤 ± 0 
167 files   ± 0     0 ❌ ± 0 

Results for commit eadc133. ± Comparison against base commit 9ff3102.

♻️ This comment has been updated with latest results.

[hellolittlej]

P0 — must fix before merge

1. TaskExecutorState.java:167 — reconnecting executor becomes permanently unschedulable
   onDisconnection now only clears availabilityState when workerId == null. Since trackIfAbsent (ExecutorStateManagerImpl.java:186) revives the archived state object, an agent that drops off while Assigned(w) comes back with availabilityState == Assigned(w). isAvailable() requires Pending; Assigned.onTaskExecutorStatusChange(Available) returns this; the Running+Available branch at :328 doesn't cover Assigned. Nothing ever moves it back.
   Telling detail: getPreviousWorkerId()/clearPreviousWorkerId() (:464, :468) have zero callers tree-wide — the reconcile-on-reconnect half of this design isn't wired up.
   Fix: on re-registration, reconcile previousWorkerId against the first report and reset availability to Pending when the executor comes back idle.

2. TaskExecutorState.java:280 (with :333) — rolling upgrade drains cluster capacity
   reservesAcceptedTask() defaults to false, i.e. every agent not yet redeployed with this PR. For those, an Available heartbeat while reconciliationState != None hits return false at :287 — no clearReconciliation(), no pending(); and :333 puts a legacy agent straight into Quarantined on its first Running+Available heartbeat. One MayHaveRun assignment failure then removes it from the pool forever. Master self-healed twice over (fell through to Running.onTaskExecutorStatusChange(Available) → pending(), and cleared a stale cancelledWorkerOnTask on mismatch).
   Fix: give legacy agents a recovery path — either clear reconciliation after N consecutive Available heartbeats, or treat !reservesAcceptedTask() as "trust the report" and go straight to Pending.

3. ExecutorStateManagerActor.java:668 — ESM actor restarts on assignment failure
   Master wrapped this block in try { … } catch (IllegalStateException e) { log.error(…) } (master :660-671); the PR deletes it. Meanwhile setCancelledWorkerOnTask/quarantineWorkerOnTask (TaskExecutorState.java:129,137) now throw when unregistered, and the new isCurrentAssignment gate checks isAssigned() && epoch && workerId — not isRegistered(). A TE that disconnects while Assigned(w) and is later revived unregistered passes the gate, and :689/:691 throws out of receive. Akka drops the message and cancels every AbstractActorWithTimers heartbeat timer.
   Fix: add isRegistered() to isCurrentAssignment, and restore the catch.

P1 — should fix

4. ExecutorStateManagerActor.java:964 and :982 — MarkExecutorTaskCancelled/TerminateWorker throw instead of replying
   Both select via isRunningOrAssigned(workerId) (TaskExecutorState.java:417), which compares only getWorkerId() and ignores registration, so they match the same unregistered phantom.setCancelledWorkerOnTask throws before sender().tell(Ack), with no try/catch — the job actor or leader-exclusive REST route gets no reply and blocks to its ask timeout, plus the actor restarts. This was a harmless field assignment before.

5. ExecutorStateManagerActor.java:993 — cancelTaskOnExecutor swallows every failure
   The whenComplete only logs: no retry, no state correction, no disconnect. But the caller has already forced Cancelling/Quarantined + tryMarkUnavailable. So on the most likely failure — TaskNotFoundException because the agent never got the task, exactly the MayHaveRun-after-timeout case — the executor stays unavailable, and per Temporarily disable subprojects. #2 a legacy agent never recovers. The wholereconciliation design depends on this call succeeding, yet its failure is a no-op.

6. ExecutorStateManagerActor.java:686 — foreign worker killed silently
   On TaskAlreadyRunningException, cancellationTarget becomes getCurrentlyRunningWorkerTask() — a worker belonging to a different job/attempt — which gets quarantined and cancelTask-ed. Theonly jobMessageRouter event emitted is WorkerLaunchFailed for expectedWorker (:674); nothing is routed for the worker actually being killed. Its job actor believes it's alive until its own heartbeat timeout expires, delaying replacement by the full window.

[hellolittlej]

Also, what triggers us to make this change?

[hellolittlej]

I gave up reading the whole summary and reviewing whole things by myself lol coz it's too agentic, therefore I delegate the review process to my agent. :)

[hellolittlej]

Do we have any metrics that can validate this change after the merge?