Skip to main content

aft/
cold_build_limiter.rs

1use std::sync::atomic::{AtomicUsize, Ordering};
2use std::sync::{Arc, LazyLock};
3use std::time::{Duration, Instant};
4
5#[cfg(not(test))]
6const DEFAULT_COLD_BUILD_LIMIT: usize = 2;
7#[cfg(test)]
8const DEFAULT_COLD_BUILD_LIMIT: usize = 1024;
9
10static GLOBAL_COLD_BUILD_LIMITER: LazyLock<Arc<ColdBuildLimiter>> =
11    LazyLock::new(|| Arc::new(ColdBuildLimiter::new(DEFAULT_COLD_BUILD_LIMIT)));
12
13pub fn try_acquire() -> Option<ColdBuildPermit> {
14    GLOBAL_COLD_BUILD_LIMITER.try_acquire()
15}
16
17/// Block until a build slot is free, then take it.
18///
19/// For build sites with no reschedule path (search-index builds spawn once per
20/// configure): skipping would strand the index, so past-cap work waits instead.
21/// Production captures showed concurrent per-root builds starving dispatch
22/// while CPU sat idle; waiting serializes that pressure at the source. Only
23/// call from dedicated background threads, never the dispatch thread or an
24/// executor worker.
25pub fn acquire_blocking(kind: &str) -> ColdBuildPermit {
26    let started = Instant::now();
27    let mut logged = false;
28    loop {
29        if let Some(permit) = GLOBAL_COLD_BUILD_LIMITER.try_acquire() {
30            if logged {
31                crate::slog_info!(
32                    "maintenance build slot acquired after {}ms wait: {}",
33                    started.elapsed().as_millis(),
34                    kind
35                );
36            }
37            return permit;
38        }
39        if !logged {
40            crate::slog_info!(
41                "maintenance build queued behind concurrency cap ({}): {}",
42                GLOBAL_COLD_BUILD_LIMITER.limit(),
43                kind
44            );
45            logged = true;
46        }
47        std::thread::sleep(Duration::from_millis(100));
48    }
49}
50
51pub fn limit() -> usize {
52    GLOBAL_COLD_BUILD_LIMITER.limit()
53}
54
55#[derive(Debug)]
56struct ColdBuildLimiter {
57    available: AtomicUsize,
58    limit: usize,
59}
60
61impl ColdBuildLimiter {
62    fn new(limit: usize) -> Self {
63        let limit = limit.max(1);
64        Self {
65            available: AtomicUsize::new(limit),
66            limit,
67        }
68    }
69
70    fn limit(&self) -> usize {
71        self.limit
72    }
73
74    fn try_acquire(self: &Arc<Self>) -> Option<ColdBuildPermit> {
75        loop {
76            let available = self.available.load(Ordering::Acquire);
77            if available == 0 {
78                return None;
79            }
80            if self
81                .available
82                .compare_exchange(
83                    available,
84                    available - 1,
85                    Ordering::AcqRel,
86                    Ordering::Acquire,
87                )
88                .is_ok()
89            {
90                return Some(ColdBuildPermit {
91                    limiter: Arc::clone(self),
92                });
93            }
94        }
95    }
96}
97
98#[derive(Debug)]
99pub struct ColdBuildPermit {
100    limiter: Arc<ColdBuildLimiter>,
101}
102
103impl Drop for ColdBuildPermit {
104    fn drop(&mut self) {
105        let previous = self.limiter.available.fetch_add(1, Ordering::Release);
106        debug_assert!(previous < self.limiter.limit);
107    }
108}
109
110#[cfg(test)]
111mod tests {
112    use super::*;
113
114    // Both tests mutate the process-global limiter; run them one at a time.
115    fn serial() -> std::sync::MutexGuard<'static, ()> {
116        static M: std::sync::OnceLock<std::sync::Mutex<()>> = std::sync::OnceLock::new();
117        M.get_or_init(|| std::sync::Mutex::new(()))
118            .lock()
119            .unwrap_or_else(std::sync::PoisonError::into_inner)
120    }
121
122    #[test]
123    fn permits_release_on_drop() {
124        let _serial = serial();
125        let before = GLOBAL_COLD_BUILD_LIMITER.available.load(Ordering::Acquire);
126        {
127            let _a = acquire_blocking("test-a");
128            let _b = acquire_blocking("test-b");
129            assert_eq!(
130                GLOBAL_COLD_BUILD_LIMITER.available.load(Ordering::Acquire),
131                before - 2
132            );
133        }
134        assert_eq!(
135            GLOBAL_COLD_BUILD_LIMITER.available.load(Ordering::Acquire),
136            before
137        );
138    }
139
140    #[test]
141    fn acquire_blocking_waits_until_release() {
142        let _serial = serial();
143        // Drain every slot, then prove a waiter blocks until one holder drops.
144        let mut held: Vec<ColdBuildPermit> = Vec::new();
145        while let Some(permit) = try_acquire() {
146            held.push(permit);
147        }
148        let waiter = std::thread::spawn(|| {
149            let _p = acquire_blocking("waiter");
150        });
151        std::thread::sleep(std::time::Duration::from_millis(250));
152        assert!(!waiter.is_finished(), "waiter must block while cap is full");
153        drop(held.pop());
154        waiter.join().expect("waiter finishes after release");
155        drop(held);
156    }
157}