libsql: Remove lazy pulling optimization

The optimization attempts to reduce pulling when an application is
writing. However, the logic is embedded in prepare(), which is totally
wrong because a statement can be reused. Let's remove the optimization
as incorrect.

Refs: tursodatabase/turso-cloud#5

Author: penberg

libsql/src/database.rs

@@ -719,7 +719,6 @@ impl Database {
                         read_your_writes: *read_your_writes,
                         context: db.sync_ctx.clone().unwrap(),
                         state: std::sync::Arc::new(Mutex::new(State::Init)),
-                        needs_pull: std::sync::atomic::AtomicBool::new(false).into(),
                     };
 
                     let conn = std::sync::Arc::new(synced);

libsql/src/sync/connection.rs

@@ -8,10 +8,7 @@ use crate::{
     sync::SyncContext,
     BatchRows, Error, Result, Statement, Transaction, TransactionBehavior,
 };
-use std::sync::{
-    atomic::{AtomicBool, Ordering},
-    Arc,
-};
+use std::sync::Arc;
 use std::time::Duration;
 use tokio::sync::Mutex;
 
@@ -24,7 +21,6 @@ pub struct SyncedConnection {
     pub read_your_writes: bool,
     pub context: Arc<Mutex<SyncContext>>,
     pub state: Arc<Mutex<State>>,
-    pub needs_pull: Arc<AtomicBool>,
 }
 
 impl SyncedConnection {
@@ -110,11 +106,6 @@ impl Conn for SyncedConnection {
 
     async fn execute_batch(&self, sql: &str) -> Result<BatchRows> {
         if self.should_execute_local(sql).await? {
-            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 {
             self.remote.execute_batch(sql).await
@@ -123,11 +114,6 @@ 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.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())
         } else {
@@ -137,28 +123,26 @@ impl Conn for SyncedConnection {
 
     async fn prepare(&self, sql: &str) -> Result<Statement> {
         if self.should_execute_local(sql).await? {
-            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(),
-                }),
+                inner: Box::new(LibsqlStmt(self.local.prepare(sql)?)),
             })
         } else {
             let stmt = Statement {
                 inner: Box::new(self.remote.prepare(sql).await?),
             };
 
             if self.read_your_writes {
-                self.needs_pull.store(true, Ordering::Relaxed);
+                Ok(Statement {
+                    inner: Box::new(SyncedStatement {
+                        conn: self.local.clone(),
+                        inner: stmt,
+                        context: self.context.clone(),
+                        pull_after: true,
+                    }),
+                })
+            } else {
+                Ok(stmt)
             }
-
-            Ok(stmt)
         }
     }
 

libsql/src/sync/statement.rs

@@ -4,14 +4,14 @@ use crate::{
     statement::Stmt,
     sync::SyncContext, Column, Result, Rows, Statement,
 };
-use std::sync::{atomic::{AtomicBool, Ordering}, Arc};
+use std::sync::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>,
+    pub pull_after: bool,
 }
 
 #[async_trait::async_trait]
@@ -21,30 +21,30 @@ impl Stmt for SyncedStatement {
     }
 
     async fn execute(&self, params: &Params) -> Result<usize> {
-        if self.needs_pull.load(Ordering::Relaxed) {
+        let result = self.inner.execute(params).await;
+        if self.pull_after {
             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
+        result
     }
 
     async fn query(&self, params: &Params) -> Result<Rows> {
-        if self.needs_pull.load(Ordering::Relaxed) {
+        let result = self.inner.query(params).await;
+        if self.pull_after {
             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
+        result
     }
 
     async fn run(&self, params: &Params) -> Result<()> {
-        if self.needs_pull.load(Ordering::Relaxed) {
+        let result = self.inner.run(params).await;
+        if self.pull_after {
             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
+        result
     }
 
     fn interrupt(&self) -> Result<()> {