wowlab_sentinel/scheduler/
dispatch.rs

1//! Load-balanced runtime claiming and realtime dispatch.
2
3use async_trait::async_trait;
4use uuid::Uuid;
5use wowlab_common::{ClaimToken, NodePublicKey, RuntimeChunkPayload};
6
7use super::{
8    JobProgressEvent, JobProgressPayload,
9    backlog::NodeBacklogs,
10    eligibility::{EligibilityContext, priority},
11    planning::PlannedJob,
12    repository::OnlineNode,
13};
14use crate::state::RuntimeState;
15
16pub(super) trait AssignmentClock: Send + Sync {
17    fn now_ms(&self) -> u64;
18}
19
20pub(super) trait ClaimTokenSource: Send + Sync {
21    fn next(&self) -> ClaimToken;
22}
23
24#[derive(Clone, Copy, Debug, Eq, PartialEq)]
25pub(super) enum PublishOutcome {
26    Published,
27    Rejected,
28}
29
30#[async_trait]
31pub(super) trait AssignmentPublisher: Send + Sync {
32    async fn job_running(&self, job_id: Uuid) -> PublishOutcome;
33    async fn chunk(
34        &self,
35        public_key: &NodePublicKey,
36        payload: &RuntimeChunkPayload,
37    ) -> PublishOutcome;
38}
39
40pub(super) struct SystemClock;
41
42impl AssignmentClock for SystemClock {
43    fn now_ms(&self) -> u64 {
44        wowlab_common::time::unix_timestamp_secs()
45            .saturating_mul(wowlab_types::constants::MS_PER_SECOND_U64)
46    }
47}
48
49pub(super) struct UuidClaimTokenSource;
50
51impl ClaimTokenSource for UuidClaimTokenSource {
52    fn next(&self) -> ClaimToken {
53        ClaimToken::new(Uuid::new_v4().to_string())
54    }
55}
56
57pub(super) struct StateAssignmentPublisher<'a> {
58    state: &'a crate::state::ServerState,
59}
60
61impl<'a> StateAssignmentPublisher<'a> {
62    pub(super) const fn new(state: &'a crate::state::ServerState) -> Self {
63        Self { state }
64    }
65}
66
67#[async_trait]
68impl AssignmentPublisher for StateAssignmentPublisher<'_> {
69    async fn job_running(&self, job_id: Uuid) -> PublishOutcome {
70        let event = JobProgressEvent {
71            r#type: "updated",
72            payload: JobProgressPayload {
73                status: "running",
74                chunks_completed: 0,
75                chunks_total: 0,
76                completed_at: None,
77                phase: None,
78                phase_count: None,
79                permutations_active: None,
80                top_mean_dps_x10: None,
81                top_k: Vec::new(),
82            },
83        };
84
85        self.state.publish(&format!("jobs:{job_id}"), &event).await;
86
87        PublishOutcome::Published
88    }
89
90    async fn chunk(
91        &self,
92        public_key: &NodePublicKey,
93        payload: &RuntimeChunkPayload,
94    ) -> PublishOutcome {
95        self.state
96            .publish(&format!("chunks:{public_key}"), payload)
97            .await;
98
99        PublishOutcome::Published
100    }
101}
102
103#[derive(Debug, Default, Eq, PartialEq)]
104pub(super) struct DispatchReport {
105    pub assigned: usize,
106    pub no_eligible_node: usize,
107    pub no_work: usize,
108    pub rejected_publications: usize,
109}
110
111pub(super) struct Dispatcher<'a, P, C, T> {
112    runtime: &'a RuntimeState,
113    publisher: P,
114    clock: C,
115    claim_tokens: T,
116}
117
118pub(super) struct DispatchDependencies<P, C, T> {
119    pub publisher: P,
120    pub clock: C,
121    pub claim_tokens: T,
122}
123
124impl<'a, P, C, T> Dispatcher<'a, P, C, T>
125where
126    P: AssignmentPublisher,
127    C: AssignmentClock,
128    T: ClaimTokenSource,
129{
130    pub(super) fn new(
131        runtime: &'a RuntimeState,
132        dependencies: DispatchDependencies<P, C, T>,
133    ) -> Self {
134        Self {
135            runtime,
136            publisher: dependencies.publisher,
137            clock: dependencies.clock,
138            claim_tokens: dependencies.claim_tokens,
139        }
140    }
141
142    pub(super) async fn dispatch(
143        &self,
144        jobs: &[PlannedJob],
145        nodes: &[OnlineNode],
146        eligibility: &EligibilityContext<'_>,
147        backlogs: &mut NodeBacklogs,
148    ) -> DispatchReport {
149        let mut report = DispatchReport::default();
150
151        for planned in jobs {
152            let Some(node) = select_node(nodes, planned, eligibility, backlogs) else {
153                report.no_eligible_node += 1;
154                continue;
155            };
156
157            let node_key = node.public_key.to_base64();
158            let batch = {
159                let mut runtimes = self.runtime.jobs.write().await;
160                let Some(runtime) = runtimes.get_mut(&planned.job.id) else {
161                    continue;
162                };
163                let max_items = runtime.suggested_batch_size();
164
165                runtime.claim_batch_with(&node_key, max_items, self.clock.now_ms(), || {
166                    self.claim_tokens.next()
167                })
168            };
169
170            let Some(batch) = batch else {
171                report.no_work += 1;
172                continue;
173            };
174
175            if planned.job.status == "pending"
176                && self.publisher.job_running(planned.job.id).await == PublishOutcome::Rejected
177            {
178                report.rejected_publications += 1;
179            }
180
181            let payload = chunk_payload(planned.job.id, batch);
182
183            if self.publisher.chunk(&node.public_key, &payload).await == PublishOutcome::Rejected {
184                report.rejected_publications += 1;
185            }
186
187            backlogs.increment(&node.public_key);
188            report.assigned += 1;
189        }
190
191        report
192    }
193}
194
195fn chunk_payload(job_id: Uuid, batch: super::runtime::AssignedBatch) -> RuntimeChunkPayload {
196    RuntimeChunkPayload {
197        job_id: job_id.to_string(),
198        chunk_id: batch.chunk_id,
199        work_context_hash: batch.work_context_hash,
200        claim_token: batch.claim_token,
201        work_items: batch.work_items,
202    }
203}
204
205fn select_node<'a>(
206    nodes: &'a [OnlineNode],
207    planned: &PlannedJob,
208    eligibility: &EligibilityContext<'_>,
209    backlogs: &NodeBacklogs,
210) -> Option<&'a OnlineNode> {
211    let eligible_nodes = nodes.iter().filter(|node| {
212        planned.config.target_nodes.is_empty()
213            || planned
214                .config
215                .target_nodes
216                .contains(&node.public_key.to_base64())
217    });
218    let candidates = eligible_nodes.filter_map(|node| {
219        let priority = priority(node, planned.job.user_id, eligibility)?;
220        let available = backlogs.available(&node.public_key, node.capacity)?;
221
222        Some((node, priority, available))
223    });
224
225    candidates
226        .max_by(
227            |(_, left_priority, left_capacity), (_, right_priority, right_capacity)| {
228                left_priority
229                    .cmp(right_priority)
230                    .then(left_capacity.cmp(right_capacity))
231            },
232        )
233        .map(|(node, _, _)| node)
234}
235
236#[cfg(test)]
237mod tests;
wowlab_sentinel/scheduler/dispatch.rs

wowlab_sentinel/scheduler/
dispatch.rs
