Merge pull request #128 from dotindustries/feat/convert-database-proxy-shim-to-async-io-issue-65-jvs2zf1sytg17vkt89pn3y8s

feat: convert database proxy shim to async I/O

Author: tervezo-ai[bot]

crates/warp-compat/src/shims/database.rs

@@ -1,4 +1,17 @@
 //! Database proxy shim — transparent TCP interception for Postgres/MySQL/Redis.
+//!
+//! The actual database proxy implementation lives in `warpgrid-host::db_proxy`.
+//! This shim serves as the compatibility layer marker.
+//!
+//! # Async I/O (US-506)
+//!
+//! The database proxy supports async I/O via `tokio::net::TcpStream`. When an
+//! `AsyncConnectionFactory` is provided to the engine, connections use non-blocking
+//! I/O that releases the pool mutex during send/recv operations, enabling concurrent
+//! query execution across multiple connections within a single Wasm instance.
+//!
+//! The sync path is preserved as a fallback — if no async factory is configured,
+//! the proxy uses blocking `std::net::TcpStream` with `block_in_place`.
 
 pub struct DatabaseShim;
 

crates/warpgrid-host/src/db_proxy.rs

@@ -277,6 +277,11 @@ impl ConnectionPoolManager {
         self.draining.load(Ordering::Relaxed)
     }
 
