# Task abandonment and the typestate lifecycle

How daisy ensures the run loop terminates even when workers die in pathological ways, and why the per-task counters are protected by a Rust enum rather than runtime guards.

## The two terminal outcomes

Every task ends in exactly one of:

1. **Done** — every block reached a terminal counter (success / permanent failure / orphan) through normal processing.
2. **Abandoned** — the runner gave up after `max_worker_restarts` failed attempts. Remaining blocks are accounted as orphaned (they were never executed) so `is_done()` can become true and the loop can exit.

The whole abandonment subsystem exists to guarantee that `is_done()` becomes true and *stays true* even when in-flight messages from dying workers race against the abandonment transition.

## Counters

Each `RunningTask` (daisy-core/src/task_state.rs) has:

| field                  | meaning                                                            |
| ---------------------- | ------------------------------------------------------------------ |
| `total_block_count`    | invariant for the whole run                                        |
| `ready_count`          | in the ready queue, not yet handed out                             |
| `processing_count`     | dispatched to a worker, awaiting `release_block`                   |
| `completed_count`      | terminal: success                                                  |
| `failed_count`         | terminal: hit `max_retries` during real processing                 |
| `orphaned_count`       | terminal: never executed (upstream failure or task abandonment)    |
| `worker_failure_count` | total dirty exits — pure stat, can exceed `max_worker_restarts`    |
| `worker_restart_count` | refill spawns the runner has performed, capped at the configured limit |

**Conservation** (Invariant 1): `total = ready + processing + completed + failed + orphaned + pending`.

**Done condition** (Invariant 2): `is_done() ⟺ ready = 0 ∧ processing = 0 ∧ pending = 0 ⟺ completed + failed + orphaned = total`.

**Stickiness** (Invariant 3): once a task reaches a terminal state, no operation may move it back to non-terminal. This is the property the typestate enforces structurally.

## The typestate

```rust
pub enum TaskState {
    Running(RunningTask),
    Done(DoneTask),
    Abandoned(AbandonedTask),
}
```

Mutation methods (`note_acquired`, `note_completed`, `note_failed_permanently`, `note_failed_for_retry`, `note_ready`, `note_orphaned`, `note_worker_died`, `note_worker_restarted`) live on `RunningTask`. `DoneTask` and `AbandonedTask` are read-only snapshots of `TaskCounters` plus a reason. The compiler enforces that mutating a task requires pattern-matching to extract `&mut RunningTask` first — terminal variants don't expose those methods at all.

`Scheduler::running_mut(&mut self, task_id)` is the single gate every counter writer goes through:

```rust
fn running_mut(&mut self, task_id: &str) -> Option<&mut RunningTask> {
    self.task_states.get_mut(task_id)?.as_running_mut()
}
```

Late events on terminal tasks return `None` here. The writer drops the message — no need to scatter `is_done()` checks.

## Transitions

- `TaskState::try_finalize_done()` — called after every successful release. If counters balance, transitions `Running → Done`. Idempotent for terminal states.
- `TaskState::abandon(reason)` — consumes the `Running` variant, accounts remaining blocks as orphaned, replaces self with `Abandoned`. Returns the count of orphaned blocks for logging. No-op for terminal states.

Both transitions use `mem::replace` with a placeholder. The HashMap key isn't changed; only the value variant.

## Worker death and restart counting

Two worker shapes, both folded into the same death/restart cycle:

**0-arg `spawn_function`**: the user owns the loop. An exception inside their function unwinds out of `PySpawnWorker::spawn`, which returns `Err`. The Rust thread exits dirty.

**1-arg `process_function`**: the Rust thread owns the loop. After every `client.release_block()`, if the block came back with status `Failed`, the thread exits dirty. This is deliberate: a process-function failure is treated as a worker death, so it bumps `worker_failure_count` and goes through the same restart-cap accounting as a 0-arg crash. The result is that a buggy `process_function` is bounded by `max_worker_restarts + 1` block failures before the runner gives up — the same invariant you get for 0-arg workers.

In both cases, the dirty exit triggers the same downstream flow:

