Async terminal connections with tokio

Author: joshua-spacetime

tools/tpcc-runner/src/client.rs

@@ -7,6 +7,7 @@ use std::time::Duration;
 use crate::config::ConnectionConfig;
 use crate::module_bindings::*;
 use spacetimedb_sdk::{DbContext, Identity, Table as _};
+use tokio::sync::oneshot;
 
 pub struct ModuleClient {
     conn: DbConnection,
@@ -60,6 +61,55 @@ impl ModuleClient {
         })
     }
 
+    pub async fn connect_async(config: &ConnectionConfig, database_identity: Identity) -> Result<Self> {
+        let (ready_tx, ready_rx) = oneshot::channel();
+        let ready_tx = Arc::new(Mutex::new(Some(ready_tx)));
+        let success_tx = Arc::clone(&ready_tx);
+        let error_tx = Arc::clone(&ready_tx);
+        let disconnect_error = Arc::new(Mutex::new(None));
+        let disconnect_error_callback = Arc::clone(&disconnect_error);
+        let mut builder = DbConnection::builder()
+            .with_uri(config.uri.clone())
+            .with_database_name(database_identity.to_string())
+            .with_confirmed_reads(config.confirmed_reads)
+            .on_connect(move |_, _, _| {
+                if let Some(tx) = success_tx.lock().expect("ready mutex poisoned").take() {
+                    let _ = tx.send(Ok::<(), anyhow::Error>(()));
+                }
+            })
+            .on_connect_error(move |_, error| {
+                if let Some(tx) = error_tx.lock().expect("ready mutex poisoned").take() {
+                    let _ = tx.send(Err(anyhow!("connection failed: {error}")));
+                }
+            })
+            .on_disconnect(move |_, error| {
+                let message = match error {
+                    Some(error) => format!("connection closed: {error}"),
+                    None => "connection closed".to_string(),
+                };
+                *disconnect_error_callback.lock().expect("disconnect mutex poisoned") = Some(message);
+            });
+
+        if let Some(token) = &config.token {
+            builder = builder.with_token(Some(token.clone()));
+        }
+
+        let conn = builder.build().context("failed to build database connection")?;
+        let thread = conn.run_threaded();
+        tokio::time::timeout(Duration::from_secs(config.timeout_secs), ready_rx)
+            .await
+            .context("timed out waiting for connection")?
+            .map_err(|_| anyhow!("connection readiness callback dropped"))??;
+
+        Ok(Self {
+            conn,
+            thread: Some(thread),
+            timeout: Duration::from_secs(config.timeout_secs),
+            disconnect_error,
+            load_state_subscription: None,
+        })
+    }
+
     pub fn subscribe_load_state(&mut self) -> Result<()> {
         if self.load_state_subscription.is_some() {
             return Ok(());
@@ -340,28 +390,27 @@ impl ModuleClient {
         }
     }
 
-    pub fn new_order(
+    pub async fn new_order_async(
         &self,
         w_id: u32,
         d_id: u8,
         c_id: u32,
         order_lines: Vec<NewOrderLineInput>,
     ) -> Result<Result<NewOrderResult, String>> {
-        let (tx, rx) = sync_channel(1);
+        let (tx, rx) = oneshot::channel();
         self.conn
             .reducers
             .new_order_then(w_id, d_id, c_id, order_lines, move |_, res| {
                 log::debug!("Got response from `new_order`: {res:?}");
                 let _ = tx.send(res);
             })?;
-        match rx.recv_timeout(self.timeout) {
-            Ok(Ok(value)) => Ok(value),
-            Ok(Err(err)) => Err(anyhow!("new_order internal error: {}", err)),
-            Err(_) => bail!("timed out waiting for new_order"),
+        match self.await_callback("new_order", rx).await? {
+            Ok(value) => Ok(value),
+            Err(err) => Err(anyhow!("new_order internal error: {}", err)),
         }
     }
 
-    pub fn payment(
+    pub async fn payment_async(
         &self,
         w_id: u32,
         d_id: u8,
@@ -370,7 +419,7 @@ impl ModuleClient {
         customer: CustomerSelector,
         payment_amount_cents: i64,
     ) -> Result<Result<PaymentResult, String>> {
-        let (tx, rx) = sync_channel(1);
+        let (tx, rx) = oneshot::channel();
         self.conn.reducers.payment_then(
             w_id,
             d_id,
@@ -383,49 +432,51 @@ impl ModuleClient {
                 let _ = tx.send(res);
             },
         )?;
-        match rx.recv_timeout(self.timeout) {
-            Ok(Ok(value)) => Ok(value),
-            Ok(Err(err)) => Err(anyhow!("payment internal error: {}", err)),
-            Err(_) => bail!("timed out waiting for payment"),
+        match self.await_callback("payment", rx).await? {
+            Ok(value) => Ok(value),
+            Err(err) => Err(anyhow!("payment internal error: {}", err)),
         }
     }
 
-    pub fn order_status(
+    pub async fn order_status_async(
         &self,
         w_id: u32,
         d_id: u8,
         customer: CustomerSelector,
     ) -> Result<Result<OrderStatusResult, String>> {
-        let (tx, rx) = sync_channel(1);
+        let (tx, rx) = oneshot::channel();
         self.conn
             .reducers
             .order_status_then(w_id, d_id, customer, move |_, res| {
                 log::debug!("Got response from `order_status`: {res:?}");
                 let _ = tx.send(res);
             })?;
-        match rx.recv_timeout(self.timeout) {
-            Ok(Ok(value)) => Ok(value),
-            Ok(Err(err)) => Err(anyhow!("order_status internal error: {}", err)),
-            Err(_) => bail!("timed out waiting for order_status"),
+        match self.await_callback("order_status", rx).await? {
+            Ok(value) => Ok(value),
+            Err(err) => Err(anyhow!("order_status internal error: {}", err)),
         }
     }
 
-    pub fn stock_level(&self, w_id: u32, d_id: u8, threshold: i32) -> Result<Result<StockLevelResult, String>> {
-        let (tx, rx) = sync_channel(1);
+    pub async fn stock_level_async(
+        &self,
+        w_id: u32,
+        d_id: u8,
+        threshold: i32,
+    ) -> Result<Result<StockLevelResult, String>> {
+        let (tx, rx) = oneshot::channel();
         self.conn
             .reducers
             .stock_level_then(w_id, d_id, threshold, move |_, res| {
                 log::debug!("Got response from `stock_level`: {res:?}");
                 let _ = tx.send(res);
             })?;
-        match rx.recv_timeout(self.timeout) {
-            Ok(Ok(value)) => Ok(value),
-            Ok(Err(err)) => Err(anyhow!("stock_level internal error: {}", err)),
-            Err(_) => bail!("timed out waiting for stock_level"),
+        match self.await_callback("stock_level", rx).await? {
+            Ok(value) => Ok(value),
+            Err(err) => Err(anyhow!("stock_level internal error: {}", err)),
         }
     }
 
-    pub fn queue_delivery(
+    pub async fn queue_delivery_async(
         &self,
         run_id: String,
         driver_id: String,
@@ -434,7 +485,7 @@ impl ModuleClient {
         w_id: u32,
         carrier_id: u8,
     ) -> Result<Result<DeliveryQueueAck, String>> {
-        let (tx, rx) = sync_channel(1);
+        let (tx, rx) = oneshot::channel();
         self.conn.reducers.queue_delivery_then(
             run_id,
             driver_id,
@@ -447,50 +498,62 @@ impl ModuleClient {
                 let _ = tx.send(res);
             },
         )?;
-        match rx.recv_timeout(self.timeout) {
-            Ok(Ok(value)) => Ok(value),
-            Ok(Err(err)) => Err(anyhow!("queue_delivery internal error: {}", err)),
-            Err(_) => bail!("timed out waiting for queue_delivery"),
+        match self.await_callback("queue_delivery", rx).await? {
+            Ok(value) => Ok(value),
+            Err(err) => Err(anyhow!("queue_delivery internal error: {}", err)),
         }
     }
 
-    pub fn delivery_progress(&self, run_id: String) -> Result<Result<DeliveryProgress, String>> {
-        let (tx, rx) = sync_channel(1);
+    pub async fn delivery_progress_async(&self, run_id: String) -> Result<Result<DeliveryProgress, String>> {
+        let (tx, rx) = oneshot::channel();
         self.conn.reducers.delivery_progress_then(run_id, move |_, res| {
             log::debug!("Got response from `delivery_progress`: {res:?}");
             let _ = tx.send(res);
         })?;
-        match rx.recv_timeout(self.timeout) {
-            Ok(Ok(value)) => Ok(value),
-            Ok(Err(err)) => Err(anyhow!("delivery_progress internal error: {}", err)),
-            Err(_) => bail!("timed out waiting for delivery_progress"),
+        match self.await_callback("delivery_progress", rx).await? {
+            Ok(value) => Ok(value),
+            Err(err) => Err(anyhow!("delivery_progress internal error: {}", err)),
         }
     }
 
-    pub fn fetch_delivery_completions(
+    pub fn shutdown(mut self) {
+        let _ = self.conn.disconnect();
+        if let Some(thread) = self.thread.take() {
+            let _ = thread.join();
+        }
+    }
+
+    pub async fn fetch_delivery_completions_async(
         &self,
         run_id: String,
         after_completion_id: u64,
         limit: u32,
     ) -> Result<Result<Vec<DeliveryCompletionView>, String>> {
-        let (tx, rx) = sync_channel(1);
+        let (tx, rx) = oneshot::channel();
         self.conn
             .reducers
             .fetch_delivery_completions_then(run_id, after_completion_id, limit, move |_, res| {
                 log::debug!("Got response from `fetch_delivery_completions`: {res:?}");
                 let _ = tx.send(res);
             })?;
-        match rx.recv_timeout(self.timeout) {
-            Ok(Ok(value)) => Ok(value),
-            Ok(Err(err)) => Err(anyhow!("fetch_delivery_completions internal error: {}", err)),
-            Err(_) => bail!("timed out waiting for fetch_delivery_completions"),
+        match self.await_callback("fetch_delivery_completions", rx).await? {
+            Ok(value) => Ok(value),
+            Err(err) => Err(anyhow!("fetch_delivery_completions internal error: {}", err)),
         }
     }
 
-    pub fn shutdown(mut self) {
-        let _ = self.conn.disconnect();
-        if let Some(thread) = self.thread.take() {
-            let _ = thread.join();
+    pub async fn shutdown_async(self) {
+        let _ = tokio::task::spawn_blocking(move || self.shutdown()).await;
+    }
+
+    async fn await_callback<T>(&self, operation: &str, rx: oneshot::Receiver<T>) -> Result<T> {
+        match tokio::time::timeout(self.timeout, rx).await {
+            Ok(Ok(value)) => Ok(value),
+            Ok(Err(_)) => Err(anyhow!("{operation} callback dropped")),
+            Err(_) => {
+                self.ensure_connected()?;
+                bail!("timed out waiting for {operation}")
+            }
         }
     }
 }