store: Make the connection timeout configurable and default to 5s

graphprotocol · Aug 27, 2021 · 6775ad9 · 6775ad9
1 parent e19e5e2
commit 6775ad9
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 1 deletion.
diff --git a/docs/environment-variables.md b/docs/environment-variables.md
@@ -156,3 +156,5 @@ those.
   given the other load management configuration settings, but never
   actually decline to run a query, instead log about load management
   decisions. Set to `true` to turn simulation on, defaults to `false`
+- `GRAPH_STORE_CONNECTION_TIMEOUT`: How long to wait to connect to a
+  database before assuming the database is down in ms. Defaults to 5000ms.
diff --git a/store/postgres/src/connection_pool.rs b/store/postgres/src/connection_pool.rs
@@ -45,6 +45,11 @@ lazy_static::lazy_static! {
             })
             .unwrap_or(0)
     };
+    static ref CONNECTION_TIMEOUT: Duration = {
+        std::env::var("GRAPH_STORE_CONNECTION_TIMEOUT").ok().map(|s| Duration::from_millis(u64::from_str(&s).unwrap_or_else(|_| {
+            panic!("GRAPH_STORE_CONNECTION_TIMEOUT must be a positive number, but is `{}`", s)
+        }))).unwrap_or(Duration::from_secs(5))
+    };
 }
 
 pub struct ForeignServer {
@@ -570,13 +575,15 @@ impl PoolInner {
         let builder: Builder<ConnectionManager<PgConnection>> = Pool::builder()
             .error_handler(error_handler.clone())
             .event_handler(event_handler.clone())
+            .connection_timeout(*CONNECTION_TIMEOUT)
             .max_size(pool_size);
         let pool = builder.build_unchecked(conn_manager);
         let fdw_pool = fdw_pool_size.map(|pool_size| {
             let conn_manager = ConnectionManager::new(postgres_url.clone());
             let builder: Builder<ConnectionManager<PgConnection>> = Pool::builder()
                 .error_handler(error_handler)
                 .event_handler(event_handler)
+                .connection_timeout(*CONNECTION_TIMEOUT)
                 .max_size(pool_size)
                 .min_idle(Some(1))
                 .idle_timeout(Some(FDW_IDLE_TIMEOUT));
@@ -707,7 +714,7 @@ impl PoolInner {
         logger: &Logger,
     ) -> Result<PooledConnection<ConnectionManager<PgConnection>>, StoreError> {
         loop {
-            match self.pool.get_timeout(Duration::from_secs(60)) {
+            match self.pool.get_timeout(*CONNECTION_TIMEOUT) {
                 Ok(conn) => return Ok(conn),
                 Err(e) => error!(logger, "Error checking out connection, retrying";
                    "error" => e.to_string(),