From 8d29df10d66ef1e6ae434eea37747ae3abd2de25 Mon Sep 17 00:00:00 2001
From: Wez Furlong <wez@wezfurlong.org>
Date: Thu, 15 Aug 2024 14:47:39 -0700
Subject: [PATCH] config: record and export latency metrics for lua event
 callbacks

---
 Cargo.lock               |  2 ++
 crates/config/Cargo.toml |  2 ++
 crates/config/src/lib.rs | 38 +++++++++++++++++++++++++++++++++++---
 3 files changed, 39 insertions(+), 3 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 21e4325e..ae9166ae 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1066,7 +1066,9 @@ dependencies = [
  "lazy_static",
  "metrics",
  "mlua",
+ "once_cell",
  "parking_lot",
+ "prometheus",
  "serde",
  "serde_json",
  "tokio",
diff --git a/crates/config/Cargo.toml b/crates/config/Cargo.toml
index 6944d9e5..12446adf 100644
--- a/crates/config/Cargo.toml
+++ b/crates/config/Cargo.toml
@@ -10,7 +10,9 @@ anyhow = "1.0"
 lazy_static= "1.4"
 metrics = {workspace=true}
 mlua = {workspace=true, features=["vendored", "lua54", "async", "send", "serialize"]}
+once_cell = "1.17"
 parking_lot = "0.12"
+prometheus = "0.13"
 serde = "1.0"
 serde_json = "1.0"
 tokio = {workspace=true, features=["fs"]}
diff --git a/crates/config/src/lib.rs b/crates/config/src/lib.rs
index 15150fd1..a220fd62 100644
--- a/crates/config/src/lib.rs
+++ b/crates/config/src/lib.rs
@@ -2,7 +2,9 @@ use crate::pool::{pool_get, pool_put};
 pub use crate::pool::{set_max_age, set_max_spare, set_max_use};
 use anyhow::Context;
 use mlua::{FromLua, FromLuaMulti, IntoLuaMulti, Lua, LuaSerdeExt, RegistryKey, Table, Value};
+use once_cell::sync::Lazy;
 use parking_lot::FairMutex as Mutex;
+use prometheus::{HistogramTimer, HistogramVec};
 use serde::Serialize;
 use std::borrow::Cow;
 use std::collections::HashSet;
@@ -32,9 +34,24 @@ lazy_static::lazy_static! {
 
 pub static VALIDATE_ONLY: AtomicBool = AtomicBool::new(false);
 pub static VALIDATION_FAILED: AtomicBool = AtomicBool::new(false);
+static LATENCY_HIST: Lazy<HistogramVec> = Lazy::new(|| {
+    prometheus::register_histogram_vec!(
+        "lua_event_latency",
+        "how long a given lua event callback took",
+        &["event"]
+    )
+    .unwrap()
+});
 
 pub type RegisterFunc = fn(&Lua) -> anyhow::Result<()>;
 
+fn latency_timer(label: &str) -> HistogramTimer {
+    LATENCY_HIST
+        .get_metric_with_label_values(&[label])
+        .expect("to get histo")
+        .start_timer()
+}
+
 #[derive(Debug)]
 struct LuaConfigInner {
     lua: Lua,
@@ -135,6 +152,7 @@ pub async fn load_config() -> anyhow::Result<LuaConfig> {
             chunk.into_function()?
         };
 
+        let _timer = latency_timer("context-creation");
         func.call_async(()).await?;
     }
     LUA_COUNT.increment(1.);
@@ -178,7 +196,10 @@ impl LuaConfig {
             .lua
             .named_registry_value::<mlua::Function>(&decorated_name)
         {
-            Ok(func) => Ok(func.call_async(args).await?),
+            Ok(func) => {
+                let _timer = latency_timer(name);
+                Ok(func.call_async(args).await?)
+            }
             _ => anyhow::bail!("{name} has not been registered"),
         }
     }
@@ -205,6 +226,7 @@ impl LuaConfig {
                 Value::Table(tbl) => {
                     for func in tbl.sequence_values::<mlua::Function>() {
                         let func = func?;
+                        let _timer = latency_timer(name);
                         let result: mlua::MultiValue = func.call_async(args.clone()).await?;
                         if result.is_empty() {
                             // Continue with other handlers
@@ -223,7 +245,10 @@ impl LuaConfig {
             .lua
             .named_registry_value::<mlua::Function>(&decorated_name)
         {
-            Ok(func) => Ok(func.call_async(args.clone()).await?),
+            Ok(func) => {
+                let _timer = latency_timer(name);
+                Ok(func.call_async(args.clone()).await?)
+            }
             _ => Ok(R::default()),
         }
     }
@@ -250,6 +275,7 @@ impl LuaConfig {
                 Value::Table(tbl) => {
                     for func in tbl.sequence_values::<mlua::Function>() {
                         let func = func?;
+                        let _timer = latency_timer(name);
                         let result: mlua::MultiValue = func.call_async(args.clone()).await?;
                         if result.is_empty() {
                             // Continue with other handlers
@@ -268,7 +294,10 @@ impl LuaConfig {
             .lua
             .named_registry_value::<mlua::Function>(&decorated_name)
         {
-            Ok(func) => Ok(func.call_async(args.clone()).await?),
+            Ok(func) => {
+                let _timer = latency_timer(name);
+                Ok(func.call_async(args.clone()).await?)
+            }
             _ => anyhow::bail!("Event {name} has not been registered"),
         }
     }
@@ -295,6 +324,7 @@ impl LuaConfig {
                 Value::Table(tbl) => {
                     for func in tbl.sequence_values::<mlua::Function>() {
                         let func = func?;
+                        let _timer = latency_timer(name);
                         let result: mlua::MultiValue = func.call_async(args.clone()).await?;
                         if result.is_empty() {
                             // Continue with other handlers
@@ -314,6 +344,7 @@ impl LuaConfig {
         match opt_func {
             Value::Nil => Ok(None),
             Value::Function(func) => {
+                let _timer = latency_timer(name);
                 let value: Value = func.call_async(args.clone()).await?;
 
                 match value {
@@ -359,6 +390,7 @@ impl LuaConfig {
             .lua
             .named_registry_value::<mlua::Function>(&decorated_name)?;
 
+        let _timer = latency_timer(name);
         let value: Value = func.call_async(args.clone()).await?;
         drop(func);