config/
pool.rs

1use crate::{get_current_epoch, LuaConfig, LuaConfigInner};
2use kumo_prometheus::declare_metric;
3use parking_lot::FairMutex as Mutex;
4use std::collections::VecDeque;
5use std::sync::atomic::{AtomicUsize, Ordering};
6use std::sync::mpsc::{channel, Sender};
7use std::sync::LazyLock;
8use std::time::Duration;
9
10static POOL: LazyLock<Mutex<Pool>> = LazyLock::new(|| Mutex::new(Pool::new()));
11
12static DROPPER: LazyLock<Sender<LuaConfigInner>> = LazyLock::new(|| {
13    let (tx, rx) = channel::<LuaConfigInner>();
14    std::thread::Builder::new()
15        .name("lua context dropper".to_string())
16        .spawn(move || {
17            for context in rx {
18                drop(context);
19            }
20        })
21        .expect("create lua context dropper thread");
22    tx
23});
24
25declare_metric! {
26/// the number of lua contexts available for reuse in the pool
27static LUA_SPARE_COUNT: IntGauge("lua_spare_count");
28}
29
30/// Maximum age of a lua context before we release it, in seconds
31static MAX_AGE: AtomicUsize = AtomicUsize::new(300);
32/// Maximum number of uses of a given lua context before we release it
33static MAX_USE: AtomicUsize = AtomicUsize::new(1024);
34/// Maximum number of spare lua contexts to maintain in the pool
35static MAX_SPARE: AtomicUsize = AtomicUsize::new(8192);
36static GC_ON_PUT: AtomicUsize = AtomicUsize::new(0);
37
38pub fn set_max_use(max_use: usize) {
39    MAX_USE.store(max_use, Ordering::Relaxed);
40}
41
42pub fn set_max_spare(max_spare: usize) {
43    MAX_SPARE.store(max_spare, Ordering::Relaxed);
44}
45
46pub fn set_max_age(max_age: usize) {
47    MAX_AGE.store(max_age, Ordering::Relaxed);
48}
49
50/// Set the gc on put percentage chance, in the range 0-100
51pub fn set_gc_on_put(v: u8) {
52    GC_ON_PUT.store(v as usize, Ordering::Relaxed);
53}
54
55fn max_age() -> Duration {
56    Duration::from_secs(MAX_AGE.load(Ordering::Relaxed) as u64)
57}
58
59/// Retire a context by handing it to the dropper thread, which closes it
60/// asynchronously instead of making the caller pay for the close.
61fn discard(config: LuaConfigInner) {
62    if let Err(err) = DROPPER.send(config) {
63        drop(err.0);
64    }
65}
66
67fn discard_all(contexts: Vec<LuaConfigInner>) {
68    for context in contexts {
69        discard(context);
70    }
71}
72
73#[derive(Default)]
74pub(crate) struct Pool {
75    pool: VecDeque<LuaConfigInner>,
76}
77
78impl Pool {
79    pub fn new() -> Self {
80        std::thread::Builder::new()
81            .name("config idler".to_string())
82            .spawn(|| loop {
83                std::thread::sleep(Duration::from_secs(30));
84                let expired = POOL.lock().expire();
85                discard_all(expired);
86            })
87            .expect("create config idler thread");
88        Self::default()
89    }
90
91    /// Remove the contexts that have aged out or belong to an earlier
92    /// epoch, returning them for the caller to discard after releasing
93    /// the lock.
94    pub fn expire(&mut self) -> Vec<LuaConfigInner> {
95        let epoch = get_current_epoch();
96        let max_age = max_age();
97        let (keep, expired): (Vec<_>, Vec<_>) = std::mem::take(&mut self.pool)
98            .into_iter()
99            .partition(|inner| inner.created.elapsed() < max_age && inner.epoch == epoch);
100        self.pool = VecDeque::from(keep);
101        if !expired.is_empty() {
102            LUA_SPARE_COUNT.sub(expired.len() as i64);
103        }
104        expired
105    }
106
107    /// Returns a reusable context, if there is one, plus any stale contexts
108    /// found while looking. The caller must discard the stale contexts
109    /// after releasing the pool lock.
110    pub fn get(&mut self) -> (Option<LuaConfigInner>, Vec<LuaConfigInner>) {
111        let max_age = max_age();
112        let mut stale = vec![];
113        while let Some(mut item) = self.pool.pop_front() {
114            LUA_SPARE_COUNT.dec();
115            if item.created.elapsed() > max_age || item.epoch != get_current_epoch() {
116                stale.push(item);
117                continue;
118            }
119            item.use_count += 1;
120            return (Some(item), stale);
121        }
122        (None, stale)
123    }
124
125    /// Put a context back into the pool, or return it to the caller
126    /// if the pool is already at capacity.  The caller must discard the
127    /// returned context.
128    pub fn put(&mut self, config: LuaConfigInner) -> Option<LuaConfigInner> {
129        if self.pool.len() + 1 > MAX_SPARE.load(Ordering::Relaxed) {
130            return Some(config);
131        }
132        self.pool.push_back(config);
133        LUA_SPARE_COUNT.inc();
134        None
135    }
136}
137
138pub(crate) fn pool_get() -> Option<LuaConfig> {
139    let (item, stale) = POOL.lock().get();
140    discard_all(stale);
141    item.map(|inner| LuaConfig { inner: Some(inner) })
142}
143
144pub(crate) fn pool_put(config: LuaConfigInner) {
145    let reusable = config.created.elapsed() <= max_age()
146        && config.use_count + 1 <= MAX_USE.load(Ordering::Relaxed);
147    if !reusable {
148        discard(config);
149        return;
150    }
151
152    // gc can be slow, so it runs before the lock is taken.
153    let prob = GC_ON_PUT.load(Ordering::Relaxed);
154    if prob != 0 {
155        let chance = (rand::random::<f32>() * 100.) as usize;
156        if chance <= prob {
157            if let Err(err) = config.lua.gc_collect() {
158                tracing::error!("Error during gc: {err:#}");
159                discard(config);
160                return;
161            }
162        }
163    }
164
165    let mut pool = POOL.lock();
166    // The optional gc pass before this lock can take long enough for another
167    // thread to bump the epoch while it runs. This recheck catches a context
168    // that went stale during that gc.
169    if config.epoch != get_current_epoch() {
170        drop(pool);
171        discard(config);
172        return;
173    }
174    if let Some(config) = pool.put(config) {
175        drop(pool);
176        discard(config);
177    }
178}
179
180#[cfg(test)]
181mod test {
182    use super::*;
183    use crate::epoch::bump_current_epoch;
184    use mlua::Lua;
185    use std::time::Instant;
186
187    fn new_inner() -> LuaConfigInner {
188        LuaConfigInner {
189            lua: Lua::new(),
190            created: Instant::now(),
191            use_count: 1,
192            epoch: get_current_epoch(),
193        }
194    }
195
196    #[test]
197    fn get_hands_back_contexts_from_an_earlier_epoch() {
198        let mut pool = Pool::default();
199        assert!(pool.put(new_inner()).is_none());
200        assert!(pool.put(new_inner()).is_none());
201
202        bump_current_epoch();
203
204        let (item, stale) = pool.get();
205        assert!(item.is_none());
206        k9::assert_equal!(stale.len(), 2);
207        k9::assert_equal!(pool.pool.len(), 0);
208    }
209
210    #[test]
211    fn expire_hands_back_contexts_from_an_earlier_epoch() {
212        let mut pool = Pool::default();
213        assert!(pool.put(new_inner()).is_none());
214
215        bump_current_epoch();
216        assert!(pool.put(new_inner()).is_none());
217
218        let expired = pool.expire();
219        k9::assert_equal!(expired.len(), 1);
220        k9::assert_equal!(pool.pool.len(), 1);
221    }
222
223    /// Guard that restores the global MAX_SPARE on drop. Tests that lower
224    /// MAX_SPARE for a capacity check use it, because tests in the same
225    /// process share the global.
226    struct RestoreMaxSpare(usize);
227
228    impl Drop for RestoreMaxSpare {
229        fn drop(&mut self) {
230            set_max_spare(self.0);
231        }
232    }
233
234    #[test]
235    fn put_hands_back_the_context_when_the_pool_is_full() {
236        let _restore = RestoreMaxSpare(MAX_SPARE.load(Ordering::Relaxed));
237        set_max_spare(1);
238        let mut pool = Pool::default();
239        assert!(pool.put(new_inner()).is_none());
240        assert!(pool.put(new_inner()).is_some());
241    }
242}
config/pool.rs

config/
pool.rs
