mod_time/
lib.rs

1use config::{get_or_create_module, get_or_create_sub_module};
2use mlua::{Lua, MetaMethod, UserData, UserDataMethods};
3use prometheus::{HistogramTimer, HistogramVec};
4use std::sync::LazyLock;
5use tokio::time::Duration;
6
7static LATENCY_HIST: LazyLock<HistogramVec> = LazyLock::new(|| {
8    prometheus::register_histogram_vec!(
9        "user_lua_latency",
10        "how long something user-defined took to run in your lua policy",
11        &["label"]
12    )
13    .unwrap()
14});
15
16pub fn register(lua: &Lua) -> anyhow::Result<()> {
17    let kumo_mod = get_or_create_module(lua, "kumo")?;
18    let time_mod = get_or_create_sub_module(lua, "time")?;
19
20    let sleep_fn = lua.create_async_function(sleep)?;
21    kumo_mod.set("sleep", sleep_fn.clone())?;
22    time_mod.set("sleep", sleep_fn)?;
23
24    time_mod.set("start_timer", lua.create_function(Timer::start)?)?;
25
26    Ok(())
27}
28
29/// A Timer keeps track of the time since it was started,
30/// and will record the duration until its done method is
31/// called, or the __close metamethod is invoked.
32struct Timer {
33    timer: Option<HistogramTimer>,
34}
35
36impl Drop for Timer {
37    fn drop(&mut self) {
38        // We might be called some time after the code is done due
39        // to gc delays and pooling. We don't want the default
40        // Drop impl for HistogramTimer to record in that case:
41        // we will only report when our done method is explicitly
42        // called in lua
43        if let Some(timer) = self.timer.take() {
44            timer.stop_and_discard();
45        }
46    }
47}
48
49impl Timer {
50    fn start(_lua: &Lua, name: String) -> mlua::Result<Self> {
51        let timer = LATENCY_HIST
52            .get_metric_with_label_values(&[&name])
53            .expect("to get histo")
54            .start_timer();
55        Ok(Self { timer: Some(timer) })
56    }
57
58    fn done(_lua: &Lua, this: &mut Self, _: ()) -> mlua::Result<Option<f64>> {
59        Ok(this.timer.take().map(|timer| timer.stop_and_record()))
60    }
61}
62
63impl UserData for Timer {
64    fn add_methods<M: UserDataMethods<Self>>(methods: &mut M) {
65        methods.add_method_mut("done", Self::done);
66        methods.add_meta_method_mut(MetaMethod::Close, Self::done);
67    }
68}
69
70async fn sleep(_lua: Lua, seconds: f64) -> mlua::Result<()> {
71    tokio::time::sleep(Duration::from_secs_f64(seconds)).await;
72    Ok(())
73}
mod_time/lib.rs

mod_time/
lib.rs
