vw_small

Hardened fork of Vaultwarden (https://github.com/dani-garcia/vaultwarden) with fewer features.
git clone https://git.philomathiclife.com/repos/vw_small
Log | Files | Refs | README

commit 729c9cff41cc74055f8397fae7f60084dcf4b71b
parent 22b9c800078f3d3cf1893857cd875306bfc632e7
Author: Daniel GarcĂ­a <dani-garcia@users.noreply.github.com>
Date:   Sat,  3 Oct 2020 22:31:52 +0200

Retry initial db connection, with adjustable option

Diffstat:
Msrc/config.rs | 3+++
Msrc/db/mod.rs | 23+++++++++++++----------
Msrc/error.rs | 5+++++
Msrc/main.rs | 2+-
Msrc/util.rs | 28+++++++++++++++++++++++++++-
5 files changed, 49 insertions(+), 12 deletions(-)

diff --git a/src/config.rs b/src/config.rs @@ -347,6 +347,9 @@ make_config! { /// that do not support WAL. Please make sure you read project wiki on the topic before changing this setting. enable_db_wal: bool, false, def, true; + /// Max database connection retries |> Number of times to retry the database connection during startup, with 1 second between each retry, set to 0 to retry indefinitely + db_connection_retries: u32, false, def, 15; + /// Bypass admin page security (Know the risks!) |> Disables the Admin Token for the admin page so you may use your own auth in-front disable_admin_token: bool, true, def, false; diff --git a/src/db/mod.rs b/src/db/mod.rs @@ -49,7 +49,7 @@ macro_rules! generate_connections { DbConnType::$name => { #[cfg($name)] { - paste::paste!{ [< $name _migrations >]::run_migrations(); } + paste::paste!{ [< $name _migrations >]::run_migrations()?; } let manager = ConnectionManager::new(&url); let pool = Pool::builder().build(manager).map_res("Failed to create pool")?; return Ok(Self::$name(pool)); @@ -242,7 +242,7 @@ mod sqlite_migrations { #[allow(unused_imports)] embed_migrations!("migrations/sqlite"); - pub fn run_migrations() { + pub fn run_migrations() -> Result<(), super::Error> { // Make sure the directory exists let url = crate::CONFIG.database_url(); let path = std::path::Path::new(&url); @@ -257,7 +257,7 @@ mod sqlite_migrations { use diesel::{Connection, RunQueryDsl}; // Make sure the database is up to date (create if it doesn't exist, or run the migrations) let connection = - diesel::sqlite::SqliteConnection::establish(&crate::CONFIG.database_url()).expect("Can't connect to DB"); + diesel::sqlite::SqliteConnection::establish(&crate::CONFIG.database_url())?; // Disable Foreign Key Checks during migration // Scoped to a connection. @@ -272,7 +272,8 @@ mod sqlite_migrations { .expect("Failed to turn on WAL"); } - embedded_migrations::run_with_output(&connection, &mut std::io::stdout()).expect("Can't run migrations"); + embedded_migrations::run_with_output(&connection, &mut std::io::stdout())?; + Ok(()) } } @@ -281,11 +282,11 @@ mod mysql_migrations { #[allow(unused_imports)] embed_migrations!("migrations/mysql"); - pub fn run_migrations() { + pub fn run_migrations() -> Result<(), super::Error> { use diesel::{Connection, RunQueryDsl}; // Make sure the database is up to date (create if it doesn't exist, or run the migrations) let connection = - diesel::mysql::MysqlConnection::establish(&crate::CONFIG.database_url()).expect("Can't connect to DB"); + diesel::mysql::MysqlConnection::establish(&crate::CONFIG.database_url())?; // Disable Foreign Key Checks during migration // Scoped to a connection/session. @@ -293,7 +294,8 @@ mod mysql_migrations { .execute(&connection) .expect("Failed to disable Foreign Key Checks during migrations"); - embedded_migrations::run_with_output(&connection, &mut std::io::stdout()).expect("Can't run migrations"); + embedded_migrations::run_with_output(&connection, &mut std::io::stdout())?; + Ok(()) } } @@ -302,11 +304,11 @@ mod postgresql_migrations { #[allow(unused_imports)] embed_migrations!("migrations/postgresql"); - pub fn run_migrations() { + pub fn run_migrations() -> Result<(), super::Error> { use diesel::{Connection, RunQueryDsl}; // Make sure the database is up to date (create if it doesn't exist, or run the migrations) let connection = - diesel::pg::PgConnection::establish(&crate::CONFIG.database_url()).expect("Can't connect to DB"); + diesel::pg::PgConnection::establish(&crate::CONFIG.database_url())?; // Disable Foreign Key Checks during migration // FIXME: Per https://www.postgresql.org/docs/12/sql-set-constraints.html, @@ -319,6 +321,7 @@ mod postgresql_migrations { .execute(&connection) .expect("Failed to disable Foreign Key Checks during migrations"); - embedded_migrations::run_with_output(&connection, &mut std::io::stdout()).expect("Can't run migrations"); + embedded_migrations::run_with_output(&connection, &mut std::io::stdout())?; + Ok(()) } } diff --git a/src/error.rs b/src/error.rs @@ -34,6 +34,8 @@ macro_rules! make_error { } use diesel::result::Error as DieselErr; +use diesel::ConnectionError as DieselConErr; +use diesel_migrations::RunMigrationsError as DieselMigErr; use diesel::r2d2::PoolError as R2d2Err; use handlebars::RenderError as HbErr; use jsonwebtoken::errors::Error as JWTErr; @@ -83,6 +85,9 @@ make_error! { AddressError(AddrErr): _has_source, _api_error, SmtpError(SmtpErr): _has_source, _api_error, FromStrError(FromStrErr): _has_source, _api_error, + + DieselConError(DieselConErr): _has_source, _api_error, + DieselMigError(DieselMigErr): _has_source, _api_error, } impl std::fmt::Debug for Error { diff --git a/src/main.rs b/src/main.rs @@ -268,7 +268,7 @@ fn check_web_vault() { } fn launch_rocket(extra_debug: bool) { - let pool = match db::DbPool::from_config() { + let pool = match util::retry_db(db::DbPool::from_config, CONFIG.db_connection_retries()) { Ok(p) => p, Err(e) => { error!("Error creating database pool: {:?}", e); diff --git a/src/util.rs b/src/util.rs @@ -410,7 +410,7 @@ fn _process_key(key: &str) -> String { // Retry methods // -pub fn retry<F, T, E>(func: F, max_tries: i32) -> Result<T, E> +pub fn retry<F, T, E>(func: F, max_tries: u32) -> Result<T, E> where F: Fn() -> Result<T, E>, { @@ -432,3 +432,29 @@ where } } } + +pub fn retry_db<F, T, E>(func: F, max_tries: u32) -> Result<T, E> +where + F: Fn() -> Result<T, E>, + E: std::error::Error, +{ + use std::{thread::sleep, time::Duration}; + let mut tries = 0; + + loop { + match func() { + ok @ Ok(_) => return ok, + Err(e) => { + tries += 1; + + if tries >= max_tries && max_tries > 0 { + return Err(e); + } + + warn!("Can't connect to database, retrying: {:?}", e); + + sleep(Duration::from_millis(1_000)); + } + } + } +}