wowlab_sentinel/scheduler/
repository.rs

1//! Persistence boundary for scheduler assignment rounds.
2
3use async_trait::async_trait;
4use sqlx::PgPool;
5use uuid::Uuid;
6use wowlab_common::NodePublicKey;
7use wowlab_types::sim::FastMap;
8
9#[derive(Debug, thiserror::Error)]
10pub(super) enum RepositoryError {
11    #[error("failed to fetch pending jobs: {0}")]
12    FetchPending(#[source] sqlx::Error),
13    #[error("failed to mark job {job_id} running: {source}")]
14    MarkRunning {
15        job_id: Uuid,
16        #[source]
17        source: sqlx::Error,
18    },
19    #[error("failed to fetch online nodes: {0}")]
20    FetchNodes(#[source] sqlx::Error),
21    #[error("failed to fetch node permissions: {0}")]
22    FetchPermissions(#[source] sqlx::Error),
23    #[error("failed to fetch user Discord identities: {0}")]
24    FetchDiscordIdentities(#[source] sqlx::Error),
25    #[error("failed to fetch friend memberships: {0}")]
26    FetchFriendMemberships(#[source] sqlx::Error),
27}
28
29#[derive(Clone, Debug)]
30pub(super) struct PendingJob {
31    pub id: Uuid,
32    pub user_id: Uuid,
33    pub sim_config: String,
34    pub sentinel_config: String,
35    pub meta: serde_json::Value,
36    pub status: String,
37}
38
39#[derive(Clone, Debug)]
40pub(super) struct OnlineNode {
41    pub public_key: NodePublicKey,
42    pub user_id: Uuid,
43    pub capacity: usize,
44}
45
46#[derive(Clone, Debug)]
47pub(super) struct NodePermission {
48    pub public_key: NodePublicKey,
49    pub access_type: String,
50    pub target_id: Option<String>,
51}
52
53#[async_trait]
54pub(super) trait AssignmentRepository: Send + Sync {
55    async fn pending_jobs(&self) -> Result<Vec<PendingJob>, RepositoryError>;
56    async fn mark_running(&self, job_id: Uuid) -> Result<u64, RepositoryError>;
57    async fn online_nodes(&self) -> Result<Vec<OnlineNode>, RepositoryError>;
58    async fn permissions(
59        &self,
60        public_keys: &[NodePublicKey],
61    ) -> Result<Vec<NodePermission>, RepositoryError>;
62    async fn user_discord_ids(
63        &self,
64        user_ids: &[Uuid],
65    ) -> Result<FastMap<Uuid, String>, RepositoryError>;
66    async fn friend_memberships(
67        &self,
68        user_ids: &[Uuid],
69    ) -> Result<FastMap<Uuid, Vec<Uuid>>, RepositoryError>;
70}
71
72pub(super) struct SqlAssignmentRepository<'a> {
73    pool: &'a PgPool,
74    batch_size: i64,
75}
76
77impl<'a> SqlAssignmentRepository<'a> {
78    pub(super) fn new(pool: &'a PgPool, batch_size: i64) -> Self {
79        Self { pool, batch_size }
80    }
81
82    async fn query_pending_jobs(&self) -> Result<Vec<PendingJob>, sqlx::Error> {
83        struct Row {
84            id: Uuid,
85            user_id: Uuid,
86            sim_config: String,
87            sentinel_config: String,
88            meta: serde_json::Value,
89            status: Option<String>,
90        }
91
92        let rows = sqlx::query_file_as!(
93            Row,
94            "queries/scheduler_fetch_pending_jobs.sql",
95            self.batch_size
96        )
97        .fetch_all(self.pool)
98        .await?;
99
100        Ok(rows
101            .into_iter()
102            .map(|row| PendingJob {
103                id: row.id,
104                user_id: row.user_id,
105                sim_config: row.sim_config,
106                sentinel_config: row.sentinel_config,
107                meta: row.meta,
108                status: row.status.unwrap_or_default(),
109            })
110            .collect())
111    }
112
113    async fn update_running(&self, job_id: Uuid) -> Result<u64, sqlx::Error> {
114        let result = sqlx::query_file!("queries/scheduler_mark_running.sql", job_id)
115            .execute(self.pool)
116            .await?;
117
118        Ok(result.rows_affected())
119    }
120
121    async fn query_online_nodes(&self) -> Result<Vec<OnlineNode>, sqlx::Error> {
122        struct OnlineNodeRow {
123            public_key: String,
124            user_id: Option<Uuid>,
125            total_cores: i32,
126            max_parallel: i32,
127            discord_id: Option<String>,
128        }
129
130        let rows: Vec<OnlineNodeRow> =
131            sqlx::query_file_as!(OnlineNodeRow, "queries/scheduler_fetch_online_nodes.sql")
132                .fetch_all(self.pool)
133                .await?;
134
135        Ok(rows
136            .into_iter()
137            .filter_map(|row| {
138                let public_key = row.public_key.parse::<NodePublicKey>().ok()?;
139                let user_id = row.user_id?;
140                let capacity = node_capacity(row.total_cores, row.max_parallel)?;
141
142                drop(row.discord_id);
143
144                Some(OnlineNode {
145                    public_key,
146                    user_id,
147                    capacity,
148                })
149            })
150            .collect())
151    }
152
153    async fn query_permissions(
154        &self,
155        public_keys: &[NodePublicKey],
156    ) -> Result<Vec<NodePermission>, sqlx::Error> {
157        struct NodePermissionRow {
158            public_key: String,
159            access_type: String,
160            target_id: Option<String>,
161        }
162
163        let keys: Vec<String> = public_keys.iter().map(NodePublicKey::to_base64).collect();
164        let rows: Vec<NodePermissionRow> = sqlx::query_file_as!(
165            NodePermissionRow,
166            "queries/scheduler_fetch_permissions.sql",
167            &keys
168        )
169        .fetch_all(self.pool)
170        .await?;
171
172        Ok(rows
173            .into_iter()
174            .filter_map(|row| {
175                let public_key = row.public_key.parse::<NodePublicKey>().ok()?;
176
177                Some(NodePermission {
178                    public_key,
179                    access_type: row.access_type,
180                    target_id: row.target_id,
181                })
182            })
183            .collect())
184    }
185
186    async fn query_user_discord_ids(
187        &self,
188        user_ids: &[Uuid],
189    ) -> Result<FastMap<Uuid, String>, sqlx::Error> {
190        struct DiscordIdRow {
191            user_id: Option<Uuid>,
192            discord_id: Option<String>,
193        }
194
195        let rows: Vec<DiscordIdRow> = sqlx::query_file_as!(
196            DiscordIdRow,
197            "queries/scheduler_fetch_user_discord_ids.sql",
198            user_ids
199        )
200        .fetch_all(self.pool)
201        .await?;
202
203        Ok(rows
204            .into_iter()
205            .filter_map(|row| Some((row.user_id?, row.discord_id?)))
206            .collect())
207    }
208
209    async fn query_friend_memberships(
210        &self,
211        user_ids: &[Uuid],
212    ) -> Result<FastMap<Uuid, Vec<Uuid>>, sqlx::Error> {
213        struct MembershipRow {
214            user_id: Uuid,
215            list_id: Uuid,
216        }
217
218        let rows: Vec<MembershipRow> = sqlx::query_file_as!(
219            MembershipRow,
220            "queries/scheduler_fetch_friend_memberships.sql",
221            user_ids
222        )
223        .fetch_all(self.pool)
224        .await?;
225
226        Ok(rows.into_iter().fold(
227            FastMap::<Uuid, Vec<Uuid>>::default(),
228            |mut memberships, row| {
229                memberships
230                    .entry(row.user_id)
231                    .or_default()
232                    .push(row.list_id);
233
234                memberships
235            },
236        ))
237    }
238}
239
240#[async_trait]
241impl AssignmentRepository for SqlAssignmentRepository<'_> {
242    async fn pending_jobs(&self) -> Result<Vec<PendingJob>, RepositoryError> {
243        match self.query_pending_jobs().await {
244            Ok(jobs) => Ok(jobs),
245            Err(source) => Err(RepositoryError::FetchPending(source)),
246        }
247    }
248
249    async fn mark_running(&self, job_id: Uuid) -> Result<u64, RepositoryError> {
250        match self.update_running(job_id).await {
251            Ok(rows) => Ok(rows),
252            Err(source) => Err(RepositoryError::MarkRunning { job_id, source }),
253        }
254    }
255
256    async fn online_nodes(&self) -> Result<Vec<OnlineNode>, RepositoryError> {
257        match self.query_online_nodes().await {
258            Ok(nodes) => Ok(nodes),
259            Err(source) => Err(RepositoryError::FetchNodes(source)),
260        }
261    }
262
263    async fn permissions(
264        &self,
265        public_keys: &[NodePublicKey],
266    ) -> Result<Vec<NodePermission>, RepositoryError> {
267        match self.query_permissions(public_keys).await {
268            Ok(permissions) => Ok(permissions),
269            Err(source) => Err(RepositoryError::FetchPermissions(source)),
270        }
271    }
272
273    async fn user_discord_ids(
274        &self,
275        user_ids: &[Uuid],
276    ) -> Result<FastMap<Uuid, String>, RepositoryError> {
277        match self.query_user_discord_ids(user_ids).await {
278            Ok(identities) => Ok(identities),
279            Err(source) => Err(RepositoryError::FetchDiscordIdentities(source)),
280        }
281    }
282
283    async fn friend_memberships(
284        &self,
285        user_ids: &[Uuid],
286    ) -> Result<FastMap<Uuid, Vec<Uuid>>, RepositoryError> {
287        match self.query_friend_memberships(user_ids).await {
288            Ok(memberships) => Ok(memberships),
289            Err(source) => Err(RepositoryError::FetchFriendMemberships(source)),
290        }
291    }
292}
293
294fn node_capacity(total_cores: i32, max_parallel: i32) -> Option<usize> {
295    let total_cores = usize::try_from(total_cores).ok()?;
296    let max_parallel = usize::try_from(max_parallel).ok()?;
297    let capacity = total_cores.min(max_parallel);
298
299    (capacity > 0).then_some(capacity)
300}
301
302#[cfg(test)]
303mod tests {
304    use googletest::prelude::*;
305
306    use super::*;
307
308    #[gtest]
309    fn node_capacity_rejects_non_positive_database_values() -> Result<()> {
310        verify_eq!(node_capacity(8, 4), Some(4))?;
311        verify_eq!(node_capacity(4, 8), Some(4))?;
312        verify_eq!(node_capacity(0, 4), None)?;
313        verify_eq!(node_capacity(4, 0), None)?;
314        verify_eq!(node_capacity(-1, 4), None)?;
315        verify_eq!(node_capacity(4, -1), None)?;
316
317        Ok(())
318    }
319}
wowlab_sentinel/scheduler/repository.rs

wowlab_sentinel/scheduler/
repository.rs
