termivar_scanner/distributed/
recovery.rs

1use std::collections::BTreeMap;
2
3use super::coordinator::{next_counter, next_u32, CoordinatorState};
4use super::lease::{terminalize_leased_task, TaskLease};
5use super::model::TaskStatus;
6use super::queue::QueueKey;
7use super::worker::{preflight_worker_release, release_worker};
8use super::DistributedError;
9
10#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
11pub struct RecoverySummary {
12    pub workers_affected: usize,
13    pub tasks_requeued: usize,
14    pub tasks_failed: usize,
15}
16
17pub(super) fn leased_task_ids<F>(state: &CoordinatorState, mut predicate: F) -> Vec<String>
18where
19    F: FnMut(&TaskLease) -> bool,
20{
21    let mut leased: Vec<(u64, String)> = state
22        .tasks
23        .iter()
24        .filter_map(|(task_id, entry)| {
25            entry
26                .task
27                .lease
28                .as_ref()
29                .filter(|lease| predicate(lease))
30                .map(|lease| (lease.lease_id, task_id.clone()))
31        })
32        .collect();
33    leased.sort();
34    leased.into_iter().map(|(_, task_id)| task_id).collect()
35}
36
37pub(super) fn preflight_recovery(
38    state: &CoordinatorState,
39    task_ids: &[String],
40) -> Result<(), DistributedError> {
41    let requeue_count = task_ids
42        .iter()
43        .filter(|task_id| {
44            state
45                .tasks
46                .get(*task_id)
47                .is_some_and(|entry| entry.task.retry_count < state.limits.max_retries)
48        })
49        .count();
50    let count = u64::try_from(requeue_count).map_err(|_| DistributedError::CounterExhausted {
51        counter: "enqueue_ordinal",
52    })?;
53    state
54        .next_enqueue_ordinal
55        .checked_add(count)
56        .ok_or(DistributedError::CounterExhausted {
57            counter: "enqueue_ordinal",
58        })?;
59    if state.queue.len().saturating_add(requeue_count) > state.limits.max_queued_tasks {
60        return Err(DistributedError::QueuedTaskCapacityReached {
61            limit: state.limits.max_queued_tasks,
62        });
63    }
64    let mut release_counts: BTreeMap<&str, u32> = BTreeMap::new();
65    for task_id in task_ids {
66        let entry = state
67            .tasks
68            .get(task_id)
69            .ok_or(DistributedError::StateInvariant {
70                reason: "recovery task is missing",
71            })?;
72        if !matches!(entry.task.status, TaskStatus::Leased | TaskStatus::Running) {
73            return Err(DistributedError::StateInvariant {
74                reason: "recovery selected a non-leased task",
75            });
76        }
77        let lease = entry
78            .task
79            .lease
80            .as_ref()
81            .ok_or(DistributedError::StateInvariant {
82                reason: "recovery task has no lease",
83            })?;
84        preflight_worker_release(state, lease)?;
85        next_counter(entry.task.record_version, "record_version")?;
86        if entry.task.retry_count < state.limits.max_retries {
87            next_counter(entry.task.task_generation, "task_generation")?;
88            next_u32(entry.task.retry_count, "retry_count")?;
89        }
90        let count = release_counts.entry(&lease.worker_id).or_default();
91        *count = next_u32(*count, "worker_release_count")?;
92    }
93    for (worker_id, releases) in release_counts {
94        let worker = state
95            .workers
96            .get(worker_id)
97            .ok_or(DistributedError::StateInvariant {
98                reason: "recovery lease references a missing worker",
99            })?;
100        if releases > worker.current_tasks {
101            return Err(DistributedError::StateInvariant {
102                reason: "recovery releases exceed worker active count",
103            });
104        }
105    }
106    Ok(())
107}
108
109#[derive(Debug, Clone, Copy, PartialEq, Eq)]
110pub(super) enum RecoveryDisposition {
111    Requeued,
112    Failed,
113}
114
115pub(super) fn apply_recovery(
116    state: &mut CoordinatorState,
117    task_id: &str,
118    now_secs: u64,
119) -> Result<RecoveryDisposition, DistributedError> {
120    let task = state
121        .tasks
122        .get(task_id)
123        .map(|entry| entry.task.clone())
124        .ok_or(DistributedError::StateInvariant {
125            reason: "recovery task disappeared",
126        })?;
127    let lease = task
128        .lease
129        .as_ref()
130        .ok_or(DistributedError::StateInvariant {
131            reason: "recovery task lost its lease",
132        })?
133        .clone();
134    let record_version = next_counter(task.record_version, "record_version")?;
135    if task.retry_count < state.limits.max_retries {
136        let task_generation = next_counter(task.task_generation, "task_generation")?;
137        let retry_count = next_u32(task.retry_count, "retry_count")?;
138        let enqueue_ordinal = state.next_enqueue_ordinal;
139        let next_ordinal = next_counter(enqueue_ordinal, "enqueue_ordinal")?;
140        release_worker(state, &lease)?;
141        let entry = state
142            .tasks
143            .get_mut(task_id)
144            .ok_or(DistributedError::StateInvariant {
145                reason: "recovery task disappeared during mutation",
146            })?;
147        entry.task.status = TaskStatus::Queued;
148        entry.task.assigned_to = None;
149        entry.task.lease = None;
150        entry.task.started_at = None;
151        entry.task.retry_count = retry_count;
152        entry.task.task_generation = task_generation;
153        entry.task.record_version = record_version;
154        let key = QueueKey::new(entry.task.priority, enqueue_ordinal, task_id.to_owned());
155        entry.queue_key = Some(key.clone());
156        state.queue.insert(key);
157        state.next_enqueue_ordinal = next_ordinal;
158        Ok(RecoveryDisposition::Requeued)
159    } else {
160        terminalize_leased_task(state, &lease, TaskStatus::Failed, now_secs, record_version)?;
161        Ok(RecoveryDisposition::Failed)
162    }
163}
termivar_scanner/distributed/recovery.rs

termivar_scanner/distributed/
recovery.rs
