store: Use AsyncPgConnection for fdw pool

lutter · lutter · commit f58d79e83aca · 2025-10-18T14:01:02.000+02:00
diff --git a/store/postgres/src/copy.rs b/store/postgres/src/copy.rs
@@ -23,11 +23,14 @@ use std::{
 };
 
 use diesel::{
-    connection::SimpleConnection as _, dsl::sql, insert_into, select, sql_query, update,
-    ExpressionMethods, OptionalExtension, QueryDsl,
+    dsl::sql, insert_into, select, sql_query, update, ExpressionMethods, OptionalExtension,
+    QueryDsl,
 };
-use diesel_async::scoped_futures::{ScopedBoxFuture, ScopedFutureExt};
-use diesel_async::RunQueryDsl;
+use diesel_async::{
+    scoped_futures::{ScopedBoxFuture, ScopedFutureExt},
+    AsyncConnection,
+};
+use diesel_async::{RunQueryDsl, SimpleAsyncConnection};
 use graph::{
     futures03::{
         future::{select_all, BoxFuture},
@@ -46,12 +49,11 @@ use itertools::Itertools;
 use crate::{
     advisory_lock, catalog, deployment,
     dynds::DataSourcesTable,
-    pool::PgConnection,
     primary::{DeploymentId, Primary, Site},
     relational::{index::IndexList, Layout, Table},
     relational_queries as rq,
     vid_batcher::{VidBatcher, VidRange},
-    AsyncConnection, AsyncPgConnection, ConnectionPool,
+    AsyncPgConnection, ConnectionPool,
 };
 
 const LOG_INTERVAL: Duration = Duration::from_secs(3 * 60);
@@ -690,12 +692,12 @@ impl From<Result<CopyTableWorker, StoreError>> for WorkerResult {
 /// This struct helps us with that. It wraps a connection and tracks whether
 /// the connection was used to acquire the copy lock
 struct LockTrackingConnection {
-    inner: PgConnection,
+    inner: AsyncPgConnection,
     has_lock: bool,
 }
 
 impl LockTrackingConnection {
-    fn new(inner: PgConnection) -> Self {
+    fn new(inner: AsyncPgConnection) -> Self {
         Self {
             inner,
             has_lock: false,
@@ -804,10 +806,10 @@ impl CopyTableWorker {
                     }
 
                     match conn
-                        .transaction_async(|conn| {
+                        .transaction(|conn| {
                             async {
                                 if let Some(timeout) = BATCH_STATEMENT_TIMEOUT.as_ref() {
-                                    conn.batch_execute(timeout)?;
+                                    conn.batch_execute(timeout).await?;
                                 }
                                 self.table.copy_batch(conn).await
                             }
@@ -850,7 +852,7 @@ impl CopyTableWorker {
                     // that is hard to predict. This mechanism ensures
                     // that if our estimation is wrong, the consequences
                     // aren't too severe.
-                    conn.transaction_async(|conn| self.table.set_batch_size(conn, 1).scope_boxed())
+                    conn.transaction(|conn| self.table.set_batch_size(conn, 1).scope_boxed())
                         .await?;
                 }
             };
@@ -972,7 +974,7 @@ impl Connection {
 
         let mut last_log = Instant::now();
         let conn = pool
-            .get_fdw_async(&logger, || {
+            .get_fdw(&logger, || {
                 if last_log.elapsed() > LOG_INTERVAL {
                     info!(&logger, "waiting for other copy operations to finish");
                     last_log = Instant::now();
@@ -1005,7 +1007,9 @@ impl Connection {
         callback: F,
     ) -> Result<BoxFuture<'conn, Result<R, StoreError>>, StoreError>
     where
-        F: for<'r> FnOnce(&'r mut PgConnection) -> ScopedBoxFuture<'a, 'r, Result<R, StoreError>>
+        F: for<'r> FnOnce(
+                &'r mut AsyncPgConnection,
+            ) -> ScopedBoxFuture<'a, 'r, Result<R, StoreError>>
             + Send
             + 'a,
         R: Send + 'a,
@@ -1017,7 +1021,7 @@ impl Connection {
             ));
         };
         let conn = &mut conn.inner;
-        Ok(conn.transaction_async(|conn| async { callback(conn).await }.scope_boxed()))
+        Ok(conn.transaction(|conn| callback(conn).scope_boxed()))
     }
 
     /// Copy private data sources if the source uses a schema version that
@@ -1080,7 +1084,7 @@ impl Connection {
         // we remove the table from the state and could drop it otherwise
         let Some(conn) = self
             .pool
-            .try_get_fdw_async(&self.logger, ENV_VARS.store.batch_worker_wait)
+            .try_get_fdw(&self.logger, ENV_VARS.store.batch_worker_wait)
             .await
         else {
             return None;
diff --git a/store/postgres/src/pool/mod.rs b/store/postgres/src/pool/mod.rs
@@ -334,30 +334,24 @@ impl ConnectionPool {
     /// The `timeout` is called every time we time out waiting for a
     /// connection. If `timeout` returns `true`, `get_fdw` returns with that
     /// error, otherwise we try again to get a connection.
-    pub fn get_fdw<F>(&self, logger: &Logger, timeout: F) -> Result<PgConnection, StoreError>
-    where
-        F: FnMut() -> bool,
-    {
-        self.get_ready()?.get_fdw(logger, timeout)
-    }
-
-    /// An async version of `get_fdw`. For now, this calls `get_fdw`
-    /// synchronously. Once `get_fdw` is not used anymore, we can make it
-    /// truly async.
-    pub async fn get_fdw_async<F>(
+    pub async fn get_fdw<F>(
         &self,
         logger: &Logger,
         timeout: F,
-    ) -> Result<PgConnection, StoreError>
+    ) -> Result<AsyncPgConnection, StoreError>
     where
         F: FnMut() -> bool,
     {
-        self.get_fdw(logger, timeout)
+        self.get_ready()?.get_fdw(logger, timeout).await
     }
 
     /// Get a connection from the pool for foreign data wrapper access if
     /// one is available
-    async fn try_get_fdw(&self, logger: &Logger, timeout: Duration) -> Option<PgConnection> {
+    pub async fn try_get_fdw(
+        &self,
+        logger: &Logger,
+        timeout: Duration,
+    ) -> Option<AsyncPgConnection> {
         let Ok(inner) = self.get_ready() else {
             return None;
         };
@@ -366,17 +360,6 @@ impl ConnectionPool {
             .await
     }
 
-    /// An async version of `try_get_fdw`. For now, this calls `try_get_fdw`
-    /// synchronously. Once `try_get_fdw` is not used anymore, we can make it
-    /// truly async.
-    pub async fn try_get_fdw_async(
-        &self,
-        logger: &Logger,
-        timeout: Duration,
-    ) -> Option<PgConnection> {
-        self.try_get_fdw(logger, timeout).await
-    }
-
     pub(crate) async fn query_permit(&self) -> QueryPermit {
         let pool = self.inner.get_unready();
         let start = Instant::now();
@@ -565,14 +548,18 @@ impl PoolInner {
     /// The `timeout` is called every time we time out waiting for a
     /// connection. If `timeout` returns `true`, `get_fdw` returns with that
     /// error, otherwise we try again to get a connection.
-    fn get_fdw<F>(&self, logger: &Logger, mut timeout: F) -> Result<PgConnection, StoreError>
+    async fn get_fdw<F>(
+        &self,
+        logger: &Logger,
+        mut timeout: F,
+    ) -> Result<AsyncPgConnection, StoreError>
     where
         F: FnMut() -> bool,
     {
         let pool = self.fdw_pool(logger)?;
         loop {
-            match graph::block_on(pool.get()) {
-                Ok(conn) => return Ok(AsyncConnectionWrapper::from(conn)),
+            match pool.get().await {
+                Ok(conn) => return Ok(conn),
                 Err(e) => {
                     if timeout() {
                         return Err(anyhow!("timeout in get_fdw: {e}").into());
@@ -585,18 +572,18 @@ impl PoolInner {
     /// 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
-    fn try_get_fdw(&self, logger: &Logger, timeout: Duration) -> Option<PgConnection> {
+    async fn try_get_fdw(&self, logger: &Logger, timeout: Duration) -> Option<AsyncPgConnection> {
         // 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.
         let Ok(fdw_pool) = self.fdw_pool(logger) else {
             return None;
         };
-        let Ok(conn) = graph::block_on(fdw_pool.get_timeout(timeout)) else {
+        let Ok(conn) = fdw_pool.get_timeout(timeout).await else {
             return None;
         };
-        Some(AsyncConnectionWrapper::from(conn))
+        Some(conn)
     }
 
     pub fn connection_detail(&self) -> Result<ForeignServer, StoreError> {
diff --git a/store/postgres/src/pool/state_tracker.rs b/store/postgres/src/pool/state_tracker.rs
@@ -50,17 +50,6 @@ impl StateTracker {
     }
 
     pub(super) async fn ignore_timeout<F, R>(&self, f: F) -> R
-    where
-        F: FnOnce() -> R,
-    {
-        self.ignore_timeout.store(true, Ordering::Relaxed);
-        let res = f();
-        self.ignore_timeout.store(false, Ordering::Relaxed);
-        res
-    }
-
-    #[allow(dead_code)]
-    pub(super) async fn ignore_timeout_async<F, R>(&self, f: F) -> R
     where
         F: AsyncFnOnce() -> R,
     {
