Skip to main content

async_geotiff/
blocking.rs

1//! Shared defaults and a shared permit pool for CPU-heavy work dispatched
2//! with `spawn_blocking` (decode, stitch, warp, render).
3
4use std::sync::Arc;
5
6use tokio::sync::{OwnedSemaphorePermit, Semaphore};
7
8use crate::error::{AgError, AgResult};
9
10pub(crate) fn default_blocking_permits() -> usize {
11    std::thread::available_parallelism()
12        .map(|n| n.get())
13        .unwrap_or(4)
14        .max(1)
15}
16
17/// A shared bound on concurrent CPU-heavy blocking jobs.
18///
19/// One `CpuLimiter` can be cloned and handed to many [`CogReader`]s and
20/// [`Tiler`]s so the *whole process* (or a whole mosaic) shares a single CPU
21/// budget instead of each reader/tiler getting its own semaphore. Without this,
22/// a mosaic that opens N assets would allow N × `available_parallelism` decode
23/// jobs plus N independent warp pools, oversubscribing the machine and thrashing
24/// caches under load.
25///
26/// Cloning is cheap (an `Arc` bump); all clones draw permits from the same pool.
27/// Each default construction creates an independent pool sized to available
28/// parallelism, or four permits if that information is unavailable. Running
29/// blocking jobs retain their permits after request cancellation until they exit.
30///
31/// **Deadlock-freedom invariant:** a job must never hold one permit while
32/// awaiting a second from the same limiter. Acquire → `spawn_blocking` → drop
33/// must be flat, not nested. The tile pipeline upholds this by releasing a
34/// reader's decode permit (inside its own `spawn_blocking`) before the warp/
35/// render permit is acquired — see `Tiler::run_tile_cpu`.
36///
37/// [`CogReader`]: crate::io::CogReader
38/// [`Tiler`]: crate::tiler::Tiler
39#[derive(Clone, Debug)]
40pub struct CpuLimiter(Arc<Semaphore>);
41
42impl CpuLimiter {
43    /// Build an independent limiter allowing `permits` concurrent blocking jobs,
44    /// clamped to `1..=tokio::sync::Semaphore::MAX_PERMITS`.
45    pub fn new(permits: usize) -> Self {
46        Self(Arc::new(Semaphore::new(
47            permits.clamp(1, Semaphore::MAX_PERMITS),
48        )))
49    }
50
51    /// Acquire one permit, awaited asynchronously so excess work yields the
52    /// Tokio worker instead of flooding the blocking pool. Hold the returned
53    /// permit for the duration of the `spawn_blocking` job.
54    pub(crate) async fn acquire(&self) -> AgResult<OwnedSemaphorePermit> {
55        Arc::clone(&self.0)
56            .acquire_owned()
57            .await
58            .map_err(|_| AgError::Warp("cpu limiter closed".into()))
59    }
60
61    /// Current number of free permits (diagnostics/tests).
62    pub fn available_permits(&self) -> usize {
63        self.0.available_permits()
64    }
65}
66
67impl Default for CpuLimiter {
68    fn default() -> Self {
69        Self::new(default_blocking_permits())
70    }
71}
72
73#[cfg(test)]
74mod tests {
75    use super::CpuLimiter;
76
77    #[test]
78    fn new_clamps_zero_permits_to_one() {
79        assert_eq!(CpuLimiter::new(0).available_permits(), 1);
80        assert_eq!(CpuLimiter::new(4).available_permits(), 4);
81    }
82
83    #[test]
84    fn oversized_cpu_limit_does_not_panic() {
85        assert_eq!(
86            CpuLimiter::new(usize::MAX).available_permits(),
87            tokio::sync::Semaphore::MAX_PERMITS
88        );
89    }
90
91    #[tokio::test]
92    async fn clones_share_one_permit_pool() {
93        let a = CpuLimiter::new(1);
94        let b = a.clone();
95        let permit = a.acquire().await.unwrap();
96        // The clone draws from the same pool: no permits left.
97        assert_eq!(b.available_permits(), 0);
98        drop(permit);
99        assert_eq!(b.available_permits(), 1);
100    }
101}