use either::Either; use libsqlite3_sys::{sqlite3_step, SQLITE_DONE, SQLITE_ROW}; use crate::error::Error; use crate::sqlite::statement::StatementHandle; #[cfg(not(feature = "runtime-tokio"))] use { libsqlite3_sys::sqlite3_stmt, sqlx_rt::yield_now, std::ptr::null_mut, std::sync::atomic::{spin_loop_hint, AtomicI32, AtomicPtr, Ordering}, std::sync::Arc, std::thread::{self, park, spawn, JoinHandle}, }; // For async-std and actix, the worker maintains a dedicated thread for each SQLite connection // All invocations of [sqlite3_step] are run on this thread // For tokio, the worker is a thin wrapper around an invocation to [block_in_place] #[cfg(not(feature = "runtime-tokio"))] const STATE_CLOSE: i32 = -1; #[cfg(not(feature = "runtime-tokio"))] const STATE_READY: i32 = 0; #[cfg(not(feature = "runtime-tokio"))] const STATE_INITIAL: i32 = 1; #[cfg(not(feature = "runtime-tokio"))] pub(crate) struct StatementWorker { statement: Arc>, status: Arc, handle: Option>, } #[cfg(feature = "runtime-tokio")] pub(crate) struct StatementWorker; #[cfg(not(feature = "runtime-tokio"))] impl StatementWorker { pub(crate) fn new() -> Self { let statement = Arc::new(AtomicPtr::new(null_mut::())); let status = Arc::new(AtomicI32::new(STATE_INITIAL)); let handle = spawn({ let statement = Arc::clone(&statement); let status = Arc::clone(&status); move || { // wait for the first command park(); 'run: while status.load(Ordering::Acquire) >= 0 { 'statement: loop { match status.load(Ordering::Acquire) { STATE_CLOSE => { // worker has been dropped; get out break 'run; } STATE_READY => { let statement = statement.load(Ordering::Acquire); if statement.is_null() { // we do not have the statement handle yet thread::yield_now(); continue; } let v = unsafe { sqlite3_step(statement) }; status.store(v, Ordering::Release); if v == SQLITE_DONE { // when a statement is _done_, we park the thread until // we need it again park(); break 'statement; } } _ => { // waits for the receiving end to be ready to receive the rows // this should take less than 1 microsecond under most conditions spin_loop_hint(); } } } } } }); Self { handle: Some(handle), statement, status, } } pub(crate) fn wake(&self) { if let Some(handle) = &self.handle { handle.thread().unpark(); } } pub(crate) fn execute(&self, statement: &StatementHandle) { // readies the worker to execute the statement // for async-std, this unparks our dedicated thread self.statement .store(statement.0.as_ptr(), Ordering::Release); } pub(crate) async fn step(&self, statement: &StatementHandle) -> Result, Error> { // storing <0> as a terminal in status releases the worker // to proceed to the next [sqlite3_step] invocation self.status.store(STATE_READY, Ordering::Release); // we then use a spin loop to wait for this to finish // 99% of the time this should be < 1 μs let status = loop { let status = self .status .compare_and_swap(STATE_READY, STATE_READY, Ordering::AcqRel); if status != STATE_READY { break status; } yield_now().await; }; match status { // a row was found SQLITE_ROW => Ok(Either::Right(())), // reached the end of the query results, // emit the # of changes SQLITE_DONE => Ok(Either::Left(statement.changes())), _ => Err(statement.last_error().into()), } } pub(crate) fn close(&mut self) { self.status.store(STATE_CLOSE, Ordering::Release); if let Some(handle) = self.handle.take() { handle.thread().unpark(); handle.join().unwrap(); } } } #[cfg(feature = "runtime-tokio")] impl StatementWorker { pub(crate) fn new() -> Self { StatementWorker } pub(crate) fn execute(&self, _statement: &StatementHandle) {} pub(crate) fn wake(&self) {} pub(crate) async fn step(&self, statement: &StatementHandle) -> Result, Error> { let statement = *statement; let status = sqlx_rt::blocking!(unsafe { sqlite3_step(statement.0.as_ptr()) }); match status { // a row was found SQLITE_ROW => Ok(Either::Right(())), // reached the end of the query results, // emit the # of changes SQLITE_DONE => Ok(Either::Left(statement.changes())), _ => Err(statement.last_error().into()), } } pub(crate) fn close(&mut self) {} } impl Drop for StatementWorker { fn drop(&mut self) { self.close(); } }