Skip to main content

aven_core/attachments/
blocking.rs

1use std::sync::{Arc, LazyLock};
2
3use anyhow::{Context, Result};
4use tokio::sync::Semaphore;
5
6const MAX_ATTACHMENT_WORKERS: usize = 2;
7const MAX_PREVIEW_WORKERS: usize = 2;
8
9static ATTACHMENT_WORKERS: LazyLock<Arc<Semaphore>> =
10    LazyLock::new(|| Arc::new(Semaphore::new(MAX_ATTACHMENT_WORKERS)));
11static PREVIEW_WORKERS: LazyLock<Arc<Semaphore>> =
12    LazyLock::new(|| Arc::new(Semaphore::new(MAX_PREVIEW_WORKERS)));
13
14pub async fn run<F, T>(work: F) -> Result<T>
15where
16    F: FnOnce() -> Result<T> + Send + 'static,
17    T: Send + 'static,
18{
19    run_with_pool(ATTACHMENT_WORKERS.clone(), work).await
20}
21
22pub async fn run_preview<F, T>(work: F) -> Result<T>
23where
24    F: FnOnce() -> Result<T> + Send + 'static,
25    T: Send + 'static,
26{
27    run_with_pool(PREVIEW_WORKERS.clone(), work).await
28}
29
30async fn run_with_pool<F, T>(workers: Arc<Semaphore>, work: F) -> Result<T>
31where
32    F: FnOnce() -> Result<T> + Send + 'static,
33    T: Send + 'static,
34{
35    let permit = workers
36        .acquire_owned()
37        .await
38        .context("image worker pool closed")?;
39    tokio::task::spawn_blocking(move || {
40        let _permit = permit;
41        work()
42    })
43    .await
44    .context("image worker failed")?
45}
46
47#[cfg(test)]
48mod tests {
49    use std::time::Duration;
50
51    use super::MAX_PREVIEW_WORKERS;
52
53    #[tokio::test]
54    async fn preview_capacity_does_not_delay_attachment_work() {
55        let mut started = Vec::new();
56        let mut previews = Vec::new();
57        for _ in 0..MAX_PREVIEW_WORKERS {
58            let (started_tx, started_rx) = tokio::sync::oneshot::channel();
59            previews.push(tokio::spawn(async move {
60                super::run_preview(move || {
61                    let _ = started_tx.send(());
62                    std::thread::sleep(Duration::from_millis(150));
63                    Ok(())
64                })
65                .await
66            }));
67            started.push(started_rx);
68        }
69        for started in started {
70            started.await.unwrap();
71        }
72
73        tokio::time::timeout(Duration::from_millis(100), super::run(|| Ok(())))
74            .await
75            .expect("attachment pool should remain available")
76            .unwrap();
77        for preview in previews {
78            preview.await.unwrap().unwrap();
79        }
80    }
81}