fix(serverless): Axiom logger (#216)

.changeset/angry-fans-live.md

@@ -0,0 +1,5 @@
+---
+'@lagon/serverless': patch
+---
+
+Fix Axiom logger

packages/serverless/Cargo.toml

@@ -17,7 +17,7 @@ serde_json = "1.0"
 metrics = "0.20.1"
 metrics-exporter-prometheus = { version = "0.11.0", features = ["http-listener"] }
 log = { version = "0.4.17", features = ["std"] }
-axiom-rs = "0.5.0"
+axiom-rs = "0.6.0"
 chrono = "0.4.22"
 lazy_static = "1.4.0"
 rand = { version = "0.8.5", features = ["std_rng"] }

packages/serverless/src/logger.rs

@@ -2,43 +2,38 @@ use axiom_rs::Client;
 use chrono::prelude::Local;
 use flume::Sender;
 use serde_json::{json, Value};
+use std::sync::{Arc, RwLock};
 
 use log::{
-    set_boxed_logger, set_max_level, warn, Level, LevelFilter, Log, Metadata, Record,
-    SetLoggerError,
+    set_boxed_logger, set_max_level, Level, LevelFilter, Log, Metadata, Record, SetLoggerError,
 };
 struct SimpleLogger {
-    tx: Sender<Value>,
+    tx: Arc<RwLock<Option<Sender<Value>>>>,
 }
 
 impl SimpleLogger {
     pub fn new() -> Self {
         let (tx, rx) = flume::unbounded();
 
         // Axiom is optional
-        if let Ok(axiom_org_id) = dotenv::var("AXIOM_ORG_ID") {
-            let axiom_token = dotenv::var("AXIOM_TOKEN").expect("AXIOM_TOKEN must be set");
-
-            let axiom_client = Client::builder()
-                .with_org_id(axiom_org_id)
-                .with_token(axiom_token)
-                .build()
-                .expect("Failed to create Axiom client");
-
-            tokio::spawn(async move {
-                // TODO: batch values
-                let value = rx.recv().unwrap();
-                axiom_client
-                    .datasets
-                    .ingest("serverless", vec![value])
-                    .await
-                    .unwrap();
-            });
-        } else {
-            warn!("Axiom is not configured. Set AXIOM_ORG_ID to enable Axiom logging");
+        match Client::new() {
+            Ok(axiom_client) => {
+                tokio::spawn(async move {
+                    axiom_client
+                        .datasets
+                        .ingest_stream("serverless", rx.into_stream())
+                        .await
+                        .unwrap();
+                });
+            }
+            Err(e) => {
+                println!("Axiom is not configured: {}", e);
+            }
         }
 
-        Self { tx }
+        Self {
+            tx: Arc::new(RwLock::new(Some(tx))),
+        }
     }
 }
 
@@ -52,24 +47,36 @@ impl Log for SimpleLogger {
             println!("{} - {} - {}", Local::now(), record.level(), record.args());
 
             // Axiom is optional, so tx can have no listeners
-            if !self.tx.is_disconnected() {
-                self.tx
-                    .send(json!({
+            let tx = self.tx.read().expect("Tx lock is poisoned");
+            if let Some(tx) = &*tx {
+                if !tx.is_disconnected() {
+                    tx.send(json!({
                         "region": dotenv::var("LAGON_REGION").expect("LAGON_REGION must be set"),
-                        "timestamp": Local::now().to_rfc3339(),
+                        "_time": Local::now().to_rfc3339(),
                         "level": record.level().to_string(),
                         "message": record.args().to_string(),
                     }))
                     .unwrap_or(())
+                }
             }
         }
     }
 
     fn flush(&self) {
-        warn!("Flushing not implemented");
+        let mut tx = self.tx.write().expect("Tx lock is poisoned");
+        tx.take();
+    }
+}
+
+pub struct FlushGuard;
+
+impl Drop for FlushGuard {
+    fn drop(&mut self) {
+        log::logger().flush()
     }
 }
 
-pub fn init_logger() -> Result<(), SetLoggerError> {
-    set_boxed_logger(Box::new(SimpleLogger::new())).map(|()| set_max_level(LevelFilter::Info))
+pub fn init_logger() -> Result<FlushGuard, SetLoggerError> {
+    set_boxed_logger(Box::new(SimpleLogger::new())).map(|()| set_max_level(LevelFilter::Info))?;
+    Ok(FlushGuard)
 }

packages/serverless/src/main.rs

@@ -226,7 +226,7 @@ async fn handle_request(
 #[tokio::main]
 async fn main() {
     dotenv::dotenv().expect("Failed to load .env file");
-    init_logger().expect("Failed to init logger");
+    let _flush_guard = init_logger().expect("Failed to init logger");
 
     let runtime = Runtime::new(RuntimeOptions::default());
     let addr = SocketAddr::from(([0, 0, 0, 0], 4000));