hyperactor_telemetry/
sqlite.rs

1/*
2 * Copyright (c) Meta Platforms, Inc. and affiliates.
3 * All rights reserved.
4 *
5 * This source code is licensed under the BSD-style license found in the
6 * LICENSE file in the root directory of this source tree.
7 */
8
9use std::collections::HashMap;
10use std::fs;
11use std::os::unix::fs::PermissionsExt;
12use std::path::PathBuf;
13use std::sync::Arc;
14use std::sync::LazyLock;
15use std::sync::Mutex;
16
17use anyhow::Result;
18use anyhow::anyhow;
19use rusqlite::Connection;
20use rusqlite::functions::FunctionFlags;
21use serde::Serialize;
22use serde_json::Value as JValue;
23use serde_rusqlite::*;
24use tracing::Event;
25use tracing::Subscriber;
26use tracing_subscriber::Layer;
27use tracing_subscriber::Registry;
28use tracing_subscriber::prelude::*;
29use tracing_subscriber::reload;
30
31pub type SqliteReloadHandle = reload::Handle<Option<SqliteLayer>, Registry>;
32
33// Reload handle allows us to include a no-op layer during init, but load
34// the layer dynamically during tests.
35static RELOAD_HANDLE: Mutex<Option<SqliteReloadHandle>> = Mutex::new(None);
36
37pub trait TableDef {
38    fn name(&self) -> &'static str;
39    fn columns(&self) -> &'static [&'static str];
40    fn create_table_stmt(&self) -> String {
41        let name = self.name();
42        let columns = self
43            .columns()
44            .iter()
45            .map(|col| format!("{col} TEXT "))
46            .collect::<Vec<String>>()
47            .join(",");
48        format!("create table if not exists {name} (seq INTEGER primary key, {columns})")
49    }
50    fn insert_stmt(&self) -> String {
51        let name = self.name();
52        let columns = self.columns().join(", ");
53        let params = self
54            .columns()
55            .iter()
56            .map(|c| format!(":{c}"))
57            .collect::<Vec<String>>()
58            .join(", ");
59        format!("insert into {name} ({columns}) values ({params})")
60    }
61}
62
63impl TableDef for (&'static str, &'static [&'static str]) {
64    fn name(&self) -> &'static str {
65        self.0
66    }
67
68    fn columns(&self) -> &'static [&'static str] {
69        self.1
70    }
71}
72
73#[derive(Clone, Debug)]
74pub struct Table {
75    pub columns: &'static [&'static str],
76    pub create_table_stmt: String,
77    pub insert_stmt: String,
78}
79
80impl From<(&'static str, &'static [&'static str])> for Table {
81    fn from(value: (&'static str, &'static [&'static str])) -> Self {
82        Self {
83            columns: value.columns(),
84            create_table_stmt: value.create_table_stmt(),
85            insert_stmt: value.insert_stmt(),
86        }
87    }
88}
89
90#[derive(Debug, Clone, Copy, PartialEq, Eq)]
91pub enum TableName {
92    ActorLifecycle,
93    Messages,
94    LogEvents,
95}
96
97impl TableName {
98    pub const ACTOR_LIFECYCLE_STR: &'static str = "actor_lifecycle";
99    pub const MESSAGES_STR: &'static str = "messages";
100    pub const LOG_EVENTS_STR: &'static str = "log_events";
101
102    pub fn as_str(&self) -> &'static str {
103        match self {
104            TableName::ActorLifecycle => Self::ACTOR_LIFECYCLE_STR,
105            TableName::Messages => Self::MESSAGES_STR,
106            TableName::LogEvents => Self::LOG_EVENTS_STR,
107        }
108    }
109
110    pub fn get_table(&self) -> &'static Table {
111        match self {
112            TableName::ActorLifecycle => &ACTOR_LIFECYCLE,
113            TableName::Messages => &MESSAGES,
114            TableName::LogEvents => &LOG_EVENTS,
115        }
116    }
117}
118
119static ACTOR_LIFECYCLE: LazyLock<Table> = LazyLock::new(|| {
120    (
121        TableName::ActorLifecycle.as_str(),
122        [
123            "actor_id",
124            "actor",
125            "name",
126            "supervised_actor",
127            "actor_status",
128            "module_path",
129            "line",
130            "file",
131        ]
132        .as_slice(),
133    )
134        .into()
135});
136
137static MESSAGES: LazyLock<Table> = LazyLock::new(|| {
138    (
139        TableName::Messages.as_str(),
140        [
141            "span_id",
142            "time_us",
143            "src",
144            "dest",
145            "payload",
146            "module_path",
147            "line",
148            "file",
149        ]
150        .as_slice(),
151    )
152        .into()
153});
154
155static LOG_EVENTS: LazyLock<Table> = LazyLock::new(|| {
156    (
157        TableName::LogEvents.as_str(),
158        [
159            "span_id",
160            "time_us",
161            "name",
162            "message",
163            "actor_id",
164            "level",
165            "line",
166            "file",
167            "module_path",
168        ]
169        .as_slice(),
170    )
171        .into()
172});
173
174pub static ALL_TABLES: LazyLock<Vec<Table>> = LazyLock::new(|| {
175    vec![
176        ACTOR_LIFECYCLE.clone(),
177        MESSAGES.clone(),
178        LOG_EVENTS.clone(),
179    ]
180});
181
182pub struct SqliteLayer {
183    conn: Arc<Mutex<Connection>>,
184}
185use tracing::field::Visit;
186
187#[derive(Debug, Clone, Default, Serialize)]
188pub struct SqlVisitor(pub HashMap<String, JValue>);
189
190impl Visit for SqlVisitor {
191    fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) {
192        self.0.insert(
193            field.name().to_string(),
194            JValue::String(format!("{:?}", value)),
195        );
196    }
197
198    fn record_str(&mut self, field: &tracing::field::Field, value: &str) {
199        self.0
200            .insert(field.name().to_string(), JValue::String(value.to_string()));
201    }
202
203    fn record_i64(&mut self, field: &tracing::field::Field, value: i64) {
204        self.0
205            .insert(field.name().to_string(), JValue::Number(value.into()));
206    }
207
208    fn record_f64(&mut self, field: &tracing::field::Field, value: f64) {
209        let n = serde_json::Number::from_f64(value).unwrap();
210        self.0.insert(field.name().to_string(), JValue::Number(n));
211    }
212
213    fn record_u64(&mut self, field: &tracing::field::Field, value: u64) {
214        self.0
215            .insert(field.name().to_string(), JValue::Number(value.into()));
216    }
217
218    fn record_bool(&mut self, field: &tracing::field::Field, value: bool) {
219        self.0.insert(field.name().to_string(), JValue::Bool(value));
220    }
221}
222
223macro_rules! insert_event {
224    ($table:expr, $conn:ident, $event:ident) => {
225        let mut v: SqlVisitor = Default::default();
226        $event.record(&mut v);
227        let meta = $event.metadata();
228        v.0.insert(
229            "module_path".to_string(),
230            meta.module_path().map(String::from).into(),
231        );
232        v.0.insert("line".to_string(), meta.line().into());
233        v.0.insert("file".to_string(), meta.file().map(String::from).into());
234        $conn.prepare_cached(&$table.insert_stmt)?.execute(
235            serde_rusqlite::to_params_named_with_fields(v, $table.columns)?
236                .to_slice()
237                .as_slice(),
238        )?;
239    };
240}
241
242/// Public helper to insert event fields into database using the same logic as the old implementation.
243/// This is used by the unified SqliteExporter to ensure identical behavior.
244pub fn insert_event_fields(conn: &Connection, table: &Table, fields: SqlVisitor) -> Result<()> {
245    conn.prepare_cached(&table.insert_stmt)?.execute(
246        serde_rusqlite::to_params_named_with_fields(fields, table.columns)?
247            .to_slice()
248            .as_slice(),
249    )?;
250    Ok(())
251}
252
253impl SqliteLayer {
254    pub fn new() -> Result<Self> {
255        let conn = Connection::open_in_memory()?;
256        Self::setup_connection(conn)
257    }
258
259    pub fn new_with_file(db_path: &str) -> Result<Self> {
260        let conn = Connection::open(db_path)?;
261        Self::setup_connection(conn)
262    }
263
264    fn setup_connection(conn: Connection) -> Result<Self> {
265        for table in ALL_TABLES.iter() {
266            conn.execute(&table.create_table_stmt, [])?;
267        }
268        conn.create_scalar_function(
269            "assert",
270            2,
271            FunctionFlags::SQLITE_UTF8 | FunctionFlags::SQLITE_DETERMINISTIC,
272            move |ctx| {
273                let condition: bool = ctx.get(0)?;
274                let message: String = ctx.get(1)?;
275
276                if !condition {
277                    return Err(rusqlite::Error::UserFunctionError(
278                        anyhow!("assertion failed:{condition} {message}",).into(),
279                    ));
280                }
281
282                Ok(condition)
283            },
284        )?;
285
286        Ok(Self {
287            conn: Arc::new(Mutex::new(conn)),
288        })
289    }
290
291    fn insert_event(&self, event: &Event<'_>) -> Result<()> {
292        let conn = self.conn.lock().unwrap();
293        match (event.metadata().target(), event.metadata().name()) {
294            (TableName::MESSAGES_STR, _) => {
295                insert_event!(TableName::Messages.get_table(), conn, event);
296            }
297            (TableName::ACTOR_LIFECYCLE_STR, _) => {
298                insert_event!(TableName::ActorLifecycle.get_table(), conn, event);
299            }
300            _ => {
301                insert_event!(TableName::LogEvents.get_table(), conn, event);
302            }
303        }
304        Ok(())
305    }
306
307    pub fn connection(&self) -> Arc<Mutex<Connection>> {
308        self.conn.clone()
309    }
310}
311
312impl<S: Subscriber> Layer<S> for SqliteLayer {
313    fn on_event(&self, event: &Event<'_>, _ctx: tracing_subscriber::layer::Context<'_, S>) {
314        self.insert_event(event).unwrap();
315    }
316}
317
318#[allow(dead_code)]
319fn print_table(conn: &Connection, table_name: TableName) -> Result<()> {
320    let table_name_str = table_name.as_str();
321
322    // Get column names
323    let mut stmt = conn.prepare(&format!("PRAGMA table_info({})", table_name_str))?;
324    let column_info = stmt.query_map([], |row| {
325        row.get::<_, String>(1) // Column name is at index 1
326    })?;
327
328    let columns: Vec<String> = column_info.collect::<Result<Vec<_>, _>>()?;
329
330    // Print header
331    println!("=== {} ===", table_name_str.to_uppercase());
332    println!("{}", columns.join(" | "));
333    println!("{}", "-".repeat(columns.len() * 10));
334
335    // Print rows
336    let mut stmt = conn.prepare(&format!("SELECT * FROM {}", table_name_str))?;
337    let rows = stmt.query_map([], |row| {
338        let mut values = Vec::new();
339        for (i, column) in columns.iter().enumerate() {
340            // Handle different column types properly
341            let value = if i == 0 && *column == "seq" {
342                // First column is always the INTEGER seq column
343                match row.get::<_, Option<i64>>(i)? {
344                    Some(v) => v.to_string(),
345                    None => "NULL".to_string(),
346                }
347            } else {
348                // All other columns are TEXT
349                match row.get::<_, Option<String>>(i)? {
350                    Some(v) => v,
351                    None => "NULL".to_string(),
352                }
353            };
354            values.push(value);
355        }
356        Ok(values.join(" | "))
357    })?;
358
359    for row in rows {
360        println!("{}", row?);
361    }
362    println!();
363    Ok(())
364}
365
366fn init_tracing_subscriber(layer: SqliteLayer) {
367    let handle = RELOAD_HANDLE.lock().unwrap();
368    if let Some(reload_handle) = handle.as_ref() {
369        let _ = reload_handle.reload(layer);
370    } else {
371        tracing_subscriber::registry().with(layer).init();
372    }
373}
374
375// === API ===
376
377// Creates a new reload handler and no-op layer for initialization
378pub fn get_reloadable_sqlite_layer() -> Result<reload::Layer<Option<SqliteLayer>, Registry>> {
379    let (layer, reload_handle) = reload::Layer::new(None);
380    let mut handle = RELOAD_HANDLE.lock().unwrap();
381    *handle = Some(reload_handle);
382    Ok(layer)
383}
384
385/// RAII guard for SQLite tracing database
386pub struct SqliteTracing {
387    db_path: Option<PathBuf>,
388    connection: Arc<Mutex<Connection>>,
389}
390
391impl SqliteTracing {
392    /// Create a new SqliteTracing with a temporary file
393    pub fn new() -> Result<Self> {
394        let temp_dir = std::env::temp_dir();
395        let file_name = format!("hyperactor_trace_{}.db", std::process::id());
396        let db_path = temp_dir.join(file_name);
397
398        let db_path_str = db_path.to_string_lossy();
399        let layer = SqliteLayer::new_with_file(&db_path_str)?;
400        let connection = layer.connection();
401
402        // Set file permissions to be readable and writable by owner and group
403        // This ensures the Python application can access the database file
404        if let Ok(metadata) = fs::metadata(&db_path) {
405            let mut permissions = metadata.permissions();
406            permissions.set_mode(0o664); // rw-rw-r--
407            let _ = fs::set_permissions(&db_path, permissions);
408        }
409
410        init_tracing_subscriber(layer);
411
412        Ok(Self {
413            db_path: Some(db_path),
414            connection,
415        })
416    }
417
418    /// Create a new SqliteTracing with in-memory database
419    pub fn new_in_memory() -> Result<Self> {
420        let layer = SqliteLayer::new()?;
421        let connection = layer.connection();
422
423        init_tracing_subscriber(layer);
424
425        Ok(Self {
426            db_path: None,
427            connection,
428        })
429    }
430
431    /// Get the path to the temporary database file (None for in-memory)
432    pub fn db_path(&self) -> Option<&PathBuf> {
433        self.db_path.as_ref()
434    }
435
436    /// Get a reference to the database connection
437    pub fn connection(&self) -> Arc<Mutex<Connection>> {
438        self.connection.clone()
439    }
440}
441
442impl Drop for SqliteTracing {
443    fn drop(&mut self) {
444        // Reset the layer to None
445        let handle = RELOAD_HANDLE.lock().unwrap();
446        if let Some(reload_handle) = handle.as_ref() {
447            let _ = reload_handle.reload(None);
448        }
449
450        // Delete the temporary file if it exists
451        if let Some(db_path) = &self.db_path
452            && db_path.exists()
453        {
454            let _ = fs::remove_file(db_path);
455        }
456    }
457}
458
459#[cfg(test)]
460mod tests {
461    use tracing::info;
462
463    use super::*;
464
465    #[test]
466    fn test_sqlite_tracing_with_file() -> Result<()> {
467        let tracing = SqliteTracing::new()?;
468        let conn = tracing.connection();
469
470        info!(target:"messages", test_field = "test_value", "Test msg");
471        info!(target:"log_events", test_field = "test_value", "Test event");
472
473        let count: i64 =
474            conn.lock()
475                .unwrap()
476                .query_row("SELECT COUNT(*) FROM messages", [], |row| row.get(0))?;
477        print_table(&conn.lock().unwrap(), TableName::LogEvents)?;
478        assert!(count > 0);
479
480        // Verify we have a file path
481        assert!(tracing.db_path().is_some());
482        let db_path = tracing.db_path().unwrap();
483        assert!(db_path.exists());
484
485        Ok(())
486    }
487
488    #[test]
489    fn test_sqlite_tracing_in_memory() -> Result<()> {
490        let tracing = SqliteTracing::new_in_memory()?;
491        let conn = tracing.connection();
492
493        info!(target:"messages", test_field = "test_value", "Test event in memory");
494
495        let count: i64 =
496            conn.lock()
497                .unwrap()
498                .query_row("SELECT COUNT(*) FROM messages", [], |row| row.get(0))?;
499        print_table(&conn.lock().unwrap(), TableName::Messages)?;
500        assert!(count > 0);
501
502        // Verify we don't have a file path for in-memory
503        assert!(tracing.db_path().is_none());
504
505        Ok(())
506    }
507
508    #[test]
509    fn test_sqlite_tracing_cleanup() -> Result<()> {
510        let db_path = {
511            let tracing = SqliteTracing::new()?;
512            let conn = tracing.connection();
513
514            info!(target:"log_events", test_field = "cleanup_test", "Test cleanup event");
515
516            let count: i64 =
517                conn.lock()
518                    .unwrap()
519                    .query_row("SELECT COUNT(*) FROM log_events", [], |row| row.get(0))?;
520            assert!(count > 0);
521
522            tracing.db_path().unwrap().clone()
523        }; // tracing goes out of scope here, triggering Drop
524
525        // File should be cleaned up after Drop
526        assert!(!db_path.exists());
527
528        Ok(())
529    }
530
531    #[test]
532    fn test_sqlite_tracing_different_targets() -> Result<()> {
533        let tracing = SqliteTracing::new_in_memory()?;
534        let conn = tracing.connection();
535
536        // Test different event targets
537        info!(target:"messages", src = "actor1", dest = "actor2", payload = "test_message", "Message event");
538        info!(target:"actor_lifecycle", actor_id = "123", actor = "TestActor", name = "test", "Lifecycle event");
539        info!(target:"log_events", test_field = "general_event", "General event");
540
541        // Check that events went to the right tables
542        let message_count: i64 =
543            conn.lock()
544                .unwrap()
545                .query_row("SELECT COUNT(*) FROM messages", [], |row| row.get(0))?;
546        assert_eq!(message_count, 1);
547
548        let lifecycle_count: i64 =
549            conn.lock()
550                .unwrap()
551                .query_row("SELECT COUNT(*) FROM actor_lifecycle", [], |row| row.get(0))?;
552        assert_eq!(lifecycle_count, 1);
553
554        let events_count: i64 =
555            conn.lock()
556                .unwrap()
557                .query_row("SELECT COUNT(*) FROM log_events", [], |row| row.get(0))?;
558        assert_eq!(events_count, 1);
559
560        Ok(())
561    }
562}
hyperactor_telemetry/sqlite.rs

hyperactor_telemetry/
sqlite.rs