+    /// Check if the pool manager has an async connection factory configured.
+    pub fn has_async_factory(&self) -> bool {
+        self.async_factory.is_some()
+    }
+
     /// Allocate the next connection handle.
     async fn allocate_handle(&self) -> u64 {
         let mut handle = self.next_handle.lock().await;

crates/warpgrid-host/src/db_proxy/host.rs

@@ -66,7 +66,11 @@ impl Host for DbProxyHost {
         let mgr = Arc::clone(&self.pool_manager);
 
         let handle = self.runtime_handle.clone();
-        tokio::task::block_in_place(|| handle.block_on(mgr.checkout(&key, password)))
+        if mgr.has_async_factory() {
+            tokio::task::block_in_place(|| handle.block_on(mgr.checkout_async(&key, password)))
+        } else {
+            tokio::task::block_in_place(|| handle.block_on(mgr.checkout(&key, password)))
+        }
     }
 
     fn send(&mut self, conn_handle: u64, data: Vec<u8>) -> Result<u32, String> {
@@ -79,8 +83,10 @@ impl Host for DbProxyHost {
         let mgr = Arc::clone(&self.pool_manager);
         let handle = self.runtime_handle.clone();
 
+        // Use send_query() which releases the mutex during I/O for concurrent access.
+        // Falls back to sync backend via block_in_place if no async backend is available.
         let sent = tokio::task::block_in_place(|| {
-            handle.block_on(mgr.send(conn_handle, &data))
+            handle.block_on(mgr.send_query(conn_handle, &data))
         })?;
 
         Ok(sent as u32)
@@ -96,8 +102,10 @@ impl Host for DbProxyHost {
         let mgr = Arc::clone(&self.pool_manager);
         let handle = self.runtime_handle.clone();
 
+        // Use receive_results() which releases the mutex during I/O.
+        // Falls back to sync backend via block_in_place if no async backend is available.
         tokio::task::block_in_place(|| {
-            handle.block_on(mgr.recv(conn_handle, max_bytes as usize))
+            handle.block_on(mgr.receive_results(conn_handle, max_bytes as usize))
         })
     }
 
@@ -118,6 +126,9 @@ impl Host for DbProxyHost {
 mod tests {
     use super::*;
     use super::super::{ConnectionBackend, ConnectionFactory, PoolConfig};
+    use super::super::async_io::{AsyncConnectionBackend, AsyncConnectionFactory, AsyncConnectFuture};
+    use std::future::Future;
+    use std::pin::Pin;
     use std::sync::atomic::{AtomicU64, Ordering};
     use std::time::Duration;
 
@@ -317,4 +328,121 @@ mod tests {
         // Reused — no new factory connect.
         assert_eq!(factory.connects.load(Ordering::Relaxed), 1);
     }
+
+    // ── Async path tests ─────────────────────────────────────────────
+
+    #[derive(Debug)]
+    struct MockAsyncBackend {
+        send_count: AtomicU64,
+    }
+
+    impl MockAsyncBackend {
+        fn new() -> Self {
+            Self {
+                send_count: AtomicU64::new(0),
+            }
+        }
+    }
+
+    impl AsyncConnectionBackend for MockAsyncBackend {
+        fn send_async<'a>(
+            &'a mut self,
+            data: &'a [u8],
+        ) -> Pin<Box<dyn Future<Output = Result<usize, String>> + Send + 'a>> {
+            self.send_count.fetch_add(1, Ordering::Relaxed);
+            let len = data.len();
+            Box::pin(async move { Ok(len) })
+        }
+
+        fn recv_async<'a>(
+            &'a mut self,
+            max_bytes: usize,
+        ) -> Pin<Box<dyn Future<Output = Result<Vec<u8>, String>> + Send + 'a>> {
+            Box::pin(async move { Ok(vec![0x42; max_bytes.min(4)]) })
+        }
+
+        fn ping_async(&mut self) -> Pin<Box<dyn Future<Output = bool> + Send + '_>> {
+            Box::pin(async { true })
+        }
+
+        fn close_async(&mut self) -> Pin<Box<dyn Future<Output = ()> + Send + '_>> {
+            Box::pin(async {})
+        }
+    }
+
+    struct TestAsyncFactory {
+        connects: AtomicU64,
+    }
+
+    impl TestAsyncFactory {
+        fn new() -> Self {
+            Self {
+                connects: AtomicU64::new(0),
+            }
+        }
+    }
+
+    impl AsyncConnectionFactory for TestAsyncFactory {
+        fn connect_async<'a>(
+            &'a self,
+            _key: &'a PoolKey,
+            _password: Option<&'a str>,
+        ) -> AsyncConnectFuture<'a> {
+            self.connects.fetch_add(1, Ordering::Relaxed);
+            Box::pin(async {
+                Ok(Box::new(MockAsyncBackend::new()) as Box<dyn AsyncConnectionBackend>)
+            })
+        }
+    }
+
+    fn make_async_host() -> DbProxyHost {
+        let sync_factory = Arc::new(TestFactory::new());
+        let async_factory = Arc::new(TestAsyncFactory::new());
+        let config = PoolConfig {
+            max_size: 5,
+            connect_timeout: Duration::from_millis(100),
+            ..PoolConfig::default()
+        };
+        let mgr = Arc::new(ConnectionPoolManager::new_with_async(
+            config,
+            sync_factory,
+            async_factory,
+        ));
+        let handle = tokio::runtime::Handle::current();
+        DbProxyHost::new(mgr, handle)
+    }
+
+    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+    async fn host_async_full_lifecycle() {
+        let mut host = make_async_host();
+
+        // Connect via async path.
+        let handle = host.connect(test_connect_config()).unwrap();
+
+        // Send via async send_query path.
+        let sent = host.send(handle, b"SELECT 1;".to_vec()).unwrap();
+        assert_eq!(sent, 9);
+
+        // Recv via async receive_results path.
+        let data = host.recv(handle, 1024).unwrap();
+        assert!(!data.is_empty());
+
+        // Close.
+        host.close(handle).unwrap();
+
+        // Handle invalid after close.
+        assert!(host.send(handle, b"x".to_vec()).is_err());
+    }
+
+    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+    async fn host_sync_fallback_when_no_async_factory() {
+        // make_host() uses sync-only factory — verify it still works.
+        let mut host = make_host();
+        let handle = host.connect(test_connect_config()).unwrap();
+        let sent = host.send(handle, b"data".to_vec()).unwrap();
+        assert_eq!(sent, 4);
+        let data = host.recv(handle, 1024).unwrap();
+        assert!(!data.is_empty());
+        host.close(handle).unwrap();
+    }
 }