spacetimedb/util/
slow.rs

1use std::time::{Duration, Instant};
2
3use spacetimedb_datastore::execution_context::WorkloadType;
4
5/// Records the execution time of some `sql`
6/// and logs when the duration goes above a specific one.
7pub struct SlowQueryLogger<'a> {
8    /// The SQL statement of the query.
9    sql: &'a str,
10    /// The start time of the query execution.
11    start: Option<Instant>,
12    /// The threshold, if any, over which execution duration would result in logging.
13    threshold: Option<Duration>,
14    /// The context the query is being run in.
15    workload: WorkloadType,
16}
17
18impl<'a> SlowQueryLogger<'a> {
19    pub fn new(sql: &'a str, threshold: Option<Duration>, workload: WorkloadType) -> Self {
20        Self {
21            sql,
22            start: threshold.map(|_| Instant::now()),
23            threshold,
24            workload,
25        }
26    }
27
28    pub fn log_guard(self) -> impl Drop + 'a {
29        scopeguard::guard(self, |logger| {
30            logger.log();
31        })
32    }
33
34    /// Log as `tracing::warn!` the query if it exceeds the threshold.
35    pub fn log(&self) -> Option<Duration> {
36        if let Some((start, threshold)) = self.start.zip(self.threshold) {
37            let elapsed = start.elapsed();
38            if elapsed > threshold {
39                tracing::warn!(workload = %self.workload, ?threshold, ?elapsed, sql = ?self.sql, "SLOW QUERY");
40                return Some(elapsed);
41            }
42        }
43        None
44    }
45}
46
47#[cfg(test)]
48mod tests {
49    use super::*;
50
51    use crate::sql::compiler::compile_sql;
52    use crate::sql::execute::tests::execute_for_testing;
53    use spacetimedb_datastore::system_tables::ST_VARNAME_SLOW_QRY;
54    use spacetimedb_datastore::system_tables::{StVarName, ST_VARNAME_SLOW_INC, ST_VARNAME_SLOW_SUB};
55    use spacetimedb_lib::error::ResultTest;
56    use spacetimedb_lib::identity::AuthCtx;
57    use spacetimedb_lib::st_var::StVarValue;
58    use spacetimedb_lib::ProductValue;
59
60    use crate::db::relational_db::tests_utils::{begin_tx, insert, with_auto_commit, TestDB};
61    use crate::db::relational_db::RelationalDB;
62    use spacetimedb_sats::{product, AlgebraicType};
63    use spacetimedb_vm::relation::MemTable;
64
65    fn run_query(db: &RelationalDB, sql: String) -> ResultTest<MemTable> {
66        let tx = begin_tx(db);
67        let q = compile_sql(db, &AuthCtx::for_testing(), &tx, &sql)?;
68        Ok(execute_for_testing(db, &sql, q)?.pop().unwrap())
69    }
70
71    fn run_query_write(db: &RelationalDB, sql: String) -> ResultTest<()> {
72        let tx = begin_tx(db);
73        let q = compile_sql(db, &AuthCtx::for_testing(), &tx, &sql)?;
74        drop(tx);
75
76        execute_for_testing(db, &sql, q)?;
77        Ok(())
78    }
79
80    #[test]
81    fn test_slow_queries() -> ResultTest<()> {
82        let db = TestDB::in_memory()?.db;
83
84        let table_id =
85            db.create_table_for_test("test", &[("x", AlgebraicType::I32), ("y", AlgebraicType::I32)], &[])?;
86
87        with_auto_commit(&db, |tx| -> ResultTest<_> {
88            for i in 0..100_000 {
89                insert(&db, tx, table_id, &product![i, i * 2])?;
90            }
91            Ok(())
92        })?;
93        let tx = begin_tx(&db);
94
95        let sql = "select * from test where x > 0";
96        let q = compile_sql(&db, &AuthCtx::for_testing(), &tx, sql)?;
97
98        let slow = SlowQueryLogger::new(sql, Some(Duration::from_millis(1)), tx.ctx.workload());
99
100        let result = execute_for_testing(&db, sql, q)?;
101        assert_eq!(result[0].data[0], product![1, 2]);
102        assert!(slow.log().is_some());
103
104        Ok(())
105    }
106
107    // Verify we can change the threshold at runtime
108    #[test]
109    fn test_runtime_config() -> ResultTest<()> {
110        let db = TestDB::in_memory()?.db;
111
112        fn fetch_row(table: MemTable) -> Option<ProductValue> {
113            table.data.into_iter().next()
114        }
115
116        // Check we can read the default config
117        let row1 = fetch_row(run_query(&db, format!("SHOW {ST_VARNAME_SLOW_QRY}"))?);
118        let row2 = fetch_row(run_query(&db, format!("SHOW {ST_VARNAME_SLOW_SUB}"))?);
119        let row3 = fetch_row(run_query(&db, format!("SHOW {ST_VARNAME_SLOW_INC}"))?);
120
121        assert_eq!(row1, None);
122        assert_eq!(row2, None);
123        assert_eq!(row3, None);
124
125        // Check we can write a new config
126        run_query_write(&db, format!("SET {ST_VARNAME_SLOW_QRY} TO 1"))?;
127        run_query_write(&db, format!("SET {ST_VARNAME_SLOW_SUB} TO 1"))?;
128        run_query_write(&db, format!("SET {ST_VARNAME_SLOW_INC} TO 1"))?;
129
130        let row1 = fetch_row(run_query(&db, format!("SHOW {ST_VARNAME_SLOW_QRY}"))?);
131        let row2 = fetch_row(run_query(&db, format!("SHOW {ST_VARNAME_SLOW_SUB}"))?);
132        let row3 = fetch_row(run_query(&db, format!("SHOW {ST_VARNAME_SLOW_INC}"))?);
133
134        assert_eq!(row1, Some(product!(StVarName::SlowQryThreshold, StVarValue::U64(1))));
135        assert_eq!(row2, Some(product!(StVarName::SlowSubThreshold, StVarValue::U64(1))));
136        assert_eq!(row3, Some(product!(StVarName::SlowIncThreshold, StVarValue::U64(1))));
137
138        // And disable the config
139        run_query_write(&db, format!("DELETE FROM st_var WHERE name = '{ST_VARNAME_SLOW_QRY}'"))?;
140        run_query_write(&db, format!("DELETE FROM st_var WHERE name = '{ST_VARNAME_SLOW_SUB}'"))?;
141        run_query_write(&db, format!("DELETE FROM st_var WHERE name = '{ST_VARNAME_SLOW_INC}'"))?;
142
143        let row1 = fetch_row(run_query(&db, format!("SHOW {ST_VARNAME_SLOW_QRY}"))?);
144        let row2 = fetch_row(run_query(&db, format!("SHOW {ST_VARNAME_SLOW_SUB}"))?);
145        let row3 = fetch_row(run_query(&db, format!("SHOW {ST_VARNAME_SLOW_INC}"))?);
146
147        assert_eq!(row1, None);
148        assert_eq!(row2, None);
149        assert_eq!(row3, None);
150        Ok(())
151    }
152}
spacetimedb/util/slow.rs

spacetimedb/util/
slow.rs
