feat: refactor event processing / auth middleware

Author: explodingcamera

Cargo.toml

@@ -99,7 +99,7 @@ axum-test={version="20"}
 [features]
 default=["geoip"]
 geoip=["dep:maxminddb"]
-_enable_seeding=[]
+__dev=[]
 
 [profile.dev]
 opt-level=1

data/licenses-npm.json

@@ -1,13 +1,13 @@
 use std::sync::Arc;
 
-use anyhow::{Result, bail};
+use anyhow::{Context, Result};
 use arc_swap::ArcSwap;
 use chrono::{DateTime, Utc};
 use rand::distr::{SampleString, StandardUniform};
-use std::sync::mpsc::Receiver;
+use tokio::sync::mpsc::Receiver;
 
 use crate::app::models::{Event, event_params};
-use crate::app::{DuckDBPool, EVENT_BATCH_INTERVAL, SqlitePool};
+use crate::app::{DuckDBPool, SqlitePool};
 
 #[derive(Clone)]
 pub struct LiwanEvents {
@@ -48,49 +48,50 @@ impl LiwanEvents {
     /// Append events in batch
     pub fn append(&self, events: impl Iterator<Item = Event>) -> Result<()> {
         let conn = self.duckdb.get()?;
-        let mut appender = conn.appender("events")?;
+        let mut appender = conn.appender("events").context("Failed to get DuckDB appender")?;
         let mut first_event_time = Utc::now();
         for event in events {
-            appender.append_row(event_params![event])?;
-            if first_event_time > event.created_at {
+            if event.created_at < first_event_time {
                 first_event_time = event.created_at;
             }
+            appender.append_row(event_params![event]).context("Failed to append event to DuckDB")?;
         }
-        appender.flush()?;
-        update_event_times(&conn, first_event_time)?;
+        appender.flush().context("Failed to flush events to DuckDB")?;
+        update_event_times(&conn, first_event_time).context("Failed to update event times in DuckDB")?;
         Ok(())
     }
 
     /// Start processing events from the given channel. Blocks until the channel is closed.
-    pub fn process(&self, events: Receiver<Event>) -> Result<()> {
-        let conn = self.duckdb.get()?;
+    pub async fn process(&self, mut events: Receiver<Event>) -> Result<()> {
+        let mut buffer = Vec::with_capacity(1024);
 
         loop {
-            match events.recv() {
-                Ok(event) => {
-                    let mut appender = conn.appender("events")?;
-                    let mut first_event_time = event.created_at;
-                    appender.append_row(event_params![event])?;
-
-                    // Non-blockingly drain the remaining events in the queue if there are any
-                    let mut count = 1;
-                    for event in events.try_iter() {
-                        appender.append_row(event_params![event])?;
-                        count += 1;
-
-                        if first_event_time > event.created_at {
-                            first_event_time = event.created_at;
-                        }
-                    }
-
-                    appender.flush()?;
-                    update_event_times(&conn, first_event_time)?;
-                    tracing::debug!("Processed {} events", count);
-
-                    // Sleep to allow more events to be received before the next batch
-                    std::thread::sleep(EVENT_BATCH_INTERVAL);
+            let count = events.recv_many(&mut buffer, 1024).await;
+
+            if count == 0 {
+                tracing::info!("Event channel closed, stopping event processing");
+                break Ok(());
+            }
+
+            let first_event_time = buffer.first().map(|e| e.created_at).unwrap_or_else(Utc::now);
+            let events = std::mem::take(&mut buffer).into_iter();
+
+            let conn = self.duckdb.clone();
+            let res = tokio::task::spawn_blocking(move || {
+                let conn = conn.get().context("Failed to get DuckDB connection")?;
+                let mut appender = conn.appender("events").context("Failed to get DuckDB appender")?;
+                for event in events {
+                    appender.append_row(event_params![event]).context("Failed to append event to DuckDB")?;
                 }
-                Err(_) => bail!("event channel closed"),
+                appender.flush().context("Failed to flush events to DuckDB")?;
+                update_event_times(&conn, first_event_time)?;
+                anyhow::Ok(())
+            })
+            .await?;
+
+            match res {
+                Err(err) => tracing::error!("Event processing task panicked: {:?}", err),
+                _ => tracing::debug!("Processed {} events", count),
             }
         }
     }

src/app/core/events.rs

@@ -38,8 +38,6 @@ mod embedded {
     pub(super) mod events { refinery::embed_migrations!("src/migrations/events"); }
 }
 
-const EVENT_BATCH_INTERVAL: std::time::Duration = std::time::Duration::from_secs(5);
-
 impl Liwan {
     pub fn try_new(config: Config) -> Result<Arc<Self>> {
         tracing::debug!("Initializing app");
@@ -114,7 +112,7 @@ impl Liwan {
     }
 }
 
-#[cfg(any(debug_assertions, test, feature = "_enable_seeding"))]
+#[cfg(any(debug_assertions, test, feature = "__dev"))]
 impl Liwan {
     pub fn seed_database(&self, count_per_entity: usize) -> Result<()> {
         use chrono::{Days, Utc};

src/app/mod.rs

@@ -32,7 +32,7 @@ pub struct Project {
     pub id: String,
     pub display_name: String,
     pub public: bool,
-    pub secret: Option<String>, // enable public access with password protection
+    pub secret: Option<String>, // currently unused
 }
 
 #[derive(Debug, Clone)]

src/app/models.rs

@@ -31,11 +31,11 @@ pub enum Command {
     UpdatePassword(UpdatePassword),
     AddUser(AddUser),
     Users(ListUsers),
-    #[cfg(any(debug_assertions, test, feature = "_enable_seeding"))]
+    #[cfg(any(debug_assertions, test, feature = "__dev"))]
     SeedDatabase(SeedDatabase),
 }
 
-#[cfg(any(debug_assertions, test, feature = "_enable_seeding"))]
+#[cfg(any(debug_assertions, test, feature = "__dev"))]
 #[derive(FromArgs)]
 #[argh(subcommand, name = "seed-database")]
 /// Seed the database with some test data
@@ -129,10 +129,10 @@ pub fn handle_command(mut config: Config, cmd: Command) -> Result<()> {
             std::fs::write(&output, DEFAULT_CONFIG)?;
             println!("Configuration file written to liwan.config.toml");
         }
-        #[cfg(any(debug_assertions, test, feature = "_enable_seeding"))]
+        #[cfg(any(debug_assertions, test, feature = "__dev"))]
         Command::SeedDatabase(_) => {
             let app = Liwan::try_new(config)?;
-            app.seed_database(1_000_000)?;
+            app.seed_database(10_000_000)?;
             println!("Database seeded with test data");
         }
     }

src/cli.rs

@@ -16,7 +16,7 @@ async fn main() -> Result<()> {
     setup_logger(args.log_level)?;
 
     let config = Config::load(args.config)?;
-    let (s, r) = std::sync::mpsc::channel::<Event>();
+    let (s, r) = tokio::sync::mpsc::channel::<Event>(1024 * 10);
 
     if let Some(cmd) = args.cmd {
         return cli::handle_command(config, cmd);
@@ -30,7 +30,7 @@ async fn main() -> Result<()> {
         biased;
         _ = liwan::utils::signals::shutdown() => app_copy.shutdown(),
         res = web::start_webserver(app.clone(), s) => res,
-        res = tokio::task::spawn_blocking(move || app.clone().events.process(r)) => res?
+        res = tokio::task::spawn(async move { app.events.process(r).await }) => res?,
     }
 }
 

src/main.rs

@@ -1,8 +1,4 @@
-use crate::{
-    app::models::{Project, UserRole},
-    web::SessionUser,
-};
-
+use crate::app::models::{Project, User, UserRole};
 pub const MAX_DATAPOINTS: u32 = 2000;
 
 pub fn is_valid_id(id: &str) -> bool {
@@ -13,8 +9,8 @@ pub fn is_valid_username(name: &str) -> bool {
     name.chars().all(|c| c.is_alphanumeric() || c == '-' || c == '_') && name.len() <= 64 && name.len() >= 3
 }
 
-pub fn can_access_project(project: &Project, user: Option<&SessionUser>) -> bool {
-    project.public || user.is_some_and(|u| u.0.role == UserRole::Admin || u.0.projects.contains(&project.id))
+pub fn can_access_project(project: &Project, user: Option<&User>) -> bool {
+    project.public || user.is_some_and(|u| u.role == UserRole::Admin || u.projects.contains(&project.id))
 }
 
 #[cfg(test)]
@@ -32,11 +28,7 @@ mod tests {
         };
         assert!(can_access_project(&project, None), "Public project should be accessible without a user.");
 
-        let user = SessionUser(User {
-            username: "test".to_string(),
-            role: UserRole::User,
-            projects: vec!["other".to_string()],
-        });
+        let user = User { username: "test".to_string(), role: UserRole::User, projects: vec!["other".to_string()] };
         assert!(can_access_project(&project, Some(&user)), "Public project should be accessible with any user.");
 
         let project = Project {
@@ -45,7 +37,7 @@ mod tests {
             secret: None,
             public: false,
         };
-        let admin_user = SessionUser(User { username: "admin".to_string(), role: UserRole::Admin, projects: vec![] });
+        let admin_user = User { username: "admin".to_string(), role: UserRole::Admin, projects: vec![] };
         assert!(can_access_project(&project, Some(&admin_user)), "Admin should have access to any project.");
 
         let project = Project {

src/utils/validate.rs

@@ -4,14 +4,15 @@ pub mod webext;
 
 use std::net::{SocketAddr, ToSocketAddrs};
 use std::ops::Deref;
-use std::sync::{Arc, mpsc::Sender};
+use std::sync::Arc;
 
 use anyhow::{Context, Result};
 use axum::handler::{Handler, HandlerWithoutStateExt};
 use rust_embed::RustEmbed;
 
 use aide::{axum::ApiRouter, openapi};
 use http::{HeaderValue, Method, header};
+use tokio::sync::mpsc::Sender;
 use tower_http::{
     compression::CompressionLayer,
     cors::{Any, CorsLayer},
@@ -21,7 +22,7 @@ use tower_http::{
 use crate::app::{Liwan, models::Event};
 use crate::web::webext::serve;
 
-pub use session::{MaybeExtract, SessionId, SessionUser};
+pub use session::MaybeSessionId;
 use webext::StaticFile;
 
 #[derive(RustEmbed, Clone)]

src/web/mod.rs

@@ -14,7 +14,8 @@ use crate::{
     app::models::{Entity, Project, UserRole},
     utils::validate::can_access_project,
     web::{
-        MaybeExtract, RouterState, SessionUser,
+        RouterState,
+        session::{Auth, MaybeAuth},
         webext::{ApiResult, AxumErrExt, empty_response, http_bail},
     },
 };
@@ -152,7 +153,7 @@ struct EntitiesResponse {
 
 async fn get_users(
     app: State<RouterState>,
-    SessionUser(user): SessionUser,
+    Auth(user): Auth,
 ) -> ApiResult<UseApi<impl IntoApiResponse, Json<UsersResponse>>> {
     if user.role != UserRole::Admin {
         http_bail!(StatusCode::FORBIDDEN, "Forbidden")
@@ -172,7 +173,7 @@ async fn get_users(
 async fn update_user(
     app: State<RouterState>,
     Path(username): Path<String>,
-    SessionUser(session_user): SessionUser,
+    Auth(session_user): Auth,
     user: Json<UpdateUserRequest>,
 ) -> ApiResult<impl IntoApiResponse> {
     if session_user.role != UserRole::Admin {
@@ -193,7 +194,7 @@ async fn update_user(
 async fn update_user_password(
     app: State<RouterState>,
     Path(username): Path<String>,
-    SessionUser(session_user): SessionUser,
+    Auth(session_user): Auth,
     password: Json<UpdatePasswordRequest>,
 ) -> ApiResult<impl IntoApiResponse> {
     if session_user.role != UserRole::Admin || username != session_user.username {
@@ -210,7 +211,7 @@ async fn update_user_password(
 async fn remove_user(
     app: State<RouterState>,
     Path(username): Path<String>,
-    SessionUser(session_user): SessionUser,
+    Auth(session_user): Auth,
 ) -> ApiResult<impl IntoApiResponse> {
     if session_user.role != UserRole::Admin {
         http_bail!(StatusCode::FORBIDDEN, "Forbidden")
@@ -227,7 +228,7 @@ async fn remove_user(
 
 async fn create_user(
     app: State<RouterState>,
-    SessionUser(session_user): SessionUser,
+    Auth(session_user): Auth,
     user: Json<CreateUserRequest>,
 ) -> ApiResult<impl IntoApiResponse> {
     if session_user.role != UserRole::Admin {
@@ -246,7 +247,7 @@ async fn create_user(
 async fn project_create_handler(
     app: State<RouterState>,
     Path(project_id): Path<String>,
-    SessionUser(user): SessionUser,
+    Auth(user): Auth,
     Json(project): Json<CreateProjectRequest>,
 ) -> ApiResult<impl IntoApiResponse> {
     if user.role != UserRole::Admin {
@@ -271,7 +272,7 @@ async fn project_create_handler(
 async fn project_update_handler(
     app: State<RouterState>,
     Path(project_id): Path<String>,
-    SessionUser(user): SessionUser,
+    Auth(user): Auth,
     Json(req): Json<UpdateProjectRequest>,
 ) -> ApiResult<impl IntoApiResponse> {
     if user.role != UserRole::Admin {
@@ -300,7 +301,7 @@ async fn project_update_handler(
 
 async fn projects_handler(
     app: State<RouterState>,
-    MaybeExtract(user): MaybeExtract<SessionUser>,
+    MaybeAuth(user): MaybeAuth,
 ) -> ApiResult<UseApi<impl IntoApiResponse, Json<ProjectsResponse>>> {
     let projects = app.projects.all().http_err("Failed to get projects", StatusCode::INTERNAL_SERVER_ERROR)?;
     let projects: Vec<Project> = projects.into_iter().filter(|p| can_access_project(p, user.as_ref())).collect();
@@ -326,7 +327,7 @@ async fn projects_handler(
 
 async fn project_handler(
     app: State<RouterState>,
-    MaybeExtract(user): MaybeExtract<SessionUser>,
+    MaybeAuth(user): MaybeAuth,
     Path(project_id): Path<String>,
 ) -> ApiResult<UseApi<impl IntoApiResponse, Json<ProjectResponse>>> {
     let project = app.projects.get(&project_id).http_status(StatusCode::NOT_FOUND)?;
@@ -353,7 +354,7 @@ async fn project_handler(
 async fn project_delete_handler(
     app: State<RouterState>,
     Path(project_id): Path<String>,
-    SessionUser(user): SessionUser,
+    Auth(user): Auth,
 ) -> ApiResult<impl IntoApiResponse> {
     let project = app.projects.get(&project_id).http_status(StatusCode::NOT_FOUND)?;
     if user.role != UserRole::Admin {
@@ -366,7 +367,7 @@ async fn project_delete_handler(
 
 async fn entities_handler(
     app: State<RouterState>,
-    SessionUser(user): SessionUser,
+    Auth(user): Auth,
 ) -> ApiResult<UseApi<impl IntoApiResponse, Json<EntitiesResponse>>> {
     if user.role != UserRole::Admin {
         http_bail!(StatusCode::FORBIDDEN, "Forbidden")
@@ -398,7 +399,7 @@ async fn entities_handler(
 
 async fn entity_create_handler(
     app: State<RouterState>,
-    SessionUser(user): SessionUser,
+    Auth(user): Auth,
     Json(entity): Json<CreateEntityRequest>,
 ) -> ApiResult<Json<EntityResponse>> {
     if user.role != UserRole::Admin {
@@ -418,7 +419,7 @@ async fn entity_create_handler(
 async fn entity_update_handler(
     app: State<RouterState>,
     Path(entity_id): Path<String>,
-    SessionUser(user): SessionUser,
+    Auth(user): Auth,
     Json(entity): Json<UpdateEntityRequest>,
 ) -> ApiResult<impl IntoApiResponse> {
     if user.role != UserRole::Admin {
@@ -443,7 +444,7 @@ async fn entity_update_handler(
 async fn entity_delete_handler(
     app: State<RouterState>,
     Path(entity_id): Path<String>,
-    SessionUser(user): SessionUser,
+    Auth(user): Auth,
 ) -> ApiResult<impl IntoApiResponse> {
     if user.role != UserRole::Admin {
         http_bail!(StatusCode::FORBIDDEN, "Forbidden")