Use a single connection pool across the whole application

src/bin/cratesfyi.rs

@@ -2,11 +2,11 @@ use std::env;
 use std::path::PathBuf;
 use std::sync::Arc;
 
-use cratesfyi::db::{self, add_path_into_database, connect_db};
+use cratesfyi::db::{self, add_path_into_database, Pool};
 use cratesfyi::utils::{add_crate_to_queue, remove_crate_priority, set_crate_priority};
-use cratesfyi::Server;
-use cratesfyi::{DocBuilder, DocBuilderOptions, RustwideBuilder};
+use cratesfyi::{Config, DocBuilder, DocBuilderOptions, RustwideBuilder, Server};
 use failure::Error;
+use postgres::Connection;
 use structopt::StructOpt;
 
 pub fn main() -> Result<(), Error> {
@@ -82,25 +82,26 @@ enum CommandLine {
 
 impl CommandLine {
     pub fn handle_args(self) -> Result<(), Error> {
-        let config = Arc::new(cratesfyi::Config::from_env()?);
+        let config = Arc::new(Config::from_env()?);
+        let pool = Pool::new(&config)?;
 
         match self {
-            Self::Build(build) => build.handle_args(),
+            Self::Build(build) => build.handle_args(pool),
             Self::StartWebServer {
                 socket_addr,
                 reload_templates,
             } => {
-                Server::start(Some(&socket_addr), reload_templates, config)?;
+                Server::start(Some(&socket_addr), reload_templates, pool, config)?;
             }
             Self::Daemon { foreground } => {
                 if foreground {
                     log::warn!("--foreground was passed, but there is no need for it anymore");
                 }
 
-                cratesfyi::utils::start_daemon(config)?;
+                cratesfyi::utils::start_daemon(config, pool)?;
             }
-            Self::Database { subcommand } => subcommand.handle_args(),
-            Self::Queue { subcommand } => subcommand.handle_args(),
+            Self::Database { subcommand } => subcommand.handle_args(&*pool.get()?),
+            Self::Queue { subcommand } => subcommand.handle_args(&*pool.get()?),
         }
 
         Ok(())
@@ -135,19 +136,18 @@ enum QueueSubcommand {
 }
 
 impl QueueSubcommand {
-    pub fn handle_args(self) {
+    pub fn handle_args(self, conn: &Connection) {
         match self {
             Self::Add {
                 crate_name,
                 crate_version,
                 build_priority,
             } => {
-                let conn = connect_db().expect("Could not connect to database");
                 add_crate_to_queue(&conn, &crate_name, &crate_version, build_priority)
                     .expect("Could not add crate to queue");
             }
 
-            Self::DefaultPriority { subcommand } => subcommand.handle_args(),
+            Self::DefaultPriority { subcommand } => subcommand.handle_args(conn),
         }
     }
 }
@@ -172,18 +172,14 @@ enum PrioritySubcommand {
 }
 
 impl PrioritySubcommand {
-    pub fn handle_args(self) {
+    pub fn handle_args(self, conn: &Connection) {
         match self {
             Self::Set { pattern, priority } => {
-                let conn = connect_db().expect("Could not connect to the database");
-
                 set_crate_priority(&conn, &pattern, priority)
                     .expect("Could not set pattern's priority");
             }
 
             Self::Remove { pattern } => {
-                let conn = connect_db().expect("Could not connect to the database");
-
                 if let Some(priority) = remove_crate_priority(&conn, &pattern)
                     .expect("Could not remove pattern's priority")
                 {
@@ -237,7 +233,7 @@ struct Build {
 }
 
 impl Build {
-    pub fn handle_args(self) {
+    pub fn handle_args(self, pool: Pool) {
         let docbuilder = {
             let mut doc_options = DocBuilderOptions::from_prefix(self.prefix);
 
@@ -253,10 +249,10 @@ impl Build {
                 .check_paths()
                 .expect("The given paths were invalid");
 
-            DocBuilder::new(doc_options)
+            DocBuilder::new(doc_options, pool.clone())
         };
 
-        self.subcommand.handle_args(docbuilder);
+        self.subcommand.handle_args(docbuilder, pool);
     }
 }
 
@@ -304,12 +300,12 @@ enum BuildSubcommand {
 }
 
 impl BuildSubcommand {
-    pub fn handle_args(self, mut docbuilder: DocBuilder) {
+    pub fn handle_args(self, mut docbuilder: DocBuilder, pool: cratesfyi::db::Pool) {
         match self {
             Self::World => {
                 docbuilder.load_cache().expect("Failed to load cache");
 
-                let mut builder = RustwideBuilder::init().unwrap();
+                let mut builder = RustwideBuilder::init(pool).unwrap();
                 builder
                     .build_world(&mut docbuilder)
                     .expect("Failed to build world");
@@ -323,7 +319,8 @@ impl BuildSubcommand {
                 local,
             } => {
                 docbuilder.load_cache().expect("Failed to load cache");
-                let mut builder = RustwideBuilder::init().expect("failed to initialize rustwide");
+                let mut builder =
+                    RustwideBuilder::init(pool).expect("failed to initialize rustwide");
 
                 if let Some(path) = local {
                     builder
@@ -345,7 +342,7 @@ impl BuildSubcommand {
 
             Self::UpdateToolchain { only_first_time } => {
                 if only_first_time {
-                    let conn = db::connect_db().unwrap();
+                    let conn = pool.get().expect("failed to get a database connection");
                     let res = conn
                         .query("SELECT * FROM config WHERE name = 'rustc_version';", &[])
                         .unwrap();
@@ -356,14 +353,14 @@ impl BuildSubcommand {
                     }
                 }
 
-                let mut builder = RustwideBuilder::init().unwrap();
+                let mut builder = RustwideBuilder::init(pool).unwrap();
                 builder
                     .update_toolchain()
                     .expect("failed to update toolchain");
             }
 
             Self::AddEssentialFiles => {
-                let mut builder = RustwideBuilder::init().unwrap();
+                let mut builder = RustwideBuilder::init(pool).unwrap();
                 builder
                     .add_essential_files()
                     .expect("failed to add essential files");
@@ -378,7 +375,7 @@ impl BuildSubcommand {
 
 #[derive(Debug, Clone, PartialEq, Eq, StructOpt)]
 enum DatabaseSubcommand {
-    /// Run database migrations
+    /// Run database migration
     Migrate {
         /// The database version to migrate to
         #[structopt(name = "VERSION")]
@@ -415,33 +412,30 @@ enum DatabaseSubcommand {
 }
 
 impl DatabaseSubcommand {
-    pub fn handle_args(self) {
+    pub fn handle_args(self, conn: &Connection) {
         match self {
             Self::Migrate { version } => {
-                let conn = connect_db().expect("failed to connect to the database");
                 db::migrate(version, &conn).expect("Failed to run database migrations");
             }
 
             Self::UpdateGithubFields => {
-                cratesfyi::utils::github_updater().expect("Failed to update github fields");
+                cratesfyi::utils::github_updater(&conn).expect("Failed to update github fields");
             }
 
             Self::AddDirectory { directory, prefix } => {
-                let conn = db::connect_db().expect("failed to connect to the database");
                 add_path_into_database(&conn, &prefix, directory)
                     .expect("Failed to add directory into database");
             }
 
             // FIXME: This is actually util command not database
-            Self::UpdateReleaseActivity => cratesfyi::utils::update_release_activity()
+            Self::UpdateReleaseActivity => cratesfyi::utils::update_release_activity(&conn)
                 .expect("Failed to update release activity"),
 
             Self::DeleteCrate { crate_name } => {
-                let conn = db::connect_db().expect("failed to connect to the database");
                 db::delete_crate(&conn, &crate_name).expect("failed to delete the crate");
             }
 
-            Self::Blacklist { command } => command.handle_args(),
+            Self::Blacklist { command } => command.handle_args(&conn),
         }
     }
 }
@@ -467,9 +461,7 @@ enum BlacklistSubcommand {
 }
 
 impl BlacklistSubcommand {
-    fn handle_args(self) {
-        let conn = db::connect_db().expect("failed to connect to the database");
-
+    fn handle_args(self, conn: &Connection) {
         match self {
             Self::List => {
                 let crates =

src/config.rs

@@ -1,37 +1,59 @@
-use failure::{bail, Error, Fail, ResultExt};
+use failure::{bail, format_err, Error, Fail, ResultExt};
 use std::env::VarError;
 use std::str::FromStr;
-use std::sync::Arc;
 
 #[derive(Debug)]
 pub struct Config {
+    // Database connection params
+    pub(crate) database_url: String,
+    pub(crate) max_pool_size: u32,
+    pub(crate) min_pool_idle: u32,
+
+    // Max size of the files served by the docs.rs frontend
     pub(crate) max_file_size: usize,
     pub(crate) max_file_size_html: usize,
 }
 
 impl Config {
     pub fn from_env() -> Result<Self, Error> {
         Ok(Self {
+            database_url: require_env("CRATESFYI_DATABASE_URL")?,
+            max_pool_size: env("DOCSRS_MAX_POOL_SIZE", 90)?,
+            min_pool_idle: env("DOCSRS_MIN_POOL_IDLE", 10)?,
+
             max_file_size: env("DOCSRS_MAX_FILE_SIZE", 50 * 1024 * 1024)?,
             max_file_size_html: env("DOCSRS_MAX_FILE_SIZE_HTML", 5 * 1024 * 1024)?,
         })
     }
 }
 
-impl iron::typemap::Key for Config {
-    type Value = Arc<Config>;
+fn env<T>(var: &str, default: T) -> Result<T, Error>
+where
+    T: FromStr,
+    T::Err: Fail,
+{
+    Ok(maybe_env(var)?.unwrap_or(default))
+}
+
+fn require_env<T>(var: &str) -> Result<T, Error>
+where
+    T: FromStr,
+    T::Err: Fail,
+{
+    maybe_env(var)?.ok_or_else(|| format_err!("configuration variable {} is missing", var))
 }
 
-fn env<T>(var: &str, default: T) -> Result<T, Error>
+fn maybe_env<T>(var: &str) -> Result<Option<T>, Error>
 where
     T: FromStr,
     T::Err: Fail,
 {
     match std::env::var(var) {
         Ok(content) => Ok(content
             .parse::<T>()
+            .map(Some)
             .with_context(|_| format!("failed to parse configuration variable {}", var))?),
-        Err(VarError::NotPresent) => Ok(default),
+        Err(VarError::NotPresent) => Ok(None),
         Err(VarError::NotUnicode(_)) => bail!("configuration variable {} is not UTF-8", var),
     }
 }

src/db/mod.rs

@@ -5,62 +5,11 @@ pub(crate) use self::add_package::add_package_into_database;
 pub use self::delete_crate::delete_crate;
 pub use self::file::add_path_into_database;
 pub use self::migrate::migrate;
-
-use failure::Fail;
-use postgres::{Connection, TlsMode};
-use std::env;
+pub use self::pool::{Pool, PoolError};
 
 mod add_package;
 pub mod blacklist;
 mod delete_crate;
 pub(crate) mod file;
 mod migrate;
-
-/// Connects to database
-pub fn connect_db() -> Result<Connection, failure::Error> {
-    let err = "CRATESFYI_DATABASE_URL environment variable is not set";
-    let db_url = env::var("CRATESFYI_DATABASE_URL").map_err(|e| e.context(err))?;
-    Connection::connect(&db_url[..], TlsMode::None).map_err(Into::into)
-}
-
-pub(crate) fn create_pool() -> r2d2::Pool<r2d2_postgres::PostgresConnectionManager> {
-    let db_url = env::var("CRATESFYI_DATABASE_URL")
-        .expect("CRATESFYI_DATABASE_URL environment variable is not exists");
-
-    let max_pool_size = env::var("DOCSRS_MAX_POOL_SIZE")
-        .map(|s| {
-            s.parse::<u32>()
-                .expect("DOCSRS_MAX_POOL_SIZE must be an integer")
-        })
-        .unwrap_or(90);
-    crate::web::metrics::MAX_DB_CONNECTIONS.set(max_pool_size as i64);
-
-    let min_pool_idle = env::var("DOCSRS_MIN_POOL_IDLE")
-        .map(|s| {
-            s.parse::<u32>()
-                .expect("DOCSRS_MIN_POOL_IDLE must be an integer")
-        })
-        .unwrap_or(10);
-
-    let manager =
-        r2d2_postgres::PostgresConnectionManager::new(&db_url[..], r2d2_postgres::TlsMode::None)
-            .expect("Failed to create PostgresConnectionManager");
-
-    r2d2::Pool::builder()
-        .max_size(max_pool_size)
-        .min_idle(Some(min_pool_idle))
-        .build(manager)
-        .expect("Failed to create r2d2 pool")
-}
-
-#[cfg(test)]
-mod test {
-    use super::*;
-
-    #[test]
-    #[ignore]
-    fn test_connect_db() {
-        let conn = connect_db();
-        assert!(conn.is_ok());
-    }
-}
+mod pool;