improve lazy pull

Author: levydsa

libsql/src/sync.rs

@@ -13,6 +13,7 @@ use uuid::Uuid;
 mod test;
 
 pub mod connection;
+pub mod statement;
 pub mod transaction;
 
 const METADATA_VERSION: u32 = 0;

libsql/src/sync/connection.rs

@@ -15,7 +15,7 @@ use std::sync::{
 use std::time::Duration;
 use tokio::sync::Mutex;
 
-use super::transaction::SyncedTx;
+use super::{statement::SyncedStatement, transaction::SyncedTx};
 
 #[derive(Clone)]
 pub struct SyncedConnection {
@@ -110,9 +110,10 @@ impl Conn for SyncedConnection {
 
     async fn execute_batch(&self, sql: &str) -> Result<BatchRows> {
         if self.should_execute_local(sql).await? {
-            if self.needs_pull.swap(false, Ordering::Relaxed) {
+            if self.needs_pull.load(Ordering::Relaxed) {
                 let mut context = self.context.lock().await;
                 crate::sync::try_pull(&mut context, &self.local).await?;
+                self.needs_pull.store(false, Ordering::Relaxed);
             }
             self.local.execute_batch(sql)
         } else {
@@ -122,9 +123,10 @@ impl Conn for SyncedConnection {
 
     async fn execute_transactional_batch(&self, sql: &str) -> Result<BatchRows> {
         if self.should_execute_local(sql).await? {
-            if self.needs_pull.swap(false, Ordering::Relaxed) {
+            if self.needs_pull.load(Ordering::Relaxed) {
                 let mut context = self.context.lock().await;
                 crate::sync::try_pull(&mut context, &self.local).await?;
+                self.needs_pull.store(false, Ordering::Relaxed);
             }
             self.local.execute_transactional_batch(sql)?;
             Ok(BatchRows::empty())
@@ -135,12 +137,17 @@ impl Conn for SyncedConnection {
 
     async fn prepare(&self, sql: &str) -> Result<Statement> {
         if self.should_execute_local(sql).await? {
-            if self.needs_pull.swap(false, Ordering::Relaxed) {
-                let mut context = self.context.lock().await;
-                crate::sync::try_pull(&mut context, &self.local).await?;
-            }
-            Ok(Statement {
+            let stmt = Statement {
                 inner: Box::new(LibsqlStmt(self.local.prepare(sql)?)),
+            };
+
+            Ok(Statement {
+                inner: Box::new(SyncedStatement {
+                    conn: self.local.clone(),
+                    inner: stmt,
+                    context: self.context.clone(),
+                    needs_pull: self.needs_pull.clone(),
+                }),
             })
         } else {
             let stmt = Statement {

libsql/src/sync/statement.rs

@@ -0,0 +1,74 @@
+use crate::{
+    local::{self},
+    params::Params,
+    statement::Stmt,
+    sync::SyncContext, Column, Result, Rows, Statement,
+};
+use std::sync::{atomic::{AtomicBool, Ordering}, Arc};
+use tokio::sync::Mutex;
+
+pub struct SyncedStatement {
+    pub conn: local::Connection,
+    pub inner: Statement,
+    pub context: Arc<Mutex<SyncContext>>,
+    pub needs_pull: Arc<AtomicBool>,
+}
+
+#[async_trait::async_trait]
+impl Stmt for SyncedStatement {
+    fn finalize(&mut self) {
+        self.inner.finalize()
+    }
+
+    async fn execute(&mut self, params: &Params) -> Result<usize> {
+        if self.needs_pull.load(Ordering::Relaxed) {
+            let mut context = self.context.lock().await;
+            crate::sync::try_pull(&mut context, &self.conn).await?;
+            self.needs_pull.store(false, Ordering::Relaxed);
+        }
+        self.inner.execute(params).await
+    }
+
+    async fn query(&mut self, params: &Params) -> Result<Rows> {
+        if self.needs_pull.load(Ordering::Relaxed) {
+            let mut context = self.context.lock().await;
+            crate::sync::try_pull(&mut context, &self.conn).await?;
+            self.needs_pull.store(false, Ordering::Relaxed);
+        }
+        self.inner.query(params).await
+    }
+
+    async fn run(&mut self, params: &Params) -> Result<()> {
+        if self.needs_pull.load(Ordering::Relaxed) {
+            let mut context = self.context.lock().await;
+            crate::sync::try_pull(&mut context, &self.conn).await?;
+            self.needs_pull.store(false, Ordering::Relaxed);
+        }
+        self.inner.run(params).await
+    }
+
+    fn interrupt(&mut self) -> Result<()> {
+        self.inner.interrupt()
+    }
+
+    fn reset(&mut self) {
+        self.inner.reset()
+    }
+
+    fn parameter_count(&self) -> usize {
+        self.inner.parameter_count()
+    }
+
+    fn parameter_name(&self, idx: i32) -> Option<&str> {
+        self.inner.parameter_name(idx)
+    }
+
+    fn column_count(&self) -> usize {
+        self.inner.column_count()
+    }
+
+    fn columns(&self) -> Vec<Column> {
+        self.inner.columns()
+    }
+}
+