1use tokio::sync::mpsc;
2use uuid::Uuid;
3use wowlab_common::RuntimeChunkPayload;
4use wowlab_types::proto::{BatchChunkCompletion, BatchWorkResult};
5
6use super::{NodeCore, NodeCoreEvent, validated_capacity};
7use crate::{ConnectionStatus, NodeState, WorkBatch, WorkBatchResult, realtime::RealtimeEvent};
8
9const MAX_EVENTS_PER_POLL: usize = 10;
10
11fn log_realtime_restart() {
12 tracing::info!("Realtime channel dropped; restarting subscription");
13}
14
15impl NodeCore {
16 pub(super) fn check_realtime_events(&mut self) {
17 let Some(ref mut rx) = self.realtime_rx else {
18 return;
19 };
20
21 let mut events = Vec::with_capacity(MAX_EVENTS_PER_POLL);
22
23 for _ in 0..MAX_EVENTS_PER_POLL {
24 match rx.try_recv() {
25 Ok(event) => events.push(event),
26 Err(mpsc::error::TryRecvError::Empty) => break,
27 Err(mpsc::error::TryRecvError::Disconnected) => {
28 self.realtime_rx = None;
29 self.set_connection(ConnectionStatus::Disconnected);
30
31 if matches!(self.state, NodeState::Running) && self.registered {
32 log_realtime_restart();
33 self.start_realtime();
34 }
35
36 return;
37 }
38 }
39 }
40
41 for event in events {
42 self.handle_realtime_event(&event);
43 }
44 }
45
46 pub(super) fn check_work_results(&mut self) {
47 let Some(ref mut rx) = self.result_rx else {
48 return;
49 };
50
51 let mut results = Vec::with_capacity(MAX_EVENTS_PER_POLL);
52 let mut disconnected = false;
53
54 for _ in 0..MAX_EVENTS_PER_POLL {
55 match rx.try_recv() {
56 Ok(result) => results.push(result),
57 Err(mpsc::error::TryRecvError::Empty) => break,
58 Err(mpsc::error::TryRecvError::Disconnected) => {
59 disconnected = true;
60 break;
61 }
62 }
63 }
64
65 if disconnected {
66 self.result_rx = None;
67 }
68
69 for result in results {
70 self.handle_work_result(&result);
71 }
72 }
73
74 fn handle_realtime_event(&mut self, event: &RealtimeEvent) {
75 if !matches!(self.state, NodeState::Running) {
76 tracing::debug!(state = ?self.state, "Ignoring realtime event outside Running state");
77
78 return;
79 }
80
81 match event {
82 RealtimeEvent::Connected => {
83 tracing::info!("Connected");
84 self.set_connection(ConnectionStatus::Connected);
85 }
86 RealtimeEvent::Disconnected => {
87 tracing::info!("Disconnected, reconnecting...");
88 self.set_connection(ConnectionStatus::Disconnected);
89 }
90 RealtimeEvent::NodeUpdated {
91 name,
92 total_cores,
93 max_parallel,
94 } => self.handle_node_update(name, *total_cores, *max_parallel),
95 RealtimeEvent::ChunkAssigned(payload) => self.handle_chunk_assigned(payload),
96 RealtimeEvent::Error(err) => {
97 tracing::warn!(error = %err, "Realtime connection error");
98 let _ = self.event_tx.try_send(NodeCoreEvent::Error(err.clone()));
99 }
100 }
101 }
102
103 fn handle_node_update(&mut self, name: &str, total_cores: i32, max_parallel: i32) {
104 let Some((total_cores, max_parallel)) = validated_capacity(total_cores, max_parallel)
105 else {
106 tracing::warn!(
107 total_cores,
108 max_parallel,
109 "Ignoring invalid node capacity update"
110 );
111
112 return;
113 };
114
115 name.clone_into(&mut self.node_name);
116 self.worker_pool.set_max_workers(max_parallel);
117 self.total_cores = total_cores;
118 }
119
120 fn handle_chunk_assigned(&mut self, payload: &RuntimeChunkPayload) {
121 let job_id = match Uuid::parse_str(&payload.job_id) {
122 Ok(id) => id,
123 Err(error) => {
124 tracing::error!(job_id = %payload.job_id, %error, "Invalid job id in chunk");
125
126 return;
127 }
128 };
129
130 let item_count = payload.work_items.len();
131 let iterations: u64 = payload
132 .work_items
133 .iter()
134 .map(|i| u64::from(i.iterations))
135 .sum();
136 #[expect(
137 clippy::cast_possible_truncation,
138 reason = "legacy display index is signed 32-bit"
139 )]
140 let chunk_index = payload.chunk_id.get() as i32;
141 #[expect(
142 clippy::cast_possible_truncation,
143 reason = "legacy display counter is signed 32-bit"
144 )]
145 let iterations_i32 = iterations.min(i32::MAX as u64) as i32;
146
147 tracing::info!(%job_id, chunk_id = payload.chunk_id.get(), item_count, iterations, "Chunk assigned");
148 let _ = self.event_tx.try_send(NodeCoreEvent::ChunkAssigned {
149 job_id,
150 chunk_index,
151 iterations: iterations_i32,
152 });
153
154 self.process_chunk(job_id, payload);
155 }
156
157 fn process_chunk(&mut self, job_id: Uuid, payload: &RuntimeChunkPayload) {
158 let chunk_id = payload.chunk_id;
159 let claim_token = payload.claim_token.clone();
160 let work_context_hash = payload.work_context_hash;
161 let items = payload.work_items.clone();
162 let cache = self.work_context_cache.clone();
163 let sentinel = self.sentinel.clone();
164 let work_tx = self.worker_pool.work_tx();
165 let event_tx = self.event_tx.clone();
166
167 self.runtime.spawn(async move {
168 let context = match cache
169 .get_or_fetch(&sentinel, job_id, work_context_hash, &claim_token)
170 .await
171 {
172 Ok(context) => context,
173 Err(error) => {
174 tracing::error!(%job_id, %error, "Failed to resolve work context");
175 #[expect(
176 clippy::cast_possible_truncation,
177 reason = "legacy display index is signed 32-bit"
178 )]
179 let chunk_index = chunk_id.get() as i32;
180 let _ = event_tx
181 .send(NodeCoreEvent::ChunkFailed {
182 job_id,
183 chunk_index,
184 error: error.to_string(),
185 })
186 .await;
187
188 return;
189 }
190 };
191
192 let batch = WorkBatch {
193 job_id,
194 chunk_id,
195 claim_token,
196 work_context_hash,
197 context,
198 items,
199 };
200
201 if let Some(tx) = work_tx {
202 if let Err(error) = tx.send(batch).await {
203 tracing::error!(%error, "Failed to submit work batch");
204 }
205 }
206 });
207 }
208
209 fn handle_work_result(&self, result: &WorkBatchResult) {
210 let job_id = result.job_id;
211 let chunk_id = result.chunk_id;
212 #[expect(
213 clippy::cast_possible_truncation,
214 reason = "legacy display index is signed 32-bit"
215 )]
216 let chunk_index = chunk_id.get() as i32;
217 let sentinel = self.sentinel.clone();
218 let event_tx = self.event_tx.clone();
219 let mean_dps = batch_mean_dps(result);
220 let completion = build_completion(result);
221
222 self.runtime.spawn(async move {
223 match sentinel.complete_batch(&completion).await {
224 Ok(()) => {
225 tracing::info!(%job_id, chunk_id = chunk_id.get(), result_count = completion.results.len(), mean_dps, "Chunk completed");
226 let _ = event_tx
227 .send(NodeCoreEvent::ChunkCompleted {
228 job_id,
229 chunk_index,
230 mean_dps,
231 })
232 .await;
233 }
234 Err(error) => {
235 tracing::error!(%job_id, chunk_id = chunk_id.get(), %error, "Failed to submit chunk result");
236 let _ = event_tx
237 .send(NodeCoreEvent::ChunkFailed {
238 job_id,
239 chunk_index,
240 error: error.to_string(),
241 })
242 .await;
243 }
244 }
245 });
246 }
247}
248
249fn build_completion(result: &WorkBatchResult) -> BatchChunkCompletion {
250 let results: Vec<BatchWorkResult> = result
251 .results
252 .iter()
253 .map(wowlab_common::RuntimeWorkResult::to_proto)
254 .collect();
255
256 BatchChunkCompletion {
257 job_id: result.job_id.to_string(),
258 chunk_id: result.chunk_id.get(),
259 claim_token: result.claim_token.as_str().to_string(),
260 work_context_hash: result.work_context_hash.as_bytes().to_vec(),
261 results,
262 }
263}
264
265#[expect(
267 clippy::cast_precision_loss,
268 clippy::cast_possible_truncation,
269 reason = "aggregate DPS is display-only f32"
270)]
271fn batch_mean_dps(result: &WorkBatchResult) -> f32 {
272 if result.results.is_empty() {
273 return 0.0;
274 }
275
276 let sum: u64 = result
277 .results
278 .iter()
279 .map(|r| u64::from(r.mean_dps_x10))
280 .sum();
281 let mean_x10 = sum as f64 / result.results.len() as f64;
282
283 (mean_x10 / wowlab_types::constants::PROTO_DPS_SCALE) as f32
284}