From 5d5e65ed4042115a42b80ca920a49abfe44968a8 Mon Sep 17 00:00:00 2001 From: Bryan Stitt Date: Thu, 13 Jul 2023 14:56:17 -0700 Subject: [PATCH] close the db connection on exit --- web3_proxy/src/app/mod.rs | 25 +++++++++++++++++++++++-- web3_proxy/src/sub_commands/proxyd.rs | 13 ++++++++++++- web3_proxy/tests/common/create_admin.rs | 4 +++- web3_proxy/tests/common/create_user.rs | 9 +++------ web3_proxy/tests/common/mysql.rs | 13 +++++++------ 5 files changed, 48 insertions(+), 16 deletions(-) diff --git a/web3_proxy/src/app/mod.rs b/web3_proxy/src/app/mod.rs index 256429e1..f6f5bfc9 100644 --- a/web3_proxy/src/app/mod.rs +++ b/web3_proxy/src/app/mod.rs @@ -49,6 +49,7 @@ use std::str::FromStr; use std::sync::atomic::{AtomicU16, Ordering}; use std::sync::{atomic, Arc}; use std::time::Duration; +use tokio::runtime::Runtime; use tokio::sync::{broadcast, mpsc, oneshot, watch, Semaphore}; use tokio::task::JoinHandle; use tokio::time::timeout; @@ -160,11 +161,31 @@ pub struct Web3ProxyAppSpawn { /// these are important and must be allowed to finish pub background_handles: FuturesUnordered>, /// config changes are sent here - pub new_top_config: watch::Sender, + pub new_top_config: Arc>, /// watch this to know when the app is ready to serve requests pub ranked_rpcs: watch::Receiver>>, } +impl Drop for Web3ProxyApp { + fn drop(&mut self) { + if let Ok(db_conn) = self.db_conn().cloned() { + /* + From the sqlx docs: + + We recommend calling .close().await to gracefully close the pool and its connections when you are done using it. + This will also wake any tasks that are waiting on an .acquire() call, + so for long-lived applications it’s a good idea to call .close() during shutdown. + */ + + let rt = Runtime::new().unwrap(); + + if let Err(err) = rt.block_on(db_conn.close()) { + error!(?err, "Unable to close db!"); + }; + } + } +} + impl Web3ProxyApp { /// The main entrypoint. pub async fn spawn( @@ -610,7 +631,7 @@ impl Web3ProxyApp { app, app_handles, background_handles: important_background_handles, - new_top_config: new_top_config_sender, + new_top_config: Arc::new(new_top_config_sender), ranked_rpcs: consensus_connections_watcher, }) } diff --git a/web3_proxy/src/sub_commands/proxyd.rs b/web3_proxy/src/sub_commands/proxyd.rs index 80ad8396..aabb6088 100644 --- a/web3_proxy/src/sub_commands/proxyd.rs +++ b/web3_proxy/src/sub_commands/proxyd.rs @@ -174,7 +174,7 @@ impl ProxydSubCommand { // start the frontend port let frontend_handle = tokio::spawn(frontend::serve( - spawned_app.app, + spawned_app.app.clone(), frontend_shutdown_receiver, frontend_shutdown_complete_sender, )); @@ -286,6 +286,17 @@ impl ProxydSubCommand { } } + if let Ok(db_conn) = spawned_app.app.db_conn().cloned() { + /* + From the sqlx docs: + + We recommend calling .close().await to gracefully close the pool and its connections when you are done using it. + This will also wake any tasks that are waiting on an .acquire() call, + so for long-lived applications it’s a good idea to call .close() during shutdown. + */ + db_conn.close().await?; + } + if background_errors.is_zero() && !exited_with_err { info!("finished"); Ok(()) diff --git a/web3_proxy/tests/common/create_admin.rs b/web3_proxy/tests/common/create_admin.rs index a181dcae..48bfebd0 100644 --- a/web3_proxy/tests/common/create_admin.rs +++ b/web3_proxy/tests/common/create_admin.rs @@ -66,9 +66,11 @@ pub async fn create_user_as_admin( }; info!(?admin_status_changer); + let db_conn = db.conn().await; + info!("Changing the status of the admin_wallet to be an admin"); // Pass on the database into it ... - admin_status_changer.main(db.conn()).await.unwrap(); + admin_status_changer.main(&db_conn).await.unwrap(); // Now log him in again, because he was just signed out // Login the admin again, because he was just signed out diff --git a/web3_proxy/tests/common/create_user.rs b/web3_proxy/tests/common/create_user.rs index 8c75cc07..a5695a55 100644 --- a/web3_proxy/tests/common/create_user.rs +++ b/web3_proxy/tests/common/create_user.rs @@ -3,14 +3,13 @@ use entities::{user, user_tier}; use ethers::prelude::{LocalWallet, Signer}; use ethers::types::Signature; use migration::sea_orm::{ - self, ActiveModelTrait, ColumnTrait, EntityTrait, IntoActiveModel, QueryFilter, + self, ActiveModelTrait, ColumnTrait, DatabaseConnection, EntityTrait, IntoActiveModel, + QueryFilter, }; use tracing::info; use web3_proxy::errors::Web3ProxyResult; use web3_proxy::frontend::users::authentication::{LoginPostResponse, PostLogin}; -use super::mysql::TestMysql; - /// Helper function to create an "ordinary" user #[allow(unused)] pub async fn create_user( @@ -57,12 +56,10 @@ pub async fn create_user( #[allow(unused)] pub async fn set_user_tier( x: &TestApp, - db: &TestMysql, + db_conn: &DatabaseConnection, user: user::Model, tier_name: &str, ) -> Web3ProxyResult { - let db_conn = db.conn(); - let ut = user_tier::Entity::find() .filter(user_tier::Column::Title.like(tier_name)) .one(db_conn) diff --git a/web3_proxy/tests/common/mysql.rs b/web3_proxy/tests/common/mysql.rs index fddb277f..e6fd0140 100644 --- a/web3_proxy/tests/common/mysql.rs +++ b/web3_proxy/tests/common/mysql.rs @@ -10,10 +10,11 @@ use tokio::{ use tracing::{info, trace, warn}; use web3_proxy::relational_db::get_migrated_db; +use migration::sea_orm::prelude; + /// on drop, the mysql docker container will be shut down pub struct TestMysql { pub url: Option, - pub conn: Option, pub container_name: String, } @@ -40,7 +41,6 @@ impl TestMysql { // create the db_data as soon as the url is known // when this is dropped, the db will be stopped let mut test_mysql = Self { - conn: None, container_name: db_container_name.clone(), url: None, }; @@ -150,9 +150,8 @@ impl TestMysql { } match get_migrated_db(db_url.clone(), 1, 1).await { - Ok(x) => { + Ok(_) => { // it worked! yey! - test_mysql.conn = Some(x); break; } Err(err) => { @@ -168,8 +167,10 @@ impl TestMysql { test_mysql } - pub fn conn(&self) -> &DatabaseConnection { - self.conn.as_ref().unwrap() + pub async fn conn(&self) -> DatabaseConnection { + get_migrated_db(self.url.clone().unwrap(), 1, 5) + .await + .unwrap() } }