Skip to content

Commit 0586ac2

Browse files
azasypkinazasypkin
committed
chore(api): improve DB connection reliability and include DB status in API status
1 parent 422a207 commit 0586ac2

9 files changed

Lines changed: 126 additions & 36 deletions

File tree

‎components/retrack‎

‎src/database.rs‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,11 @@ impl Database {
1717

1818
Ok(Database { pool })
1919
}
20+
21+
/// Returns `true` if the database is reachable.
22+
pub async fn is_alive(&self) -> bool {
23+
sqlx::query("SELECT 1").execute(&self.pool).await.is_ok()
24+
}
2025
}
2126

2227
impl AsRef<Database> for Database {

‎src/server.rs‎

Lines changed: 38 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,8 @@ mod ui_state;
77
#[cfg(test)]
88
pub use self::app_state::tests;
99
pub use self::ui_state::{
10-
Status, StatusLevel, SubscriptionState, UiPlatformState, UiState, WebhookUrlType,
10+
DatabaseStatus, Status, StatusLevel, SubscriptionState, UiPlatformState, UiState,
11+
WebhookUrlType,
1112
};
1213

1314
use crate::{
@@ -27,8 +28,8 @@ pub use app_state::AppState;
2728
use handlers::SecutilsOpenApi;
2829
use serde_json::json;
2930
use sqlx::postgres::PgPoolOptions;
30-
use std::sync::Arc;
31-
use tracing::info;
31+
use std::{sync::Arc, time::Duration};
32+
use tracing::{info, warn};
3233
use tracing_actix_web::TracingLogger;
3334
use utoipa::OpenApi;
3435
use utoipa_rapidoc::RapiDoc;
@@ -57,22 +58,43 @@ pub async fn run(config: Config, http_port: u16) -> Result<(), anyhow::Error> {
5758
config.db.port,
5859
urlencoding::encode(&config.db.name)
5960
);
60-
let database = Database::create(
61-
PgPoolOptions::new()
62-
.max_connections(config.db.max_connections)
63-
.min_connections(config.db.min_connections)
64-
.acquire_timeout(config.db.acquire_timeout)
65-
.max_lifetime(config.db.max_lifetime)
66-
.idle_timeout(config.db.idle_timeout)
67-
.test_before_acquire(true)
68-
.connect(&db_url)
69-
.await?,
70-
)
71-
.await?;
61+
let pool_options = PgPoolOptions::new()
62+
.max_connections(config.db.max_connections)
63+
.min_connections(config.db.min_connections)
64+
.acquire_timeout(config.db.acquire_timeout)
65+
.max_lifetime(config.db.max_lifetime)
66+
.idle_timeout(config.db.idle_timeout)
67+
.test_before_acquire(true);
68+
69+
let pool = {
70+
const MAX_RETRIES: u32 = 5;
71+
let mut attempt = 0;
72+
loop {
73+
match pool_options.clone().connect(&db_url).await {
74+
Ok(pool) => break pool,
75+
Err(err) => {
76+
attempt += 1;
77+
if attempt >= MAX_RETRIES {
78+
return Err(err).context(format!(
79+
"Failed to connect to database after {MAX_RETRIES} attempts"
80+
));
81+
}
82+
83+
let delay = Duration::from_secs(1 << attempt.min(4));
84+
warn!(
85+
attempt,
86+
max_retries = MAX_RETRIES,
87+
"Failed to connect to database: {err}. Retrying in {delay:?}…"
88+
);
89+
tokio::time::sleep(delay).await;
90+
}
91+
}
92+
}
93+
};
7294

7395
let api = Arc::new(Api::new(
7496
config.clone(),
75-
database,
97+
Database::create(pool).await?,
7698
search_index,
7799
Network::create(&config)?,
78100
create_templates()?,

‎src/server/app_state.rs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ use crate::{
22
api::Api,
33
config::Config,
44
network::{DnsResolver, EmailTransport, TokioDnsResolver},
5-
server::{Status, StatusLevel},
5+
server::{DatabaseStatus, Status, StatusLevel},
66
};
77
use dashmap::DashMap;
88
use lettre::{AsyncSmtpTransport, Tokio1Executor};
@@ -28,6 +28,7 @@ impl<DR: DnsResolver, ET: EmailTransport> AppState<DR, ET> {
2828
status: RwLock::new(Status {
2929
version: env!("CARGO_PKG_VERSION").to_string(),
3030
level: StatusLevel::Available,
31+
db: DatabaseStatus { operational: true },
3132
}),
3233
api,
3334
responder_semaphores: DashMap::new(),

‎src/server/handlers.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -245,6 +245,7 @@ pub(crate) async fn resolve_shared_user(
245245
crate::users::UserSettings,
246246
crate::users::UserSettingsSetter,
247247
// Status
248+
crate::server::DatabaseStatus,
248249
crate::server::Status,
249250
crate::server::StatusLevel,
250251
// Subscription
@@ -466,6 +467,7 @@ mod tests {
466467
"ContentSecurityPolicyTrustedTypesDirectiveValue",
467468
"ContentSecurityPolicyWebrtcDirectiveValue",
468469
"DataFileSecret",
470+
"DatabaseStatus",
469471
"EmailParams",
470472
"EntityTag",
471473
"ExportFormat",

‎src/server/handlers/status_get.rs‎

Lines changed: 19 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,9 @@
1-
use crate::{error::Error as SecutilsError, server::app_state::AppState};
1+
use crate::{
2+
error::Error as SecutilsError,
3+
server::{DatabaseStatus, StatusLevel, app_state::AppState},
4+
};
25
use actix_web::{HttpResponse, get, web};
36
use anyhow::anyhow;
4-
use std::ops::Deref;
57
use tracing::error;
68

79
/// Returns the current server status.
@@ -14,12 +16,19 @@ use tracing::error;
1416
)]
1517
#[get("/api/status")]
1618
pub async fn status_get(state: web::Data<AppState>) -> Result<HttpResponse, SecutilsError> {
17-
state
18-
.status
19-
.read()
20-
.map(|status| HttpResponse::Ok().json(status.deref()))
21-
.map_err(|err| {
22-
error!("Failed to read status: {err}");
23-
SecutilsError::from(anyhow!("Status is not available."))
24-
})
19+
let mut status = state.status.read().map(|s| s.clone()).map_err(|err| {
20+
error!("Failed to read status: {err}");
21+
SecutilsError::from(anyhow!("Status is not available."))
22+
})?;
23+
24+
let db_operational = state.api.db.is_alive().await;
25+
status.db = DatabaseStatus {
26+
operational: db_operational,
27+
};
28+
if !db_operational {
29+
error!("Database is not reachable.");
30+
status.level = StatusLevel::Unavailable;
31+
}
32+
33+
Ok(HttpResponse::Ok().json(status))
2534
}

‎src/server/ui_state.rs‎

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,13 @@
1+
mod database_status;
12
mod status;
23
mod status_level;
34
mod webhook_url_type;
45

56
mod subscription_state;
67

78
pub use self::{
8-
status::Status, status_level::StatusLevel, subscription_state::SubscriptionState,
9-
webhook_url_type::WebhookUrlType,
9+
database_status::DatabaseStatus, status::Status, status_level::StatusLevel,
10+
subscription_state::SubscriptionState, webhook_url_type::WebhookUrlType,
1011
};
1112
use crate::{
1213
users::{ClientUserShare, User, UserSettings},
@@ -51,7 +52,7 @@ mod tests {
5152

5253
use crate::{
5354
server::{
54-
Status, StatusLevel, UiPlatformState, UiState, WebhookUrlType,
55+
DatabaseStatus, Status, StatusLevel, UiPlatformState, UiState, WebhookUrlType,
5556
ui_state::subscription_state::SubscriptionState,
5657
},
5758
tests::{mock_config, mock_user},
@@ -70,6 +71,7 @@ mod tests {
7071
status: &Status {
7172
version: "1.0.0-alpha.4".to_string(),
7273
level: StatusLevel::Available,
74+
db: DatabaseStatus { operational: true },
7375
},
7476
user: Some(user),
7577
subscription: Some(SubscriptionState {
@@ -106,7 +108,10 @@ mod tests {
106108
{
107109
"status": {
108110
"version": "1.0.0-alpha.4",
109-
"level": "available"
111+
"level": "available",
112+
"db": {
113+
"operational": true
114+
}
110115
},
111116
"user": {
112117
"email": "dev-00000000-0000-0000-0000-000000000001@secutils.dev",
@@ -168,6 +173,7 @@ mod tests {
168173
status: &Status {
169174
version: "1.0.0-alpha.4".to_string(),
170175
level: StatusLevel::Available,
176+
db: DatabaseStatus { operational: true },
171177
},
172178
user: None,
173179
subscription: Default::default(),
@@ -183,7 +189,10 @@ mod tests {
183189
{
184190
"status": {
185191
"version": "1.0.0-alpha.4",
186-
"level": "available"
192+
"level": "available",
193+
"db": {
194+
"operational": true
195+
}
187196
},
188197
"utils": [],
189198
"webhookUrlType": "subdomain",

‎src/server/ui_state/database_status.rs‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
use serde::Serialize;
2+
use utoipa::ToSchema;
3+
4+
/// Database connectivity status.
5+
#[derive(Clone, Serialize, ToSchema)]
6+
pub struct DatabaseStatus {
7+
/// Indicates if the database is reachable.
8+
pub operational: bool,
9+
}
10+
11+
#[cfg(test)]
12+
mod tests {
13+
use crate::server::DatabaseStatus;
14+
use insta::assert_json_snapshot;
15+
16+
#[test]
17+
fn serialization() -> anyhow::Result<()> {
18+
assert_json_snapshot!(DatabaseStatus {
19+
operational: true,
20+
}, @r###"
21+
{
22+
"operational": true
23+
}
24+
"###);
25+
26+
assert_json_snapshot!(DatabaseStatus {
27+
operational: false,
28+
}, @r###"
29+
{
30+
"operational": false
31+
}
32+
"###);
33+
34+
Ok(())
35+
}
36+
}

‎src/server/ui_state/status.rs‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
use crate::server::StatusLevel;
1+
use crate::server::{DatabaseStatus, StatusLevel};
22
use serde::Serialize;
33
use utoipa::ToSchema;
44

@@ -9,22 +9,28 @@ pub struct Status {
99
pub version: String,
1010
/// Current availability level.
1111
pub level: StatusLevel,
12+
/// Status of the database connection.
13+
pub db: DatabaseStatus,
1214
}
1315

1416
#[cfg(test)]
1517
mod tests {
16-
use crate::server::{Status, StatusLevel};
18+
use crate::server::{DatabaseStatus, Status, StatusLevel};
1719
use insta::assert_json_snapshot;
1820

1921
#[test]
2022
fn serialization() -> anyhow::Result<()> {
2123
assert_json_snapshot!(Status {
2224
version: "1.0.0-alpha.4".to_string(),
2325
level: StatusLevel::Available,
26+
db: DatabaseStatus { operational: true },
2427
}, @r###"
2528
{
2629
"version": "1.0.0-alpha.4",
27-
"level": "available"
30+
"level": "available",
31+
"db": {
32+
"operational": true
33+
}
2834
}
2935
"###);
3036

0 commit comments

Comments
 (0)