store: Make ConnectionPool.try_get_fdw async

We might be blocking the tokio thread otherwise

Author: lutter

store/postgres/src/connection_pool.rs

@@ -641,16 +641,22 @@ impl ConnectionPool {
 
     /// Get a connection from the pool for foreign data wrapper access if
     /// one is available
-    pub fn try_get_fdw(
+    pub async fn try_get_fdw(
         &self,
         logger: &Logger,
         timeout: Duration,
     ) -> Option<PooledConnection<ConnectionManager<PgConnection>>> {
-        let Ok(inner) = self.get_ready() else {
-            return None;
-        };
-        self.state_tracker
-            .ignore_timeout(|| inner.try_get_fdw(logger, timeout))
+        let pool = self.clone();
+        let logger = logger.cheap_clone();
+        tokio::task::spawn_blocking(move || {
+            let Ok(inner) = pool.get_ready() else {
+                return None;
+            };
+            pool.state_tracker
+                .ignore_timeout(|| inner.try_get_fdw(&logger, timeout))
+        })
+        .await
+        .unwrap_or(None)
     }
 
     pub(crate) async fn query_permit(&self) -> QueryPermit {

store/postgres/src/copy.rs

@@ -1033,19 +1033,17 @@ impl Connection {
     /// Opportunistically create an extra worker if we have more tables to
     /// copy and there are idle fdw connections. If there are no more tables
     /// or no idle connections, this will return `None`.
-    fn extra_worker(
+    async fn extra_worker(
         &mut self,
         state: &mut CopyState,
         progress: &Arc<CopyProgress>,
     ) -> Option<Pin<Box<dyn Future<Output = WorkerResult>>>> {
         // 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
+        let conn = self
             .pool
             .try_get_fdw(&self.logger, ENV_VARS.store.batch_worker_wait)
-        else {
-            return None;
-        };
+            .await?;
         let Some(table) = state.unfinished.pop() else {
             return None;
         };
@@ -1134,7 +1132,7 @@ impl Connection {
                 if workers.len() >= self.workers {
                     break;
                 }
-                let Some(worker) = self.extra_worker(&mut state, &progress) else {
+                let Some(worker) = self.extra_worker(&mut state, &progress).await else {
                     break;
                 };
                 workers.add(worker);