Merge rust-lang#1344

1344: Allow database timeouts to be configured r=ashleygwilliams

We're consistently seeing the endpoint for crate downloads H12. The logs
are showing that we are *eventually* finishing without error, it's just
taking >30s, so Heroku is killing the connection.

Unfortunately, this means that we don't have visibility into *why* it's
timing out. It's either waiting for a long time for a DB connection,
updating the `version_downloads` table (my guess), or talking to S3. If
it's database related, this will turn those requests into actual errors.

This setting will increase our error rates, as we're seeing the 99th
percentile push up to 5s at some points, and we do have requests that
take 15s to complete (but are completing). Eventually I think we should
be able to consider 10s response times unacceptable, but for now I'm
only going to leave this at 10s for long enough for me to see *what* is
failing, at which point I will set it back to 30s until we no longer see
requests pushing 10s regularly.

Author: bors-voyager[bot]

src/app.rs

@@ -3,6 +3,7 @@
 use std::env;
 use std::path::PathBuf;
 use std::sync::{Arc, Mutex};
+use std::time::Duration;
 
 use diesel::r2d2;
 use git2;
@@ -71,11 +72,19 @@ impl App {
             _ => 1,
         };
 
+        let db_connection_timeout = match (env::var("DB_TIMEOUT"), config.env) {
+            (Ok(num), _) => num.parse().expect("couldn't parse DB_TIMEOUT"),
+            (_, ::Env::Production) => 10,
+            _ => 30,
+        };
+
         let thread_pool = Arc::new(ScheduledThreadPool::new(db_helper_threads));
 
         let diesel_db_config = r2d2::Pool::builder()
             .max_size(db_pool_size)
             .min_idle(db_min_idle)
+            .connection_timeout(Duration::from_secs(db_connection_timeout))
+            .connection_customizer(Box::new(db::SetStatementTimeout(db_connection_timeout)))
             .thread_pool(thread_pool);
 
         let repo = git2::Repository::open(&config.git_repo_checkout).unwrap();

src/db.rs

@@ -1,8 +1,8 @@
 use std::env;
 
 use conduit::Request;
-use diesel::prelude::{ConnectionResult, PgConnection};
-use diesel::r2d2::{self, ConnectionManager};
+use diesel::prelude::*;
+use diesel::r2d2::{self, ConnectionManager, CustomizeConnection};
 use url::Url;
 
 use middleware::app::RequestApp;
@@ -44,3 +44,17 @@ impl<T: Request + ?Sized> RequestTransaction for T {
         self.app().diesel_database.get().map_err(Into::into)
     }
 }
+
+#[derive(Debug, Clone, Copy)]
+pub struct SetStatementTimeout(pub u64);
+
+impl CustomizeConnection<PgConnection, r2d2::Error> for SetStatementTimeout {
+    fn on_acquire(&self, conn: &mut PgConnection) -> Result<(), r2d2::Error> {
+        use diesel::sql_query;
+
+        sql_query(format!("SET statement_timeout = {}", self.0 * 1000))
+            .execute(conn)
+            .map_err(r2d2::Error::QueryError)?;
+        Ok(())
+    }
+}