Skip to main content

flams_system/building/
queue_manager.rs

1use flams_math_archives::backend::{AnyBackend, SandboxedBackend, SandboxedRepository};
2use flams_utils::vecmap::VecMap;
3use std::{fmt::Display, num::NonZeroU32, sync::atomic::AtomicU8};
4
5use crate::settings::Settings;
6
7use super::queue::{Queue, QueueName, QueueState, RunningQueue};
8
9#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
10pub struct QueueId(NonZeroU32);
11impl QueueId {
12    #[must_use]
13    pub fn global() -> Self {
14        Self(NonZeroU32::new(1).unwrap_or_else(|| unreachable!()))
15    }
16}
17impl Display for QueueId {
18    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
19        write!(f, "queue {}", self.0)
20    }
21}
22impl From<QueueId> for NonZeroU32 {
23    #[inline]
24    fn from(id: QueueId) -> Self {
25        id.0
26    }
27}
28impl From<NonZeroU32> for QueueId {
29    #[inline]
30    fn from(id: NonZeroU32) -> Self {
31        Self(id)
32    }
33}
34
35#[derive(Debug)]
36pub struct QueueManager {
37    inner: parking_lot::RwLock<VecMap<QueueId, Queue>>,
38    threads: Semaphore,
39}
40
41static QUEUE_MANAGER: std::sync::OnceLock<QueueManager> = std::sync::OnceLock::new();
42
43impl QueueManager {
44    #[inline]
45    pub fn clear() {
46        if let Some(m) = QUEUE_MANAGER.get() {
47            m.inner.write().0.clear();
48        }
49    }
50    pub fn initialize(num_threads: u8) {
51        QUEUE_MANAGER.get_or_init(|| {
52            let init = if crate::settings::Settings::get().admin_pwd.is_some() {
53                VecMap::new()
54            } else {
55                vec![(
56                    QueueId::global(),
57                    Queue::new(QueueId::global(), QueueName::Global, AnyBackend::Global),
58                )]
59                .into()
60            };
61            Self {
62                inner: parking_lot::RwLock::new(init),
63                threads: {
64                    #[cfg(feature = "tokio")]
65                    {
66                        match num_threads {
67                            0 => Semaphore::Linear,
68                            i => Semaphore::Counting {
69                                inner: std::sync::Arc::new(tokio::sync::Semaphore::new(i as usize)),
70                                _num: flams_utils::triomphe::Arc::new(AtomicU8::new(i)),
71                            },
72                        }
73                    }
74                    #[cfg(not(feature = "tokio"))]
75                    {
76                        Semaphore::Linear
77                    }
78                },
79            }
80        });
81    }
82    /// ### Panics
83    pub fn get() -> &'static Self {
84        QUEUE_MANAGER.get().expect("Queue manager not initialized")
85    }
86
87    pub fn new_queue(&self, queue_name: &str) -> QueueId {
88        super::BUILD_QUEUE_SPAN.in_scope(move || {
89      let mut inner = self.inner.write();
90      let mut count = 0;
91      loop {
92        if inner.0.iter().any(|(_,q)| matches!(q.name(),QueueName::Sandbox{name,idx} if &**name == queue_name && *idx == count)) {
93          count += 1;
94        } else {break}
95      }
96      let sbname = format!("{queue_name}_{count}");
97      tracing::info_span!("Build Queue",name = &sbname).in_scope(|| {
98        let id = QueueId(NonZeroU32::new(inner.0.iter().map(|(k,_)| k.0.get()).max().unwrap_or_default() + 1).unwrap_or_else(|| unreachable!()));
99        let backend = AnyBackend::Sandbox(SandboxedBackend::new(&sbname,&Settings::get().temp_dir()));
100        inner.insert(id,
101          Queue::new(id,
102            QueueName::Sandbox{name:queue_name.to_string().into(),idx:count},
103            backend
104          )
105        );
106        id
107      })
108    })
109    }
110
111    pub fn all_queues(&self) -> Vec<(QueueId, QueueName, Option<Vec<SandboxedRepository>>)> {
112        let inner = self.inner.read();
113        inner
114            .iter()
115            .map(|(k, v)| {
116                (
117                    *k,
118                    v.name().clone(),
119                    if let AnyBackend::Sandbox(sb) = v.backend() {
120                        Some(sb.get_repos())
121                    } else {
122                        None
123                    },
124                )
125            })
126            .collect()
127    }
128
129    pub fn with_all_queues<R>(&self, f: impl FnOnce(&[(QueueId, Queue)]) -> R) -> R {
130        f(&self.inner.read().0)
131    }
132
133    pub fn queues_for_user(
134        &self,
135        user_name: &str,
136    ) -> Vec<(QueueId, QueueName, Option<Vec<SandboxedRepository>>)> {
137        let inner = self.inner.read();
138        inner
139            .iter()
140            .filter_map(|(k, v)| {
141                if let QueueName::Sandbox { name, .. } = v.name() {
142                    if name
143                        .strip_prefix(user_name)
144                        .is_some_and(|rest| rest.as_bytes().iter().all(|b| b.is_ascii_digit()))
145                    {
146                        Some((
147                            *k,
148                            v.name().clone(),
149                            if let AnyBackend::Sandbox(sb) = v.backend() {
150                                Some(sb.get_repos())
151                            } else {
152                                None
153                            },
154                        ))
155                    } else {
156                        None
157                    }
158                } else {
159                    None
160                }
161            })
162            .collect()
163    }
164
165    pub fn with_queue<R>(&self, id: QueueId, f: impl FnOnce(Option<&Queue>) -> R) -> R {
166        let inner = self.inner.read();
167        f(inner.get(&id))
168    }
169
170    /// #### Errors
171    #[cfg(feature = "tokio")]
172    pub fn migrate<R>(
173        &self,
174        id: QueueId,
175        then: impl FnOnce(&SandboxedBackend) -> eyre::Result<R>,
176    ) -> eyre::Result<(R, usize)> {
177        self.migrate_i(id, then).map_err(|e| {
178            tracing::error!("Error migrating: {e:#}");
179            e
180        })
181    }
182
183    #[cfg(feature = "tokio")]
184    #[allow(clippy::significant_drop_tightening)]
185    fn migrate_i<R>(
186        &self,
187        id: QueueId,
188        then: impl FnOnce(&SandboxedBackend) -> eyre::Result<R>,
189    ) -> eyre::Result<(R, usize)> {
190        use eyre::eyre;
191        use flams_math_archives::utils::SyncEngine;
192        use flams_utils::impossible;
193
194        let mut inner = self.inner.write();
195        let r = if let Some(queue) = inner.get(&id) {
196            if !matches!(&*queue.0.state.read(), QueueState::Finished { .. }) {
197                return Err(eyre!("Queue {id} not finished"));
198            }
199            if !matches!(queue.backend(), AnyBackend::Sandbox(_)) {
200                return Err(eyre!("Global Queue can not be migrated"));
201            }
202            let AnyBackend::Sandbox(sandbox) = queue.backend() else {
203                unreachable!()
204            };
205            then(sandbox)?
206        } else {
207            return Err(eyre!("No queue {id} found"));
208        };
209        let Some(queue) = inner.remove(&id) else {
210            unreachable!()
211        };
212        let AnyBackend::Sandbox(sandbox) = queue.backend() else {
213            impossible!()
214        };
215        let num = sandbox.migrate::<SyncEngine>()?;
216
217        let _ = tokio::task::spawn_blocking(|| {
218            for e in crate::iter::<crate::FlamsExtension>() {
219                (e.on_reload)();
220            }
221        });
222        Ok((r, num))
223    }
224
225    #[allow(clippy::significant_drop_tightening)]
226    pub fn delete(&self, id: QueueId) {
227        let mut inner = self.inner.write();
228        if let Some(q) = inner.remove(&id) {
229            let mut s = q.0.state.write();
230            if let QueueState::Running(RunningQueue { queue, blocked, .. }) = &mut *s {
231                queue.clear();
232                blocked.clear();
233            }
234            if matches!(q.name(), QueueName::Global) {
235                inner.insert(id, Queue::new(id, QueueName::Global, AnyBackend::Global));
236            }
237        }
238    }
239
240    /// ### Errors
241    /// if no queue with that id exists
242    #[allow(clippy::result_unit_err)]
243    pub fn start_queue(&self, id: QueueId) -> Result<(), ()> {
244        let sem = self.threads.clone();
245        self.with_queue(id, |q| {
246            let Some(q) = q else { return Err(()) };
247            q.start(sem);
248            Ok(())
249        })
250    }
251
252    /// ### Errors
253    /// if no queue with that id exists
254    #[allow(clippy::result_unit_err)]
255    pub fn requeue_failed(&self, id: QueueId) -> Result<(), ()> {
256        self.with_queue(id, |q| {
257            q.map_or(Err(()), |q| {
258                q.requeue_failed();
259                Ok(())
260            })
261        })
262    }
263
264    pub fn with_global<R>(&self, f: impl FnOnce(&Queue) -> R) -> R {
265        let inner = self.inner.read();
266        f(inner
267            .get(&QueueId::global())
268            .unwrap_or_else(|| unreachable!()))
269    }
270}
271
272#[derive(Debug, Clone)]
273pub(crate) enum Semaphore {
274    Linear,
275    #[cfg(feature = "tokio")]
276    Counting {
277        inner: std::sync::Arc<tokio::sync::Semaphore>,
278        _num: flams_utils::triomphe::Arc<AtomicU8>,
279    },
280}