From e73cc8252a14753aef2aa5638996eaf8013e913e Mon Sep 17 00:00:00 2001 From: "Levy A." Date: Tue, 24 Jun 2025 14:46:32 -0300 Subject: [PATCH 1/2] fix: pull on read --- libsql/src/database.rs | 1 + libsql/src/sync.rs | 1 - libsql/src/sync/connection.rs | 34 +++++++++++------- libsql/src/sync/statement.rs | 66 ----------------------------------- 4 files changed, 23 insertions(+), 79 deletions(-) delete mode 100644 libsql/src/sync/statement.rs diff --git a/libsql/src/database.rs b/libsql/src/database.rs index bb07bb189d..7069799caa 100644 --- a/libsql/src/database.rs +++ b/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); diff --git a/libsql/src/sync.rs b/libsql/src/sync.rs index 9a74b118b0..d7fdf5c48f 100644 --- a/libsql/src/sync.rs +++ b/libsql/src/sync.rs @@ -13,7 +13,6 @@ use uuid::Uuid; mod test; pub mod connection; -pub mod statement; pub mod transaction; const METADATA_VERSION: u32 = 0; diff --git a/libsql/src/sync/connection.rs b/libsql/src/sync/connection.rs index 807e49a2f4..73019d03be 100644 --- a/libsql/src/sync/connection.rs +++ b/libsql/src/sync/connection.rs @@ -8,11 +8,14 @@ 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; -use super::{statement::SyncedStatement, transaction::SyncedTx}; +use super::transaction::SyncedTx; #[derive(Clone)] pub struct SyncedConnection { @@ -21,6 +24,7 @@ pub struct SyncedConnection { pub read_your_writes: bool, pub context: Arc>, pub state: Arc>, + pub needs_pull: Arc, } impl SyncedConnection { @@ -89,7 +93,7 @@ impl SyncedConnection { _ => { *state = predicted_end_state; false - }, + } }; Ok(should_execute_local) @@ -106,6 +110,10 @@ impl Conn for SyncedConnection { async fn execute_batch(&self, sql: &str) -> Result { 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?; + } self.local.execute_batch(sql) } else { self.remote.execute_batch(sql).await @@ -114,6 +122,10 @@ impl Conn for SyncedConnection { async fn execute_transactional_batch(&self, sql: &str) -> Result { 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?; + } self.local.execute_transactional_batch(sql)?; Ok(BatchRows::empty()) } else { @@ -123,6 +135,10 @@ impl Conn for SyncedConnection { async fn prepare(&self, sql: &str) -> Result { 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 { inner: Box::new(LibsqlStmt(self.local.prepare(sql)?)), }) @@ -132,16 +148,10 @@ impl Conn for SyncedConnection { }; 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) } } diff --git a/libsql/src/sync/statement.rs b/libsql/src/sync/statement.rs deleted file mode 100644 index 679933298b..0000000000 --- a/libsql/src/sync/statement.rs +++ /dev/null @@ -1,66 +0,0 @@ -use crate::{ - local::{self}, - params::Params, - statement::Stmt, - sync::SyncContext, Column, Result, Rows, Statement, -}; -use std::sync::Arc; -use tokio::sync::Mutex; - -pub struct SyncedStatement { - pub conn: local::Connection, - pub context: Arc>, - pub inner: Statement, -} - -#[async_trait::async_trait] -impl Stmt for SyncedStatement { - fn finalize(&mut self) { - self.inner.finalize() - } - - async fn execute(&mut self, params: &Params) -> Result { - let result = self.inner.execute(params).await; - let mut context = self.context.lock().await; - crate::sync::try_pull(&mut context, &self.conn).await?; - result - } - - async fn query(&mut self, params: &Params) -> Result { - let result = self.inner.query(params).await; - let mut context = self.context.lock().await; - crate::sync::try_pull(&mut context, &self.conn).await?; - result - } - - 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 - } - - 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 { - self.inner.columns() - } -} From f7e1c86132288c43abc1f33bf6b9bb9b8a9d4b4d Mon Sep 17 00:00:00 2001 From: "Levy A." Date: Thu, 26 Jun 2025 16:50:03 -0300 Subject: [PATCH 2/2] improve lazy pull --- libsql/src/sync.rs | 1 + libsql/src/sync/connection.rs | 23 +++++++---- libsql/src/sync/statement.rs | 74 +++++++++++++++++++++++++++++++++++ 3 files changed, 90 insertions(+), 8 deletions(-) create mode 100644 libsql/src/sync/statement.rs diff --git a/libsql/src/sync.rs b/libsql/src/sync.rs index d7fdf5c48f..9a74b118b0 100644 --- a/libsql/src/sync.rs +++ b/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; diff --git a/libsql/src/sync/connection.rs b/libsql/src/sync/connection.rs index 73019d03be..c2809c0f57 100644 --- a/libsql/src/sync/connection.rs +++ b/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 { 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 { 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 { 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 { diff --git a/libsql/src/sync/statement.rs b/libsql/src/sync/statement.rs new file mode 100644 index 0000000000..ad2183b0f4 --- /dev/null +++ b/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>, + pub needs_pull: Arc, +} + +#[async_trait::async_trait] +impl Stmt for SyncedStatement { + fn finalize(&mut self) { + self.inner.finalize() + } + + async fn execute(&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.execute(params).await + } + + async fn query(&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.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 { + self.inner.columns() + } +} +