Skip to main content

wowlab_sentinel/scheduler/
reclaim.rs

1#![expect(
2    clippy::cast_precision_loss,
3    clippy::cast_sign_loss,
4    reason = "database chunk counters are constrained non-negative and metrics accept f64 samples"
5)]
6
7use async_trait::async_trait;
8use chrono::Utc;
9use uuid::Uuid;
10use wowlab_types::constants::MS_PER_SECOND_U64;
11
12use super::{JobProgressEvent, JobProgressPayload};
13use crate::{cron::CronJob, scheduler::runtime::JobRuntimeStatus, state::ServerState};
14
15const MS_PER_MINUTE: u64 = 60_000;
16
17#[derive(Debug, Default)]
18struct JobReclaimStats {
19    reclaimed: u64,
20    failed: u64,
21}
22
23/// Reclaims stale in-flight chunks from nodes that have gone offline.
24#[derive(Debug)]
25pub(crate) struct ReclaimChunksJob {
26    schedule: String,
27}
28
29impl ReclaimChunksJob {
30    pub(crate) fn new(schedule: &str) -> Self {
31        Self {
32            schedule: schedule.to_string(),
33        }
34    }
35}
36
37#[async_trait]
38impl CronJob for ReclaimChunksJob {
39    fn name(&self) -> &'static str {
40        "reclaim_stale_chunks"
41    }
42
43    fn schedule(&self) -> &str {
44        &self.schedule
45    }
46
47    async fn run(&self, state: &ServerState) {
48        if let Err(e) = do_reclaim(state).await {
49            tracing::error!(error = %e, "Failed to reclaim stale chunks");
50        }
51    }
52}
53
54pub(crate) async fn do_reclaim(
55    state: &ServerState,
56) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
57    let timeout_ms = (state.config.reclaim_timeout_minutes.max(0) as u64) * MS_PER_MINUTE;
58    let max_attempts = state.config.reclaim_max_attempts.max(0) as u32;
59    let now_ms = wowlab_common::time::unix_timestamp_secs().saturating_mul(MS_PER_SECOND_U64);
60
61    let work = {
62        let mut runtimes = state.runtime.jobs.write().await;
63
64        collect_reclaim_work(&mut runtimes, now_ms, timeout_ms, max_attempts)
65    };
66
67    for event in &work.events {
68        write_reclaim_event(state, event).await;
69    }
70
71    let completed_at = Utc::now().to_rfc3339();
72
73    for job_id in work.failed_jobs {
74        mark_job_failed(state, job_id, &completed_at).await;
75    }
76
77    if work.total_reclaimed > 0 {
78        tracing::info!(count = work.total_reclaimed, "Reclaimed stale chunks");
79        metrics::counter!(crate::telemetry::CHUNKS_RECLAIMED).increment(work.total_reclaimed);
80        metrics::gauge!(crate::telemetry::CHUNKS_RUNNING).decrement(work.total_reclaimed as f64);
81    }
82
83    if work.total_failed > 0 {
84        tracing::warn!(
85            count = work.total_failed,
86            "Failed chunks (exceeded max attempts)"
87        );
88        metrics::counter!(crate::telemetry::CHUNKS_FAILED).increment(work.total_failed);
89    }
90
91    Ok(())
92}
93
94type ReclaimEvent = (Uuid, String, Option<String>, Option<String>);
95
96struct ReclaimWork {
97    events: Vec<ReclaimEvent>,
98    failed_jobs: Vec<Uuid>,
99    total_reclaimed: u64,
100    total_failed: u64,
101}
102
103fn collect_reclaim_work(
104    runtimes: &mut crate::scheduler::runtime::JobRuntimeStore,
105    now_ms: u64,
106    timeout_ms: u64,
107    max_attempts: u32,
108) -> ReclaimWork {
109    let mut work = ReclaimWork {
110        events: Vec::new(),
111        failed_jobs: Vec::new(),
112        total_reclaimed: 0,
113        total_failed: 0,
114    };
115
116    for runtime in runtimes.iter_mut() {
117        let outcome = runtime.reclaim_stale(now_ms, timeout_ms, max_attempts);
118
119        if outcome.reclaimed.is_empty() && outcome.failed.is_empty() {
120            continue;
121        }
122
123        let stats = JobReclaimStats {
124            reclaimed: outcome.reclaimed.len() as u64,
125            failed: outcome.failed.len() as u64,
126        };
127
128        for claim in &outcome.failed {
129            work.events
130                .push(reclaim_event(runtime.job_id, claim, max_attempts));
131        }
132
133        if !runtime.has_outstanding_work() && !runtime.is_complete() && runtime.has_failures() {
134            runtime.status = JobRuntimeStatus::Failed;
135            work.failed_jobs.push(runtime.job_id);
136        }
137
138        work.total_reclaimed += stats.reclaimed;
139        work.total_failed += stats.failed;
140    }
141
142    for job_id in &work.failed_jobs {
143        runtimes.remove(job_id);
144    }
145
146    work
147}
148
149fn reclaim_event(
150    job_id: Uuid,
151    claim: &crate::scheduler::runtime::ClaimRecord,
152    max_attempts: u32,
153) -> ReclaimEvent {
154    (
155        job_id,
156        "failure".into(),
157        Some(claim.node_public_key.clone()),
158        Some(format!("exceeded {max_attempts} reclaim attempts")),
159    )
160}
161
162async fn write_reclaim_event(state: &ServerState, event: &ReclaimEvent) {
163    let (job_id, event_type, node_pk, error) = event;
164
165    if let Err(db_error) = sqlx::query_file!(
166        "queries/chunk_event_log_insert.sql",
167        *job_id,
168        event_type,
169        Option::<Uuid>::None,
170        node_pk.as_deref(),
171        error.as_deref(),
172    )
173    .execute(state.dbs.get::<crate::scheduler::SchedulerDb>())
174    .await
175    {
176        tracing::error!(%job_id, error = %db_error, "Failed to write reclaim event log");
177    }
178}
179
180async fn mark_job_failed(state: &ServerState, job_id: Uuid, completed_at: &str) {
181    if let Err(error) = sqlx::query_file!("queries/reclaim_mark_failed.sql", job_id)
182        .execute(state.dbs.get::<crate::scheduler::SchedulerDb>())
183        .await
184    {
185        tracing::error!(%job_id, %error, "Failed to mark job failed");
186
187        return;
188    }
189
190    tracing::warn!(%job_id, "Job marked failed after reclaim");
191    let event = JobProgressEvent {
192        r#type: "updated",
193        payload: JobProgressPayload {
194            status: "failed",
195            chunks_completed: 0,
196            chunks_total: 0,
197            completed_at: Some(completed_at.to_owned()),
198            phase: None,
199            phase_count: None,
200            permutations_active: None,
201            top_mean_dps_x10: None,
202            top_k: Vec::new(),
203        },
204    };
205
206    state.publish(&format!("jobs:{job_id}"), &event).await;
207    state.publish("jobs:all", &event).await;
208}