Pull lazily on read for read-your-writes (#2109)

#2105

Before:
```cmd
$ hyperfine target/debug/libsql-probe
Benchmark 1: target/debug/libsql-probe
  Time (mean ± σ):     11.488 s ±  3.327 s    [User: 0.367 s, System: 0.111 s]
  Range (min … max):    6.270 s … 17.406 s    10 runs
```

After:
```cmd
$ hyperfine target/debug/libsql-probe
Benchmark 1: target/debug/libsql-probe
  Time (mean ± σ):      3.156 s ±  1.063 s    [User: 0.406 s, System: 0.066 s]
  Range (min … max):    1.764 s …  5.268 s    10 runs
```

On this kind of workload, running the server locally:

```rs
for i in 0..100 {
    conn.execute("insert into test values (?)", libsql::params!(i)).await?;
}

let mut rows = conn.query("select * from test", ()).await?;
```

Author: levydsa

libsql/src/database.rs

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

libsql/src/sync/statement.rs

@@ -4,13 +4,14 @@ use crate::{
     statement::Stmt,
     sync::SyncContext, Column, Result, Rows, Statement,
 };
-use std::sync::Arc;
+use std::sync::{atomic::{AtomicBool, Ordering}, Arc};
 use tokio::sync::Mutex;
 
 pub struct SyncedStatement {
     pub conn: local::Connection,
-    pub context: Arc<Mutex<SyncContext>>,
     pub inner: Statement,
+    pub context: Arc<Mutex<SyncContext>>,
+    pub needs_pull: Arc<AtomicBool>,
 }
 
 #[async_trait::async_trait]
@@ -20,24 +21,30 @@ impl Stmt for SyncedStatement {
     }
 
     async fn execute(&mut self, params: &Params) -> Result<usize> {
-        let result = self.inner.execute(params).await;
-        let mut context = self.context.lock().await;
-        crate::sync::try_pull(&mut context, &self.conn).await?;
-        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.execute(params).await
     }
 
     async fn query(&mut self, params: &Params) -> Result<Rows> {
-        let result = self.inner.query(params).await;
-        let mut context = self.context.lock().await;
-        crate::sync::try_pull(&mut context, &self.conn).await?;
-        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.query(params).await
     }
 
     async fn run(&mut self, params: &Params) -> Result<()> {
-        let result = self.inner.run(params).await;
-        let mut context = self.context.lock().await;
-        crate::sync::try_pull(&mut context, &self.conn).await?;
-        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<()> {
@@ -64,3 +71,4 @@ impl Stmt for SyncedStatement {
         self.inner.columns()
     }
 }
+