graph, store: Wait for a little bit for an addl fdw connection

If we don't wait at all, we can use fewer connections than are available
since the pool might be below its capacity but has no idle connections open
currently

Author: lutter

graph/src/env/store.rs

@@ -91,6 +91,12 @@ pub struct EnvVarsStore {
     /// least 1
     pub batch_workers: usize,
 
+    /// How long to wait to get an additional connection for a batch worker.
+    /// This should just be big enough to allow the connection pool to
+    /// establish a connection. Set by `GRAPH_STORE_BATCH_WORKER_WAIT`.
+    /// Value is in ms and defaults to 2000ms
+    pub batch_worker_wait: Duration,
+
     /// Prune tables where we will remove at least this fraction of entity
     /// versions by rebuilding the table. Set by
     /// `GRAPH_STORE_HISTORY_REBUILD_THRESHOLD`. The default is 0.5
@@ -182,6 +188,7 @@ impl TryFrom<InnerStore> for EnvVarsStore {
             batch_target_duration: Duration::from_secs(x.batch_target_duration_in_secs),
             batch_timeout: x.batch_timeout_in_secs.map(Duration::from_secs),
             batch_workers: x.batch_workers,
+            batch_worker_wait: Duration::from_millis(x.batch_worker_wait),
             rebuild_threshold: x.rebuild_threshold.0,
             delete_threshold: x.delete_threshold.0,
             history_slack_factor: x.history_slack_factor.0,
@@ -251,6 +258,8 @@ pub struct InnerStore {
     batch_timeout_in_secs: Option<u64>,
     #[envconfig(from = "GRAPH_STORE_BATCH_WORKERS", default = "1")]
     batch_workers: usize,
+    #[envconfig(from = "GRAPH_STORE_BATCH_WORKER_WAIT", default = "2000")]
+    batch_worker_wait: u64,
     #[envconfig(from = "GRAPH_STORE_HISTORY_REBUILD_THRESHOLD", default = "0.5")]
     rebuild_threshold: ZeroToOneF64,
     #[envconfig(from = "GRAPH_STORE_HISTORY_DELETE_THRESHOLD", default = "0.05")]

store/postgres/src/connection_pool.rs

@@ -529,8 +529,9 @@ impl ConnectionPool {
     pub fn try_get_fdw(
         &self,
         logger: &Logger,
+        timeout: Duration,
     ) -> Result<Option<PooledConnection<ConnectionManager<PgConnection>>>, StoreError> {
-        self.get_ready()?.try_get_fdw(logger)
+        self.get_ready()?.try_get_fdw(logger, timeout)
     }
 
     pub fn connection_detail(&self) -> Result<ForeignServer, StoreError> {
@@ -1034,12 +1035,22 @@ impl PoolInner {
         }
     }
 
-    /// Get a connection from the fdw pool if one is available
+    /// Get a connection from the fdw pool if one is available. We wait for
+    /// `timeout` for a connection which should be set just big enough to
+    /// allow establishing a connection
     pub fn try_get_fdw(
         &self,
         logger: &Logger,
+        timeout: Duration,
     ) -> Result<Option<PooledConnection<ConnectionManager<PgConnection>>>, StoreError> {
-        Ok(self.fdw_pool(logger)?.try_get())
+        // Any error trying to get a connection is treated as "couldn't get
+        // a connection in time". If there is a serious error with the
+        // database, e.g., because it's not available, the next database
+        // operation will run into it and report it.
+        self.fdw_pool(logger)?
+            .get_timeout(timeout)
+            .map(|conn| Some(conn))
+            .or_else(|_| Ok(None))
     }
 
     pub fn connection_detail(&self) -> Result<ForeignServer, StoreError> {

store/postgres/src/copy.rs

@@ -914,7 +914,10 @@ impl Connection {
     ) -> Result<Option<Pin<Box<dyn Future<Output = CopyTableWorker>>>>, StoreError> {
         // It's important that we get the connection before the table since
         // we remove the table from the state and could drop it otherwise
-        let Some(conn) = self.pool.try_get_fdw(&self.logger)? else {
+        let Some(conn) = self
+            .pool
+            .try_get_fdw(&self.logger, ENV_VARS.store.batch_worker_wait)?
+        else {
             return Ok(None);
         };
         let Some(table) = state.unfinished.pop() else {