Stream new logs from NATS

Author: max-lt

src/db.rs

@@ -1,6 +1,7 @@
 use chrono::{DateTime, Utc};
 use serde::Serialize;
 use sqlx::{postgres::PgPoolOptions, PgPool};
+use std::str::FromStr;
 use uuid::Uuid;
 
 #[derive(Clone, Debug, PartialEq, sqlx::Type, Serialize)]
@@ -14,6 +15,22 @@ pub enum LogLevel {
     Trace,
 }
 
+impl FromStr for LogLevel {
+    type Err = ();
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match s {
+            "error" => Ok(LogLevel::Error),
+            "warn" => Ok(LogLevel::Warn),
+            "info" => Ok(LogLevel::Info),
+            "log" => Ok(LogLevel::Log),
+            "debug" => Ok(LogLevel::Debug),
+            "trace" => Ok(LogLevel::Trace),
+            _ => Ok(LogLevel::Info), // Default to Info for unknown levels
+        }
+    }
+}
+
 #[derive(Debug, Clone, Serialize)]
 pub struct LogEntry {
     pub date: DateTime<Utc>,

src/main.rs

@@ -5,12 +5,10 @@ mod routes;
 use actix_web::{middleware, web, App, HttpServer};
 use chrono::Utc;
 use futures::StreamExt;
-use std::sync::Arc;
-use tokio::sync::broadcast;
 use uuid::Uuid;
 
 use db::{create_pool, insert_log, LogEntry, LogLevel};
-use routes::{get_worker_logs, health, stream_worker_logs, AppState};
+use routes::{health, stream_worker_logs, AppState};
 
 #[actix_web::main]
 async fn main() -> std::io::Result<()> {
@@ -22,21 +20,18 @@ async fn main() -> std::io::Result<()> {
     // Connect to database
     let pool = create_pool().await.expect("Failed to connect to database");
 
-    // Create broadcast channel for SSE
-    let (tx, _rx) = broadcast::channel::<LogEntry>(100);
-    let broadcaster = Arc::new(tx);
+    // Connect to NATS
+    let nats_client = nats::nats_connect().await;
 
-    // Spawn NATS subscriber task
+    // Spawn NATS subscriber task for database persistence
     let pool_clone = pool.clone();
-    let broadcaster_clone = broadcaster.clone();
+    let nats_clone = nats_client.clone();
 
     tokio::spawn(async move {
-        log::info!("Starting NATS subscriber...");
-
-        let nc = nats::nats_connect().await;
+        log::info!("Starting NATS subscriber for database persistence...");
 
         // Subscribe to all worker logs: *.console.*
-        let mut sub = nc
+        let mut sub = nats_clone
             .subscribe("*.console.*")
             .await
             .expect("Failed to subscribe to NATS");
@@ -63,18 +58,7 @@ async fn main() -> std::io::Result<()> {
                 }
             };
 
-            let level = match level_str {
-                "error" => LogLevel::Error,
-                "warn" => LogLevel::Warn,
-                "info" => LogLevel::Info,
-                "log" => LogLevel::Log,
-                "debug" => LogLevel::Debug,
-                "trace" => LogLevel::Trace,
-                _ => {
-                    log::warn!("Unknown log level: {}", level_str);
-                    LogLevel::Info
-                }
-            };
+            let level = level_str.parse().unwrap_or(LogLevel::Info);
 
             let message = match String::from_utf8(msg.payload.to_vec()) {
                 Ok(m) => m,
@@ -92,12 +76,9 @@ async fn main() -> std::io::Result<()> {
             };
 
             // Insert into database
-            if let Err(e) = insert_log(&pool_clone, log_entry.clone()).await {
+            if let Err(e) = insert_log(&pool_clone, log_entry).await {
                 log::error!("Failed to insert log: {:?}", e);
             }
-
-            // Broadcast to SSE subscribers
-            let _ = broadcaster_clone.send(log_entry);
         }
 
         log::warn!("NATS subscriber stopped");
@@ -115,11 +96,10 @@ async fn main() -> std::io::Result<()> {
         App::new()
             .app_data(web::Data::new(AppState {
                 pool: pool.clone(),
-                broadcaster: broadcaster.clone(),
+                nats_client: nats_client.clone(),
             }))
             .wrap(middleware::Logger::default())
             .service(health)
-            .service(get_worker_logs)
             .service(stream_worker_logs)
     })
     .bind(("0.0.0.0", port))?

src/routes.rs

@@ -1,63 +1,113 @@
 use actix_web::{get, web, HttpResponse, Responder};
 use actix_web_lab::sse::{self, Sse};
+use futures::StreamExt;
 use sqlx::PgPool;
-use std::sync::Arc;
-use tokio::sync::broadcast;
 use uuid::Uuid;
 
 use crate::db::{get_logs, LogEntry};
 
 pub struct AppState {
     pub pool: PgPool,
-    pub broadcaster: Arc<broadcast::Sender<LogEntry>>,
+    pub nats_client: async_nats::Client,
 }
 
-#[get("/health")]
+// Helper struct for creating SSE log events from raw data
+struct LogEventData {
+    date: i64,
+    level: String,
+    message: String,
+}
+
+impl Into<sse::Data> for LogEntry {
+    fn into(self) -> sse::Data {
+        sse::Data::new_json(serde_json::json!({
+            "date": self.date.timestamp_millis(),
+            "level": format!("{:?}", self.level).to_lowercase(),
+            "message": self.message
+        }))
+        .unwrap()
+        .event("log")
+    }
+}
+
+impl Into<sse::Data> for LogEventData {
+    fn into(self) -> sse::Data {
+        sse::Data::new_json(serde_json::json!({
+            "date": self.date,
+            "level": self.level,
+            "message": self.message
+        }))
+        .unwrap()
+        .event("log")
+    }
+}
+
+#[get("/api/v1/health")]
 pub async fn health() -> impl Responder {
     HttpResponse::Ok().json(serde_json::json!({
         "status": "ok",
         "service": "openworkers-logs"
     }))
 }
 
-#[get("/logs/{worker_id}")]
-pub async fn get_worker_logs(
+#[get("/api/v1/workers/{worker_id}/logs")]
+pub async fn stream_worker_logs(
     worker_id: web::Path<Uuid>,
     data: web::Data<AppState>,
 ) -> impl Responder {
-    match get_logs(&data.pool, *worker_id, 100).await {
-        Ok(logs) => HttpResponse::Ok().json(logs),
+    let worker_id = *worker_id;
+
+    // Fetch last 10 logs from database
+    let historical_logs = match get_logs(&data.pool, worker_id, 10).await {
+        Ok(logs) => logs,
         Err(e) => {
-            log::error!("Failed to fetch logs: {:?}", e);
-            HttpResponse::InternalServerError().json(serde_json::json!({
-                "error": "Failed to fetch logs"
-            }))
+            log::error!("Failed to fetch historical logs: {:?}", e);
+            vec![]
         }
+    };
+
+    // Subscribe to NATS for this specific worker's logs
+    let subject = format!("{}.console.>", worker_id);
+    let nats_sub = data.nats_client.subscribe(subject).await.ok();
+
+    if nats_sub.is_none() {
+        log::error!("Failed to subscribe to NATS");
     }
-}
 
-#[get("/logs/{worker_id}/stream")]
-pub async fn stream_worker_logs(
-    worker_id: web::Path<Uuid>,
-    data: web::Data<AppState>,
-) -> impl Responder {
-    let worker_id = *worker_id;
-    let mut rx = data.broadcaster.subscribe();
+    let mut nats_sub = nats_sub.unwrap();
+
+    let mut id_counter: u64 = 0;
 
     let stream = async_stream::stream! {
-        while let Ok(log) = rx.recv().await {
-            // Only send logs for this worker
-            if log.worker_id == worker_id {
-                let event = sse::Event::Data(
-                    sse::Data::new_json(serde_json::json!({
-                        "date": log.date,
-                        "message": log.message,
-                        "level": log.level
-                    }))
-                    .unwrap()
-                );
-                yield Ok::<_, actix_web::Error>(event);
-            }
+        // First, yield historical logs in reverse order (oldest first)
+        for log in historical_logs.into_iter().rev() {
+            let mut data: sse::Data = log.into();
+            data.set_id(format!("{}", id_counter));
+            id_counter += 1;
+            yield Ok::<_, actix_web::Error>(sse::Event::Data(data));
+        }
+
+        // Then stream new logs in real-time from NATS
+        while let Some(msg) = nats_sub.next().await {
+            // Parse level from subject: {worker_id}.console.{level}
+            let level_str = msg.subject.split('.').nth(2).unwrap_or("info");
+
+            let message = match String::from_utf8(msg.payload.to_vec()) {
+                Ok(m) => m,
+                Err(_) => continue,
+            };
+
+            let data = LogEventData {
+                date: chrono::Utc::now().timestamp_millis(),
+                level: level_str.to_string(),
+                message,
+            };
+
+            let mut data: sse::Data = data.into();
+            data.set_id(format!("{}", id_counter));
+            id_counter += 1;
+
+            yield Ok::<_, actix_web::Error>(sse::Event::Data(data));
         }
     };