From 593666135c9de8b4e1ad4ee6d68f2f15c1e7fc68 Mon Sep 17 00:00:00 2001 From: xfy Date: Thu, 4 Jun 2026 10:03:56 +0800 Subject: [PATCH] feat: add DB connection retry logic with get_conn() helper --- src/api/auth.rs | 10 +++++----- src/api/posts.rs | 24 ++++++++++++------------ src/db/mod.rs | 4 ++++ src/db/pool.rs | 21 +++++++++++++++++++++ src/tasks/session_cleanup.rs | 4 ++-- 5 files changed, 44 insertions(+), 19 deletions(-) diff --git a/src/api/auth.rs b/src/api/auth.rs index ea3d80f..0dec85b 100644 --- a/src/api/auth.rs +++ b/src/api/auth.rs @@ -5,7 +5,7 @@ use dioxus::prelude::*; use http::header::{HeaderValue, SET_COOKIE}; use crate::auth::{password, session}; -use crate::db::pool::DB_POOL; +use crate::db::pool::get_conn; use crate::models::user::{PublicUser, User, UserRole}; #[allow(dead_code)] @@ -85,7 +85,7 @@ pub async fn register( }); } - let client = DB_POOL.get().await.map_err(|e| { + let client = get_conn().await.map_err(|e| { tracing::error!("Register DB connection failed: {:?}", e); ServerFnError::new(format!("数据库连接失败: {}", e)) })?; @@ -142,7 +142,7 @@ pub async fn register( #[server(Login, "/api")] pub async fn login(username: String, password: String) -> Result { - let client = DB_POOL.get().await.map_err(|e| { + let client = get_conn().await.map_err(|e| { tracing::error!("Login DB connection failed: {:?}", e); ServerFnError::new(format!("数据库连接失败: {}", e)) })?; @@ -228,7 +228,7 @@ pub async fn logout() -> Result { None }; - let client = DB_POOL.get().await.map_err(|e| { + let client = get_conn().await.map_err(|e| { tracing::error!("Logout DB connection failed: {:?}", e); ServerFnError::new(format!("数据库连接失败: {}", e)) })?; @@ -282,7 +282,7 @@ pub async fn get_current_user() -> Result { return Ok(CurrentUserResponse { user: None }); }; - let client = DB_POOL.get().await.map_err(|e| { + let client = get_conn().await.map_err(|e| { tracing::error!("GetCurrentUser DB connection failed: {:?}", e); ServerFnError::new(format!("数据库连接失败: {}", e)) })?; diff --git a/src/api/posts.rs b/src/api/posts.rs index 7deedc3..a64cad5 100644 --- a/src/api/posts.rs +++ b/src/api/posts.rs @@ -2,7 +2,7 @@ use dioxus::prelude::*; -use crate::db::pool::DB_POOL; +use crate::db::pool::get_conn; use crate::models::post::{Post, PostStats, PostStatus, Tag}; use crate::models::user::{User, UserRole}; @@ -42,7 +42,7 @@ async fn get_current_admin_user() -> Result { return Err(ServerFnError::new("未登录")); }; - let client = DB_POOL.get().await.map_err(|e| { + let client = get_conn().await.map_err(|e| { tracing::error!("DB connection failed: {:?}", e); ServerFnError::new(format!("数据库连接失败: {}", e)) })?; @@ -782,7 +782,7 @@ pub async fn create_post( _ => slugify(&title), }; - let mut client = DB_POOL.get().await.map_err(|e| { + let mut client = get_conn().await.map_err(|e| { tracing::error!("DB connection failed: {:?}", e); ServerFnError::new(format!("数据库连接失败: {}", e)) })?; @@ -910,7 +910,7 @@ pub async fn update_post( ) -> Result { let user = get_current_admin_user().await?; - let mut client = DB_POOL.get().await.map_err(|e| { + let mut client = get_conn().await.map_err(|e| { tracing::error!("DB connection failed: {:?}", e); ServerFnError::new(format!("数据库连接失败: {}", e)) })?; @@ -1086,7 +1086,7 @@ pub async fn update_post( #[server(GetPostBySlug, "/api")] pub async fn get_post_by_slug(slug: String) -> Result { - let client = DB_POOL.get().await.map_err(|e| { + let client = get_conn().await.map_err(|e| { tracing::error!("DB connection failed: {:?}", e); ServerFnError::new(format!("数据库连接失败: {}", e)) })?; @@ -1137,7 +1137,7 @@ pub async fn list_published_posts( page: i32, per_page: i32, ) -> Result { - let client = DB_POOL.get().await.map_err(|e| { + let client = get_conn().await.map_err(|e| { tracing::error!("DB connection failed: {:?}", e); ServerFnError::new(format!("数据库连接失败: {}", e)) })?; @@ -1171,7 +1171,7 @@ pub async fn list_published_posts( pub async fn list_posts() -> Result { let _user = get_current_admin_user().await?; - let client = DB_POOL.get().await.map_err(|e| { + let client = get_conn().await.map_err(|e| { tracing::error!("DB connection failed: {:?}", e); ServerFnError::new(format!("数据库连接失败: {}", e)) })?; @@ -1202,7 +1202,7 @@ pub async fn list_posts() -> Result { pub async fn delete_post(post_id: i32) -> Result { let _user = get_current_admin_user().await?; - let client = DB_POOL.get().await.map_err(|e| { + let client = get_conn().await.map_err(|e| { tracing::error!("DB connection failed: {:?}", e); ServerFnError::new(format!("数据库连接失败: {}", e)) })?; @@ -1237,7 +1237,7 @@ pub async fn delete_post(post_id: i32) -> Result Result { - let client = DB_POOL.get().await.map_err(|e| { + let client = get_conn().await.map_err(|e| { tracing::error!("DB connection failed: {:?}", e); ServerFnError::new(format!("数据库连接失败: {}", e)) })?; @@ -1272,7 +1272,7 @@ pub async fn list_tags() -> Result { #[server(GetPostsByTag, "/api")] pub async fn get_posts_by_tag(tag_name: String) -> Result { - let client = DB_POOL.get().await.map_err(|e| { + let client = get_conn().await.map_err(|e| { tracing::error!("DB connection failed: {:?}", e); ServerFnError::new(format!("数据库连接失败: {}", e)) })?; @@ -1305,7 +1305,7 @@ pub async fn get_posts_by_tag(tag_name: String) -> Result Result { let _user = get_current_admin_user().await?; - let client = DB_POOL.get().await.map_err(|e| { + let client = get_conn().await.map_err(|e| { tracing::error!("DB connection failed: {:?}", e); ServerFnError::new(format!("数据库连接失败: {}", e)) })?; @@ -1345,7 +1345,7 @@ pub async fn get_post_stats() -> Result { #[server(SearchPosts, "/api")] pub async fn search_posts(query: String) -> Result { - let client = DB_POOL.get().await.map_err(|e| { + let client = get_conn().await.map_err(|e| { tracing::error!("DB connection failed: {:?}", e); ServerFnError::new(format!("数据库连接失败: {}", e)) })?; diff --git a/src/db/mod.rs b/src/db/mod.rs index 612046a..08c09bd 100644 --- a/src/db/mod.rs +++ b/src/db/mod.rs @@ -11,4 +11,8 @@ pub mod pool { } } pub static DB_POOL: DummyPool = DummyPool; + + pub async fn get_conn() -> Result<(), ()> { + Err(()) + } } diff --git a/src/db/pool.rs b/src/db/pool.rs index 90b975e..804ed31 100644 --- a/src/db/pool.rs +++ b/src/db/pool.rs @@ -1,4 +1,5 @@ use std::sync::LazyLock; +use std::time::Duration; use deadpool_postgres::{Manager, ManagerConfig, Pool, RecyclingMethod}; use tokio_postgres::NoTls; @@ -19,3 +20,23 @@ pub static DB_POOL: LazyLock = LazyLock::new(|| { .build() .expect("Failed to create database connection pool") }); + +const MAX_RETRIES: u32 = 3; +const RETRY_DELAY: Duration = Duration::from_secs(2); + +pub async fn get_conn() -> Result { + let mut last_err = None; + for attempt in 0..=MAX_RETRIES { + match DB_POOL.get().await { + Ok(conn) => return Ok(conn), + Err(e) => { + if attempt < MAX_RETRIES { + tracing::warn!("DB connection attempt {} failed, retrying in {:?}: {:?}", attempt + 1, RETRY_DELAY, e); + tokio::time::sleep(RETRY_DELAY).await; + } + last_err = Some(e); + } + } + } + Err(last_err.unwrap()) +} diff --git a/src/tasks/session_cleanup.rs b/src/tasks/session_cleanup.rs index b5df719..df32c5c 100644 --- a/src/tasks/session_cleanup.rs +++ b/src/tasks/session_cleanup.rs @@ -2,13 +2,13 @@ use std::time::Duration; use tokio::time::interval; -use crate::db::pool::DB_POOL; +use crate::db::pool::get_conn; pub async fn run_cleanup() { let mut ticker = interval(Duration::from_secs(3600)); loop { ticker.tick().await; - match DB_POOL.get().await { + match get_conn().await { Ok(client) => { if let Err(e) = client .execute("DELETE FROM sessions WHERE expires_at < NOW()", &[])