1use async_trait::async_trait;
4use serde_json::json;
5use uuid::Uuid;
6use wowlab_common::{
7 ClaimToken, NodePublicKey, RuntimeChunkId, RuntimeWorkResult, WorkContextHash,
8 node_http::NodeChunkCompletionResponse, sim::intent,
9};
10use wowlab_types::{game::SpecId, proto};
11
12use crate::{
13 scheduler::{
14 JobProgressEvent, JobProgressPayload, JobProgressTopPermutation,
15 runtime::{CompletionOutcome, JobRuntimeStatus, ProgressSnapshot},
16 },
17 state::{RuntimeState, ServerState},
18 strategy::{FinalOutput, StrategyError},
19};
20
21const WORK_CONTEXT_HASH_LEN: usize = 32;
22
23#[derive(Debug, thiserror::Error)]
24pub(in crate::http) enum CompletionError {
25 #[error("invalid job_id")]
26 InvalidJobId(#[source] uuid::Error),
27 #[error("invalid work_context_hash length")]
28 InvalidWorkContextHashLength {
29 actual: usize,
30 #[source]
31 source: std::array::TryFromSliceError,
32 },
33 #[error("job not found or not running")]
34 JobNotFound,
35 #[error("completion ingest failed")]
36 Ingest(#[source] StrategyError),
37 #[error("{0}")]
38 Conflict(String),
39 #[error("{0}")]
40 Stale(String),
41 #[error("finalize failed")]
42 Finalize(#[source] StrategyError),
43 #[error("{0}")]
44 Persistence(
45 #[from]
46 #[source]
47 CompletionPersistenceError,
48 ),
49}
50
51impl From<uuid::Error> for CompletionError {
52 fn from(error: uuid::Error) -> Self {
53 Self::InvalidJobId(error)
54 }
55}
56
57#[derive(Debug, thiserror::Error)]
58pub(in crate::http) enum CompletionPersistenceError {
59 #[error("failed to begin completion transaction")]
60 Begin(#[source] sqlx::Error),
61 #[error("failed to save result")]
62 Save(#[source] sqlx::Error),
63 #[error("failed to commit completion transaction")]
64 Commit(#[source] sqlx::Error),
65}
66
67pub(in crate::http) struct CompletionInput {
68 job_id: Uuid,
69 chunk_id: RuntimeChunkId,
70 claim_token: ClaimToken,
71 work_context_hash: WorkContextHash,
72 node_public_key: String,
73 results: Vec<RuntimeWorkResult>,
74}
75
76impl CompletionInput {
77 pub(in crate::http) fn parse(
78 node_public_key: &NodePublicKey,
79 completion: &proto::BatchChunkCompletion,
80 ) -> Result<Self, CompletionError> {
81 let job_id = completion.job_id.parse()?;
82 let work_context_hash = parse_work_context_hash(&completion.work_context_hash)?;
83
84 Ok(Self {
85 job_id,
86 chunk_id: RuntimeChunkId::new(completion.chunk_id),
87 claim_token: ClaimToken::new(completion.claim_token.clone()),
88 work_context_hash,
89 node_public_key: node_public_key.to_base64(),
90 results: completion
91 .results
92 .iter()
93 .map(RuntimeWorkResult::from_proto)
94 .collect(),
95 })
96 }
97}
98
99fn parse_work_context_hash(bytes: &[u8]) -> Result<WorkContextHash, CompletionError> {
100 let raw: [u8; WORK_CONTEXT_HASH_LEN] =
101 bytes
102 .try_into()
103 .map_err(|source| CompletionError::InvalidWorkContextHashLength {
104 actual: bytes.len(),
105 source,
106 })?;
107
108 Ok(WorkContextHash::from_bytes(raw))
109}
110
111pub(in crate::http) struct FinalizedJob {
112 job_id: Uuid,
113 sim_config: String,
114 output: FinalOutput,
115}
116
117pub(in crate::http) struct ChunkEvent<'a> {
118 job_id: Uuid,
119 event_type: &'static str,
120 claim_token: &'a ClaimToken,
121 error: &'a str,
122}
123
124#[async_trait]
125pub(in crate::http) trait CompletionRepository: Send + Sync {
126 async fn persist(&self, job: FinalizedJob) -> Result<(), CompletionPersistenceError>;
127 async fn record_event(&self, event: ChunkEvent<'_>) -> Result<(), sqlx::Error>;
128}
129
130#[async_trait]
131pub(in crate::http) trait CompletionObserver: Send + Sync {
132 async fn accepted(
133 &self,
134 job_id: Uuid,
135 snapshot: Option<&ProgressSnapshot>,
136 node_public_key: &str,
137 );
138 async fn finalized(&self, job_id: Uuid, strategy: &str, node_public_key: &str);
139}
140
141pub(in crate::http) struct SqlCompletionRepository<'a> {
142 pool: &'a sqlx::PgPool,
143}
144
145impl SqlCompletionRepository<'_> {
146 async fn persist_transaction(
147 &self,
148 job: FinalizedJob,
149 ) -> Result<(), CompletionPersistenceError> {
150 let mut tx = match self.pool.begin().await {
151 Ok(tx) => tx,
152 Err(source) => {
153 tracing::error!(error = %source, "Failed to begin transaction");
154
155 return Err(CompletionPersistenceError::Begin(source));
156 }
157 };
158 let mut meta_map = serde_json::Map::new();
159
160 meta_map.insert("status".into(), json!("completed"));
161 meta_map.insert(
162 "completed_at".into(),
163 json!(chrono::Utc::now().to_rfc3339()),
164 );
165
166 if let Some(season) = resolve_active_season(&mut tx).await {
167 meta_map.insert("season_id".into(), json!(season));
168 } else {
169 tracing::warn!(job_id = %job.job_id, "No active season found");
170 }
171
172 if let Some(spec) = resolve_spec_id_from_config(&job.sim_config) {
173 meta_map.insert("spec_id".into(), json!(spec));
174 } else {
175 tracing::warn!(job_id = %job.job_id, "Could not resolve spec_id from sim_config");
176 }
177
178 let meta = serde_json::Value::Object(meta_map);
179
180 if let Err(source) = sqlx::query_file!(
181 "queries/jobs_finalize.sql",
182 meta,
183 job.output.result_pb.as_slice(),
184 job.output.timeline_pb.as_deref(),
185 job.job_id,
186 )
187 .execute(&mut *tx)
188 .await
189 {
190 tracing::error!(error = %source, "Failed to finalize job row");
191
192 return Err(CompletionPersistenceError::Save(source));
193 }
194
195 if let Err(source) = tx.commit().await {
196 tracing::error!(error = %source, "Failed to commit transaction");
197
198 return Err(CompletionPersistenceError::Commit(source));
199 }
200
201 Ok(())
202 }
203}
204
205#[async_trait]
206impl CompletionRepository for SqlCompletionRepository<'_> {
207 async fn persist(&self, job: FinalizedJob) -> Result<(), CompletionPersistenceError> {
208 self.persist_transaction(job).await
209 }
210
211 async fn record_event(&self, event: ChunkEvent<'_>) -> Result<(), sqlx::Error> {
212 let token_uuid: Option<Uuid> = event.claim_token.as_str().parse().ok();
213
214 sqlx::query_file!(
215 "queries/chunk_event_log_insert.sql",
216 event.job_id,
217 event.event_type,
218 token_uuid,
219 Option::<String>::None,
220 Some(event.error),
221 )
222 .execute(self.pool)
223 .await?;
224
225 Ok(())
226 }
227}
228
229pub(in crate::http) struct StateCompletionObserver<'a> {
230 state: &'a ServerState,
231}
232
233impl StateCompletionObserver<'_> {
234 async fn publish_finalized(&self, job_id: Uuid, strategy: &str, node_public_key: &str) {
235 tracing::info!(job_id = %job_id, strategy, node = node_public_key, "Job finalized");
236 metrics::counter!(
237 crate::telemetry::CHUNKS_COMPLETED,
238 "node_public_key" => node_public_key.to_string()
239 )
240 .increment(1);
241 metrics::gauge!(crate::telemetry::CHUNKS_RUNNING).decrement(1.0);
242
243 let event = JobProgressEvent {
244 r#type: "updated",
245 payload: JobProgressPayload {
246 status: "completed",
247 chunks_completed: 0,
248 chunks_total: 0,
249 completed_at: Some(chrono::Utc::now().to_rfc3339()),
250 phase: None,
251 phase_count: None,
252 permutations_active: None,
253 top_mean_dps_x10: None,
254 top_k: Vec::new(),
255 },
256 };
257
258 self.state.publish(&format!("jobs:{job_id}"), &event).await;
259 self.state.publish("jobs:all", &event).await;
260 }
261}
262
263#[async_trait]
264impl CompletionObserver for StateCompletionObserver<'_> {
265 async fn accepted(
266 &self,
267 job_id: Uuid,
268 snapshot: Option<&ProgressSnapshot>,
269 node_public_key: &str,
270 ) {
271 publish_progress(self.state, job_id, snapshot).await;
272 metrics::counter!(
273 crate::telemetry::CHUNKS_COMPLETED,
274 "node_public_key" => node_public_key.to_string()
275 )
276 .increment(1);
277 }
278
279 async fn finalized(&self, job_id: Uuid, strategy: &str, node_public_key: &str) {
280 self.publish_finalized(job_id, strategy, node_public_key)
281 .await;
282 }
283}
284
285pub(in crate::http) struct CompletionWorkflow<'a, R, O> {
286 runtime: &'a RuntimeState,
287 repository: R,
288 observer: O,
289}
290
291pub(in crate::http) fn completion_workflow(
292 state: &ServerState,
293) -> CompletionWorkflow<'_, SqlCompletionRepository<'_>, StateCompletionObserver<'_>> {
294 CompletionWorkflow {
295 runtime: &state.runtime,
296 repository: SqlCompletionRepository {
297 pool: state.dbs.get::<crate::http::HttpDb>(),
298 },
299 observer: StateCompletionObserver { state },
300 }
301}
302
303impl<R, O> CompletionWorkflow<'_, R, O>
304where
305 R: CompletionRepository,
306 O: CompletionObserver,
307{
308 pub(in crate::http) async fn complete(
309 &self,
310 input: CompletionInput,
311 ) -> Result<NodeChunkCompletionResponse, CompletionError> {
312 let (outcome, snapshot) = {
313 let mut runtimes = self.runtime.jobs.write().await;
314 let runtime = runtimes
315 .get_mut(&input.job_id)
316 .ok_or(CompletionError::JobNotFound)?;
317 let outcome = match runtime.complete(
318 input.chunk_id,
319 &input.claim_token,
320 &input.work_context_hash,
321 &input.node_public_key,
322 &input.results,
323 ) {
324 Ok(outcome) => outcome,
325 Err(source) => {
326 tracing::error!(job_id = %input.job_id, error = %source, "Completion ingest failed");
327
328 return Err(CompletionError::Ingest(source));
329 }
330 };
331 let snapshot = runtime.progress_snapshot();
332
333 (outcome, snapshot)
334 };
335
336 match outcome {
337 CompletionOutcome::Accepted {
338 job_complete: false,
339 } => {
340 self.observer
341 .accepted(input.job_id, snapshot.as_ref(), &input.node_public_key)
342 .await;
343
344 Ok(completion_response(false, None))
345 }
346 CompletionOutcome::Accepted { job_complete: true } => {
347 let _ = self.finalize(input.job_id, &input.node_public_key).await?;
348
349 Ok(completion_response(true, None))
350 }
351 CompletionOutcome::Idempotent => {
352 let job_complete = self.finalize(input.job_id, &input.node_public_key).await?;
353
354 Ok(completion_response(job_complete, Some(true)))
355 }
356 CompletionOutcome::Conflict(message) => {
357 self.record_event(
358 input.job_id,
359 "idempotency_conflict",
360 &input.claim_token,
361 &message,
362 )
363 .await;
364
365 Err(CompletionError::Conflict(message))
366 }
367 CompletionOutcome::Stale(message) => {
368 self.record_event(
369 input.job_id,
370 "stale_completion",
371 &input.claim_token,
372 &message,
373 )
374 .await;
375
376 Err(CompletionError::Stale(message))
377 }
378 }
379 }
380
381 async fn finalize(&self, job_id: Uuid, node_public_key: &str) -> Result<bool, CompletionError> {
382 let (finalized, strategy) = {
383 let mut runtimes = self.runtime.jobs.write().await;
384
385 let Some(runtime) = runtimes.get_mut(&job_id) else {
386 return Ok(false);
387 };
388
389 if runtime.status == JobRuntimeStatus::Finalizing {
390 return Ok(false);
391 }
392
393 if runtime.status != JobRuntimeStatus::Completed {
394 return Ok(false);
395 }
396
397 let output = match runtime.finalize() {
398 Ok(output) => output,
399 Err(source) => {
400 tracing::error!(job_id = %job_id, error = %source, "Finalize failed");
401
402 return Err(CompletionError::Finalize(source));
403 }
404 };
405
406 runtime.status = JobRuntimeStatus::Finalizing;
407
408 (
409 FinalizedJob {
410 job_id,
411 sim_config: runtime.base_sim_config.clone(),
412 output,
413 },
414 runtime.strategy_name.clone(),
415 )
416 };
417
418 if let Err(error) = self.repository.persist(finalized).await {
419 if let Some(runtime) = self.runtime.jobs.write().await.get_mut(&job_id) {
420 runtime.status = JobRuntimeStatus::Completed;
421 }
422
423 return Err(error.into());
424 }
425
426 self.runtime.jobs.write().await.remove(&job_id);
427 self.observer
428 .finalized(job_id, &strategy, node_public_key)
429 .await;
430
431 Ok(true)
432 }
433
434 async fn record_event(
435 &self,
436 job_id: Uuid,
437 event_type: &'static str,
438 claim_token: &ClaimToken,
439 error: &str,
440 ) {
441 if let Err(source) = self
442 .repository
443 .record_event(ChunkEvent {
444 job_id,
445 event_type,
446 claim_token,
447 error,
448 })
449 .await
450 {
451 tracing::error!(job_id = %job_id, error = %source, "Failed to write chunk_event_log");
452 }
453 }
454}
455
456const fn completion_response(
457 job_complete: bool,
458 already_completed: Option<bool>,
459) -> NodeChunkCompletionResponse {
460 NodeChunkCompletionResponse {
461 success: true,
462 already_completed,
463 job_complete,
464 }
465}
466
467async fn publish_progress(state: &ServerState, job_id: Uuid, snapshot: Option<&ProgressSnapshot>) {
468 let payload = match snapshot {
469 Some(snapshot) => JobProgressPayload {
470 status: "running",
471 chunks_completed: 0,
472 chunks_total: 0,
473 completed_at: None,
474 phase: Some(snapshot.phase),
475 phase_count: Some(snapshot.phase_count),
476 permutations_active: Some(snapshot.permutations_active),
477 top_mean_dps_x10: Some(snapshot.top_mean_dps_x10),
478 top_k: snapshot
479 .top_k
480 .iter()
481 .map(|entry| JobProgressTopPermutation {
482 perm_idx: entry.perm_idx,
483 mean_dps_x10: entry.mean_dps_x10,
484 iterations: entry.iterations,
485 })
486 .collect(),
487 },
488 None => JobProgressPayload {
489 status: "running",
490 chunks_completed: 0,
491 chunks_total: 0,
492 completed_at: None,
493 phase: None,
494 phase_count: None,
495 permutations_active: None,
496 top_mean_dps_x10: None,
497 top_k: Vec::new(),
498 },
499 };
500 let event = JobProgressEvent {
501 r#type: "updated",
502 payload,
503 };
504
505 state.publish(&format!("jobs:{job_id}"), &event).await;
506}
507
508async fn resolve_active_season(tx: &mut sqlx::Transaction<'_, sqlx::Postgres>) -> Option<String> {
509 struct Row {
510 season_id: String,
511 }
512
513 sqlx::query_file_as!(Row, "queries/resolve_active_season.sql")
514 .fetch_optional(&mut **tx)
515 .await
516 .ok()
517 .flatten()
518 .map(|row| row.season_id)
519}
520
521fn resolve_spec_id_from_config(sim_config: &str) -> Option<i32> {
522 let parsed = intent::parse_sim_config(sim_config).ok()?;
523 let spec = SpecId::parse_wow_spec_id(parsed.spec).ok()?;
524
525 i32::try_from(spec.wow_spec_id()).ok()
526}
527
528#[cfg(test)]
529mod tests;