Merge pull request #759 from Kixiron/r2u2

Author: Pietro Albini

src/db/mod.rs

@@ -28,10 +28,19 @@ pub fn connect_db() -> Result<Connection, failure::Error> {
 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 pool_size = env::var("DOCSRS_POOL_SIZE")
+        .map(|s| {
+            s.parse::<u32>()
+                .expect("DOCSRS_POOL_SIZE must be an integer")
+        })
+        .unwrap_or(90);
+
     let manager =
         r2d2_postgres::PostgresConnectionManager::new(&db_url[..], r2d2_postgres::TlsMode::None)
             .expect("Failed to create PostgresConnectionManager");
+
     r2d2::Pool::builder()
+        .max_size(pool_size)
         .build(manager)
         .expect("Failed to create r2d2 pool")
 }

src/web/builds.rs

@@ -65,7 +65,7 @@ pub fn build_list_handler(req: &mut Request) -> IronResult<Response> {
     let version = cexpect!(router.find("version"));
     let req_build_id: i32 = router.find("id").unwrap_or("0").parse().unwrap_or(0);
 
-    let conn = extension!(req, Pool).get();
+    let conn = extension!(req, Pool).get()?;
     let limits = ctry!(Limits::for_crate(&conn, name));
 
     let mut build_list: Vec<Build> = Vec::new();

src/web/crate_details.rs

@@ -315,7 +315,7 @@ pub fn crate_details_handler(req: &mut Request) -> IronResult<Response> {
     let name = cexpect!(router.find("name"));
     let req_version = router.find("version");
 
-    let conn = extension!(req, Pool).get();
+    let conn = extension!(req, Pool).get()?;
 
     match match_version(&conn, &name, req_version).and_then(|m| m.assume_exact()) {
         Some(MatchSemver::Exact((version, _))) => {

src/web/file.rs

@@ -47,7 +47,7 @@ pub struct DatabaseFileHandler;
 impl Handler for DatabaseFileHandler {
     fn handle(&self, req: &mut Request) -> IronResult<Response> {
         let path = req.url.path().join("/");
-        let conn = extension!(req, Pool).get();
+        let conn = extension!(req, Pool).get()?;
         if let Some(file) = File::from_path(&conn, &path) {
             Ok(file.serve())
         } else {

src/web/metrics.rs

@@ -65,10 +65,22 @@ lazy_static::lazy_static! {
         &["route"]
     )
     .unwrap();
+
+    pub static ref FAILED_DB_CONNECTIONS: IntCounter = register_int_counter!(
+        "docsrs_failed_db_connections",
+        "Number of attempted and failed connections to the database"
+    )
+    .unwrap();
+
+    pub static ref CONCURRENT_DB_CONNECTIONS: IntGauge = register_int_gauge!(
+        "docsrs_db_connections",
+        "The number of currently used database connections"
+    )
+    .unwrap();
 }
 
 pub fn metrics_handler(req: &mut Request) -> IronResult<Response> {
-    let conn = extension!(req, Pool).get();
+    let conn = extension!(req, Pool).get()?;
 
     QUEUED_CRATES_COUNT.set(
         ctry!(conn.query("SELECT COUNT(*) FROM queue WHERE attempt < 5;", &[]))
@@ -81,6 +93,9 @@ pub fn metrics_handler(req: &mut Request) -> IronResult<Response> {
             .get(0),
     );
 
+    let pool = extension!(req, Pool);
+    CONCURRENT_DB_CONNECTIONS.set(pool.connections() as i64);
+
     let mut buffer = Vec::new();
     let families = prometheus::gather();
     ctry!(TextEncoder::new().encode(&families, &mut buffer));

src/web/mod.rs

@@ -84,9 +84,6 @@ const OPENSEARCH_XML: &[u8] = include_bytes!("opensearch.xml");
 
 const DEFAULT_BIND: &str = "0.0.0.0:3000";
 
-type PoolFactoryFn = dyn Fn() -> Pool + Send + Sync;
-type PoolFactory = Box<PoolFactoryFn>;
-
 fn handlebars_engine() -> Result<HandlebarsEngine, SourceError> {
     // TODO: Use DocBuilderOptions for paths
     let mut hbse = HandlebarsEngine::new();
@@ -103,25 +100,25 @@ struct CratesfyiHandler {
     router_handler: Box<dyn Handler>,
     database_file_handler: Box<dyn Handler>,
     static_handler: Box<dyn Handler>,
-    pool_factory: PoolFactory,
+    pool: Pool,
 }
 
 impl CratesfyiHandler {
-    fn chain<H: Handler>(pool_factory: &PoolFactoryFn, base: H) -> Chain {
+    fn chain<H: Handler>(pool: Pool, base: H) -> Chain {
         let hbse = handlebars_engine().expect("Failed to load handlebar templates");
 
         let mut chain = Chain::new(base);
-        chain.link_before(pool_factory());
+        chain.link_before(pool);
         chain.link_after(hbse);
         chain
     }
 
-    fn new(pool_factory: PoolFactory) -> CratesfyiHandler {
+    fn new(pool: Pool) -> CratesfyiHandler {
         let routes = routes::build_routes();
         let blacklisted_prefixes = routes.page_prefixes();
 
-        let shared_resources = Self::chain(&pool_factory, rustdoc::SharedResourceHandler);
-        let router_chain = Self::chain(&pool_factory, routes.iron_router());
+        let shared_resources = Self::chain(pool.clone(), rustdoc::SharedResourceHandler);
+        let router_chain = Self::chain(pool.clone(), routes.iron_router());
         let prefix = PathBuf::from(
             env::var("CRATESFYI_PREFIX")
                 .expect("the CRATESFYI_PREFIX environment variable is not set"),
@@ -138,7 +135,7 @@ impl CratesfyiHandler {
                 Box::new(file::DatabaseFileHandler),
             )),
             static_handler: Box::new(static_handler),
-            pool_factory,
+            pool,
         }
     }
 }
@@ -191,7 +188,7 @@ impl Handler for CratesfyiHandler {
                     debug!("Path not found: {}", DebugPath(&req.url));
                 }
 
-                Self::chain(&self.pool_factory, err).handle(req)
+                Self::chain(self.pool.clone(), err).handle(req)
             })
     }
 }
@@ -354,28 +351,26 @@ pub struct Server {
 
 impl Server {
     pub fn start(addr: Option<&str>) -> Self {
-        let server = Self::start_inner(addr.unwrap_or(DEFAULT_BIND), Box::new(Pool::new));
+        let server = Self::start_inner(addr.unwrap_or(DEFAULT_BIND), Pool::new());
         info!("Running docs.rs web server on http://{}", server.addr());
         server
     }
 
     #[cfg(test)]
     pub(crate) fn start_test(conn: Arc<Mutex<Connection>>) -> Self {
-        Self::start_inner(
-            "127.0.0.1:0",
-            Box::new(move || Pool::new_simple(conn.clone())),
-        )
+        Self::start_inner("127.0.0.1:0", Pool::new_simple(conn.clone()))
     }
 
-    fn start_inner(addr: &str, pool_factory: PoolFactory) -> Self {
+    fn start_inner(addr: &str, pool: Pool) -> Self {
         // poke all the metrics counters to instantiate and register them
         metrics::TOTAL_BUILDS.inc_by(0);
         metrics::SUCCESSFUL_BUILDS.inc_by(0);
         metrics::FAILED_BUILDS.inc_by(0);
         metrics::NON_LIBRARY_BUILDS.inc_by(0);
         metrics::UPLOADED_FILES_TOTAL.inc_by(0);
+        metrics::FAILED_DB_CONNECTIONS.inc_by(0);
 
-        let cratesfyi = CratesfyiHandler::new(pool_factory);
+        let cratesfyi = CratesfyiHandler::new(pool);
         let inner = Iron::new(cratesfyi)
             .http(addr)
             .unwrap_or_else(|_| panic!("Failed to bind to socket on {}", addr));

src/web/pool.rs

@@ -1,10 +1,12 @@
 use crate::db::create_pool;
-use iron::{typemap, BeforeMiddleware, IronResult, Request};
+use iron::{status::Status, typemap, BeforeMiddleware, IronError, IronResult, Request};
 use postgres::Connection;
+use std::marker::PhantomData;
 
 #[cfg(test)]
 use std::sync::{Arc, Mutex, MutexGuard};
 
+#[derive(Debug, Clone)]
 pub(crate) enum Pool {
     R2D2(r2d2::Pool<r2d2_postgres::PostgresConnectionManager>),
     #[cfg(test)]
@@ -20,43 +22,55 @@ impl Pool {
     pub(crate) fn new_simple(conn: Arc<Mutex<Connection>>) -> Self {
         Pool::Simple(conn)
     }
-}
 
-impl typemap::Key for Pool {
-    type Value = PoolConnection;
-}
+    pub(super) fn get<'a>(&'a self) -> IronResult<DerefConnection<'a>> {
+        match self {
+            Self::R2D2(conn) => {
+                let conn = conn.get().map_err(|err| {
+                    log::error!("Error getting db connection: {:?}", err);
+                    super::metrics::FAILED_DB_CONNECTIONS.inc();
+
+                    IronError::new(err, Status::InternalServerError)
+                })?;
+
+                Ok(DerefConnection::Connection(conn, PhantomData))
+            }
 
-impl BeforeMiddleware for Pool {
-    fn before(&self, req: &mut Request) -> IronResult<()> {
-        req.extensions.insert::<Pool>(match self {
-            Self::R2D2(pool) => PoolConnection::R2D2(pool.get().unwrap()),
             #[cfg(test)]
-            Self::Simple(mutex) => PoolConnection::Simple(mutex.clone()),
-        });
-        Ok(())
+            Self::Simple(mutex) => Ok(DerefConnection::Guard(
+                mutex.lock().expect("failed to lock the connection"),
+            )),
+        }
     }
-}
-
-pub(crate) enum PoolConnection {
-    R2D2(r2d2::PooledConnection<r2d2_postgres::PostgresConnectionManager>),
-    #[cfg(test)]
-    Simple(Arc<Mutex<Connection>>),
-}
 
-impl PoolConnection {
-    pub(super) fn get(&self) -> DerefConnection<'_> {
+    pub(crate) fn connections(&self) -> u32 {
         match self {
-            Self::R2D2(conn) => DerefConnection::Connection(&conn),
+            Self::R2D2(conn) => conn.state().connections,
+
             #[cfg(test)]
-            Self::Simple(mutex) => {
-                DerefConnection::Guard(mutex.lock().expect("failed to lock the connection"))
-            }
+            Self::Simple(..) => 0,
         }
     }
 }
 
+impl typemap::Key for Pool {
+    type Value = Pool;
+}
+
+impl BeforeMiddleware for Pool {
+    fn before(&self, req: &mut Request) -> IronResult<()> {
+        req.extensions.insert::<Pool>(self.clone());
+
+        Ok(())
+    }
+}
+
 pub(crate) enum DerefConnection<'a> {
-    Connection(&'a Connection),
+    Connection(
+        r2d2::PooledConnection<r2d2_postgres::PostgresConnectionManager>,
+        PhantomData<&'a ()>,
+    ),
+
     #[cfg(test)]
     Guard(MutexGuard<'a, Connection>),
 }
@@ -66,7 +80,8 @@ impl<'a> std::ops::Deref for DerefConnection<'a> {
 
     fn deref(&self) -> &Connection {
         match self {
-            Self::Connection(conn) => conn,
+            Self::Connection(conn, ..) => conn,
+
             #[cfg(test)]
             Self::Guard(guard) => &guard,
         }

src/web/releases.rs

@@ -335,7 +335,7 @@ fn get_search_results(
 }
 
 pub fn home_page(req: &mut Request) -> IronResult<Response> {
-    let conn = extension!(req, Pool).get();
+    let conn = extension!(req, Pool).get()?;
     let packages = get_releases(&conn, 1, RELEASES_IN_HOME, Order::ReleaseTime);
     Page::new(packages)
         .set_true("show_search_form")
@@ -344,7 +344,7 @@ pub fn home_page(req: &mut Request) -> IronResult<Response> {
 }
 
 pub fn releases_feed_handler(req: &mut Request) -> IronResult<Response> {
-    let conn = extension!(req, Pool).get();
+    let conn = extension!(req, Pool).get()?;
     let packages = get_releases(&conn, 1, RELEASES_IN_FEED, Order::ReleaseTime);
     let mut resp = ctry!(Page::new(packages).to_resp("releases_feed"));
     resp.headers.set(::iron::headers::ContentType(
@@ -391,7 +391,7 @@ pub fn recent_releases_handler(req: &mut Request) -> IronResult<Response> {
         .unwrap_or("1")
         .parse()
         .unwrap_or(1);
-    let conn = extension!(req, Pool).get();
+    let conn = extension!(req, Pool).get()?;
     let packages = get_releases(&conn, page_number, RELEASES_IN_RELEASES, Order::ReleaseTime);
     releases_handler(
         packages,
@@ -408,7 +408,7 @@ pub fn releases_by_stars_handler(req: &mut Request) -> IronResult<Response> {
         .unwrap_or("1")
         .parse()
         .unwrap_or(1);
-    let conn = extension!(req, Pool).get();
+    let conn = extension!(req, Pool).get()?;
     let packages = get_releases(&conn, page_number, RELEASES_IN_RELEASES, Order::GithubStars);
     releases_handler(
         packages,
@@ -425,7 +425,7 @@ pub fn releases_recent_failures_handler(req: &mut Request) -> IronResult<Respons
         .unwrap_or("1")
         .parse()
         .unwrap_or(1);
-    let conn = extension!(req, Pool).get();
+    let conn = extension!(req, Pool).get()?;
     let packages = get_releases(
         &conn,
         page_number,
@@ -447,7 +447,7 @@ pub fn releases_failures_by_stars_handler(req: &mut Request) -> IronResult<Respo
         .unwrap_or("1")
         .parse()
         .unwrap_or(1);
-    let conn = extension!(req, Pool).get();
+    let conn = extension!(req, Pool).get()?;
     let packages = get_releases(
         &conn,
         page_number,
@@ -468,7 +468,7 @@ pub fn author_handler(req: &mut Request) -> IronResult<Response> {
     // page number of releases
     let page_number: i64 = router.find("page").unwrap_or("1").parse().unwrap_or(1);
 
-    let conn = extension!(req, Pool).get();
+    let conn = extension!(req, Pool).get()?;
 
     #[allow(clippy::or_fun_call)]
     let author = ctry!(router
@@ -518,7 +518,7 @@ pub fn search_handler(req: &mut Request) -> IronResult<Response> {
     let params = ctry!(req.get::<Params>());
     let query = params.find(&["query"]);
 
-    let conn = extension!(req, Pool).get();
+    let conn = extension!(req, Pool).get()?;
     if let Some(&Value::String(ref query)) = query {
         // check if I am feeling lucky button pressed and redirect user to crate page
         // if there is a match
@@ -618,7 +618,7 @@ pub fn search_handler(req: &mut Request) -> IronResult<Response> {
 }
 
 pub fn activity_handler(req: &mut Request) -> IronResult<Response> {
-    let conn = extension!(req, Pool).get();
+    let conn = extension!(req, Pool).get()?;
     let release_activity_data: Json = ctry!(conn.query(
         "SELECT value FROM config WHERE name = 'release_activity'",
         &[]
@@ -635,7 +635,7 @@ pub fn activity_handler(req: &mut Request) -> IronResult<Response> {
 }
 
 pub fn build_queue_handler(req: &mut Request) -> IronResult<Response> {
-    let conn = extension!(req, Pool).get();
+    let conn = extension!(req, Pool).get()?;
     let mut crates: Vec<(String, String, i32)> = Vec::new();
     for krate in &conn
         .query(