1. `worker_exit_rx` fires on the server's `select!`.
2. `check_thread_health` joins the thread, decrements `allocator.alive`, increments `worker_failure_count` (dirty exits only — it is a reporting stat).
3. `rebalance_workers` decides whether to refill, governed by the **start budget** below.
4. If the budget is exhausted and `alive == 0`, `abandon_exhausted_tasks` transitions the task to `Abandoned` and BFS-propagates the same transition (with reason `UpstreamAbandoned`) to all transitively-downstream tasks.

Because each block-failure is a death, in 1-arg mode the budget also bounds total block failures: a buggy `process_function` will fail at most `max_worker_restarts + 1` blocks before the runner gives up.

## The worker-start budget

A task may ever start at most **`max_workers + max_worker_restarts`** workers, total, regardless of how or why previous workers exited — dirty, clean, or productive alike (`worker_start_count` on `RunningTask`; enforced in `rebalance_workers`). Starts beyond the first `max_workers` are accounted as restarts (`worker_restart_count`), so the abandonment condition and the tqdm `♻=N` counter keep their meaning.

Why no exemptions: daisy workers are expected to be **long-running** — they may hold large models in memory and process many blocks over their lifetime. Recycling workers mid-task (exit after N blocks, rely on respawn) is not a supported pattern; fix the leak, or size `max_worker_restarts` generously. Exempting "good" exits would also reopen an unbounded-respawn hole: a spawn function whose worker silently fails to start (e.g. `subprocess.run(..., check=False)` around a command that can't run on the node) exits cleanly with zero progress and would otherwise respawn forever.

The budget assumes the **blocking-spawn contract**: a spawn function runs
for the lifetime of the worker it starts (`sbatch --wait`, `bsub -K`,
`subprocess.run`). `alive` counts spawn-function threads — daisy has no
other window onto external workers. A submit-and-return spawn function
makes every submission look like an instant clean death: the budget is
consumed by resubmissions while the real jobs (invisible until they
connect) pile up, and the abandonment condition (`no alive workers +
budget exhausted + blocks remaining`) can fire while queued jobs are
still pending. Counting active per-task client connections toward
`alive` would relax this; it is on the roadmap.

Two practical consequences:

- Workers that exit because the queue drained consume no further budget — with no ready blocks, `rebalance_workers` never refills, so over-provisioning `max_workers` past the block count cannot abandon a completed task.
- On preemptible or walltime-limited infrastructure, size `max_worker_restarts` to the number of worker deaths you expect over the run, or rely on done markers and simply start a fresh run to resume.

## Race windows the typestate closes

The motivating problem: when a worker crashes holding a block, three messages travel over independent transports (TCP `ReleaseBlock`, TCP `Disconnect`, mpsc `WorkerStats`) and arrive at the server in any order. Without the typestate gate, the abandonment transition could run while a stale `release_block` is still in the pipe — the late release would push counters past `total_block_count` and flip `is_done()` back to false, hanging the run loop forever.

Four such races, all closed by the gate:

1. **Late `ReleaseBlock` after direct abandonment** — `release_block` early-returns when `as_running_mut()` is `None`.
2. **Downstream ready-block generation after transitive abandonment** — `queue_ready_block` early-returns for non-running destination tasks. Without this, a healthy upstream's release could repopulate `ready_count` on an already-abandoned downstream.
3. **Retry-path requeue after abandonment** — the retry path goes through `release_block` and hits the same gate.
4. **Lost-block release after abandonment** — bookkeeper-driven lost releases also go through `release_block`.

All four are now structural — *unrepresentable* rather than *runtime-checked*.

## Tests

- `tests/test_worker_restarts.py` covers the cycle end-to-end:
  - `test_restart_cap_terminates_task_and_marks_remaining_failed` — direct abandonment with full restart usage
  - `test_task_abandonment_does_not_block_other_tasks` — abandonment of one task doesn't stall an independent peer
  - `test_abandoned_upstream_unblocks_downstream` — transitive orphan propagation
  - `test_abandonment_handles_in_flight_block_release_race` — Race 1 stress test (5 iterations)
  - `test_block_function_failure_kills_worker_and_drives_abandonment` — 1-arg mode hooks into the same cycle
  - `test_healthy_upstream_does_not_repopulate_abandoned_downstream` — Race 2 regression

- `tests/test_keyboard_interrupt.py` — SIGINT during the run aborts cleanly via the same shutdown path.