diff --git a/Cargo.toml b/Cargo.toml index fb1385a..9b3598d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -11,7 +11,7 @@ crate-type = ["cdylib"] [dependencies] glob-match = "0.2" -libsql = { version = "0.10.0-pre.3", features = ["encryption"] } +libsql = { version = "0.10.0-pre.4", features = ["encryption"] } napi = { version = "2", default-features = false, features = ["napi6", "tokio_rt", "async"] } napi-derive = "2" once_cell = "1.18.0" diff --git a/integration-tests/tests/async.test.js b/integration-tests/tests/async.test.js index 03644e8..a7ac3c2 100644 --- a/integration-tests/tests/async.test.js +++ b/integration-tests/tests/async.test.js @@ -548,32 +548,35 @@ test.serial("Timeout on Statement.get() does not leak into later prepare()/EXPLA db.close(); }); -test.serial("Query timeout covers wait time on the execution lock", async (t) => { - t.timeout(15_000); - const [db, errorType] = await connect(":memory:", { defaultQueryTimeout: 200 }); - await db.exec("CREATE TABLE t(x INTEGER)"); - const insert = await db.prepare("INSERT INTO t VALUES (?)"); - for (let i = 0; i < 5000; i++) { - await insert.run(i); - } - const stmt = await db.prepare("SELECT * FROM t ORDER BY x ASC"); +test.serial("Query timeout interrupts long-running queries under concurrency", async (t) => { + t.timeout(30_000); + const [db] = await connect(":memory:"); + + // Run many never-ending queries concurrently, each on its own statement so + // they genuinely run in parallel. Every one must be interrupted by its own + // query timeout, leaving no query running. + const CONCURRENCY = 20; + const sql = + "WITH RECURSIVE inf(n) AS (SELECT 1 UNION ALL SELECT n + 1 FROM inf) SELECT * FROM inf"; let interrupts = 0; const promises = []; - for (let i = 0; i < 50; i++) { + for (let i = 0; i < CONCURRENCY; i++) { promises.push( - stmt.all().catch((err) => { - if (err.code === "SQLITE_INTERRUPT") { - interrupts++; - } else { - throw err; - } - }) + db.prepare(sql).then((stmt) => + stmt.all(undefined, { queryTimeout: 100 }).catch((err) => { + if (err.code === "SQLITE_INTERRUPT") { + interrupts++; + } else { + throw err; + } + }) + ) ); } await Promise.all(promises); - t.true(interrupts > 0, `expected some queries to timeout, got ${interrupts}`); + t.is(interrupts, CONCURRENCY, `expected every query to time out, got ${interrupts}`); db.close(); }); diff --git a/src/lib.rs b/src/lib.rs index 2977870..16b0e89 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -37,7 +37,7 @@ use std::{ atomic::{AtomicBool, Ordering}, Arc, Mutex, }, - time::{Duration, Instant}, + time::Duration, }; use tokio::runtime::Runtime; use tracing_subscriber::{filter::LevelFilter, EnvFilter}; @@ -238,8 +238,6 @@ pub struct Database { memory: bool, // Maximum time in milliseconds that a query is allowed to run. query_timeout: Option, - // Ensures only one operation executes per connection at a time. - execution_lock: Arc>, } impl Drop for Database { @@ -341,14 +339,12 @@ pub async fn connect(path: String, opts: Option) -> Result { .as_ref() .and_then(|o| o.defaultQueryTimeout) .and_then(query_timeout_duration); - let execution_lock = Arc::new(tokio::sync::Mutex::new(())); Ok(Database { db: Some(db), conn: Some(conn), default_safe_integers, memory, query_timeout, - execution_lock, }) } @@ -404,9 +400,7 @@ impl Database { )); } }; - let (_execution_guard, deadline) = - acquire_execution_lock(&self.execution_lock, self.query_timeout).await?; - let timeout_guard = register_remaining_timeout(&conn, deadline)?; + let timeout_guard = register_timeout(&conn, self.query_timeout); let stmt = match conn.prepare(&sql).await { Ok(stmt) => stmt, Err(err) if is_sqlite_interrupt(&err) => { @@ -414,7 +408,7 @@ impl Database { // can't fire conn.interrupt() for our id mid-probe. drop(timeout_guard); clear_stale_interrupt(&conn).await; - let _retry_guard = register_remaining_timeout(&conn, deadline)?; + let _retry_guard = register_timeout(&conn, self.query_timeout); conn.prepare(&sql).await.map_err(Error::from)? } Err(err) => return Err(Error::from(err).into()), @@ -425,13 +419,7 @@ impl Database { pluck: false.into(), timing: false.into(), }; - Ok(Statement::new( - conn, - stmt, - mode, - self.query_timeout, - self.execution_lock.clone(), - )) + Ok(Statement::new(conn, stmt, mode, self.query_timeout)) } /// Sets the authorizer for the database. @@ -573,9 +561,7 @@ impl Database { Some(timeout_ms) => query_timeout_duration(timeout_ms), None => self.query_timeout, }; - let (_execution_guard, deadline) = - acquire_execution_lock(&self.execution_lock, query_timeout).await?; - let _timeout_guard = register_remaining_timeout(&conn, deadline)?; + let _timeout_guard = register_timeout(&conn, query_timeout); conn.execute_batch(&sql).await.map_err(Error::from)?; Ok(()) } @@ -865,58 +851,18 @@ fn is_sqlite_interrupt(err: &libsql::Error) -> bool { ) } -fn query_timeout_error() -> napi::Error { - throw_sqlite_error( - "interrupted".to_string(), - "SQLITE_INTERRUPT".to_string(), - libsql::ffi::SQLITE_INTERRUPT, - ) -} - -fn timeout_deadline(timeout: Duration) -> Instant { - let now = Instant::now(); - now.checked_add(timeout) - .unwrap_or_else(|| now + Duration::from_secs(86400)) -} - -fn register_remaining_timeout( - conn: &Arc, - deadline: Option, -) -> Result> { - match deadline { - Some(deadline) => match deadline.checked_duration_since(Instant::now()) { - Some(remaining) if remaining > Duration::ZERO => { - Ok(Some(QueryTimeoutManager::global().register(conn, remaining))) - } - _ => Err(query_timeout_error()), - }, - None => Ok(None), - } -} -/// Acquire the per-connection execution lock, enforcing the query timeout -/// across the queue-wait. Returns the lock guard and the absolute deadline -/// for the current operation. If the timeout elapses while waiting, returns -/// a SQLITE_INTERRUPT error so the caller sees the same behaviour as when a -/// timeout fires during execution. -async fn acquire_execution_lock( - lock: &Arc>, +/// Register a query timeout against `target` for the operation about to run, +/// returning a guard that cancels the timeout when dropped. The target is the +/// statement for per-statement operations (so the timer interrupts only that +/// statement, leaving concurrent operations on the same connection untouched) +/// or the connection for connection-wide operations. Returns `None` when no +/// timeout is in effect. +fn register_timeout( + target: &Arc, timeout: Option, -) -> Result<(tokio::sync::OwnedMutexGuard<()>, Option)> { - let lock = lock.clone(); - match timeout { - None => Ok((lock.lock_owned().await, None)), - Some(t) => { - let deadline = timeout_deadline(t); - match tokio::time::timeout(t, lock.lock_owned()).await { - Ok(guard) => match deadline.checked_duration_since(Instant::now()) { - Some(remaining) if remaining > Duration::ZERO => Ok((guard, Some(deadline))), - _ => Err(query_timeout_error()), - }, - Err(_) => Err(query_timeout_error()), - } - } - } +) -> Option { + timeout.map(|t| QueryTimeoutManager::global().register(target, t)) } async fn clear_stale_interrupt(conn: &Arc) { @@ -944,8 +890,6 @@ pub struct Statement { mode: AccessMode, // Maximum time in milliseconds that a query is allowed to run. query_timeout: Option, - // Shared per-connection execution lock. - execution_lock: Arc>, } #[napi] @@ -962,7 +906,6 @@ impl Statement { stmt: libsql::Statement, mode: AccessMode, query_timeout: Option, - execution_lock: Arc>, ) -> Self { let column_names: Vec = stmt .columns() @@ -976,7 +919,6 @@ impl Statement { column_names, mode, query_timeout, - execution_lock, } } @@ -999,12 +941,9 @@ impl Statement { let stmt = self.stmt.clone(); let conn = self.conn.clone(); let query_timeout = self.resolve_query_timeout(query_options); - let execution_lock = self.execution_lock.clone(); let future = async move { - let (_execution_guard, deadline) = - acquire_execution_lock(&execution_lock, query_timeout).await?; - let _timeout_guard = register_remaining_timeout(&conn, deadline)?; + let _timeout_guard = register_timeout(&stmt, query_timeout); stmt.run(params).await.map_err(Error::from)?; let changes = if conn.total_changes() == total_changes_before { 0 @@ -1052,14 +991,10 @@ impl Statement { let stmt = self.stmt.clone(); let stmt_fut = stmt.clone(); - let conn = self.conn.clone(); let query_timeout = self.resolve_query_timeout(query_options); - let execution_lock = self.execution_lock.clone(); let future = async move { - let (_execution_guard, deadline) = - acquire_execution_lock(&execution_lock, query_timeout).await?; let result: std::result::Result<(Option, Option), Error> = { - let _timeout_guard = register_remaining_timeout(&conn, deadline)?; + let _timeout_guard = register_timeout(&stmt_fut, query_timeout); async { let mut rows = stmt_fut.query(params).await.map_err(Error::from)?; let row = rows.next().await.map_err(Error::from)?; @@ -1137,32 +1072,24 @@ impl Statement { let params = map_params(&stmt, params)?; let stmt_for_query = self.stmt.clone(); let stmt_for_iter = stmt_for_query.clone(); - let conn = self.conn.clone(); let query_timeout = self.resolve_query_timeout(query_options); - let execution_lock = self.execution_lock.clone(); let future = async move { - let (execution_guard, deadline) = - acquire_execution_lock(&execution_lock, query_timeout).await?; - let timeout_guard = register_remaining_timeout(&conn, deadline)?; + let timeout_guard = register_timeout(&stmt_for_query, query_timeout); let rows = stmt_for_query.query(params).await.map_err(Error::from)?; - Ok::<_, napi::Error>((rows, execution_guard, timeout_guard)) + Ok::<_, napi::Error>((rows, timeout_guard)) }; let column_names = self.column_names.clone(); - env.execute_tokio_future( - future, - move |&mut _env, (result, execution_guard, timeout_guard)| { - Ok(RowsIterator::new( - Arc::new(tokio::sync::Mutex::new(result)), - stmt_for_iter, - column_names, - safe_ints, - raw, - pluck, - timeout_guard, - execution_guard, - )) - }, - ) + env.execute_tokio_future(future, move |&mut _env, (result, timeout_guard)| { + Ok(RowsIterator::new( + Arc::new(tokio::sync::Mutex::new(result)), + stmt_for_iter, + column_names, + safe_ints, + raw, + pluck, + timeout_guard, + )) + }) } #[napi] @@ -1274,12 +1201,9 @@ pub fn statement_get_sync( let rt = runtime()?; let query_timeout = stmt.resolve_query_timeout(query_options); - let execution_lock = stmt.execution_lock.clone(); let result: Result<(Option, Option)> = { rt.block_on(async move { - let (_execution_guard, deadline) = - acquire_execution_lock(&execution_lock, query_timeout).await?; - let _timeout_guard = register_remaining_timeout(&stmt.conn, deadline)?; + let _timeout_guard = register_timeout(&stmt.stmt, query_timeout); let params = map_params(&stmt.stmt, params)?; let mut rows = stmt.stmt.query(params).await.map_err(Error::from)?; let row = rows.next().await.map_err(Error::from)?; @@ -1318,11 +1242,8 @@ pub fn statement_run_sync( stmt.stmt.reset(); let rt = runtime()?; let query_timeout = stmt.resolve_query_timeout(query_options); - let execution_lock = stmt.execution_lock.clone(); rt.block_on(async move { - let (_execution_guard, deadline) = - acquire_execution_lock(&execution_lock, query_timeout).await?; - let _timeout_guard = register_remaining_timeout(&stmt.conn, deadline)?; + let _timeout_guard = register_timeout(&stmt.stmt, query_timeout); let params = map_params(&stmt.stmt, params)?; let total_changes_before = stmt.conn.total_changes(); let start = std::time::Instant::now(); @@ -1355,14 +1276,10 @@ pub fn statement_iterate_sync( let raw = stmt.mode.raw.load(Ordering::SeqCst); let pluck = stmt.mode.pluck.load(Ordering::SeqCst); let query_timeout = stmt.resolve_query_timeout(query_options); - let conn = stmt.conn.clone(); - let execution_lock = stmt.execution_lock.clone(); let inner_stmt = stmt.stmt.clone(); let iter_stmt = inner_stmt.clone(); - let (rows, column_names, execution_guard, timeout_guard) = rt.block_on(async move { - let (execution_guard, deadline) = - acquire_execution_lock(&execution_lock, query_timeout).await?; - let timeout_guard = register_remaining_timeout(&conn, deadline)?; + let (rows, column_names, timeout_guard) = rt.block_on(async move { + let timeout_guard = register_timeout(&inner_stmt, query_timeout); inner_stmt.reset(); let params = map_params(&inner_stmt, params)?; let rows = inner_stmt.query(params).await.map_err(Error::from)?; @@ -1371,7 +1288,7 @@ pub fn statement_iterate_sync( column_names .push(std::ffi::CString::new(rows.column_name(i).unwrap().to_string()).unwrap()); } - Ok::<_, napi::Error>((rows, column_names, execution_guard, timeout_guard)) + Ok::<_, napi::Error>((rows, column_names, timeout_guard)) })?; Ok(RowsIterator::new( Arc::new(tokio::sync::Mutex::new(rows)), @@ -1381,7 +1298,6 @@ pub fn statement_iterate_sync( raw, pluck, timeout_guard, - execution_guard, )) } @@ -1530,7 +1446,6 @@ pub struct RowsIterator { raw: bool, pluck: bool, timeout_guard: Mutex>, - execution_guard: Mutex>>, } #[napi] @@ -1543,7 +1458,6 @@ impl RowsIterator { raw: bool, pluck: bool, timeout_guard: Option, - execution_guard: tokio::sync::OwnedMutexGuard<()>, ) -> Self { Self { rows, @@ -1553,7 +1467,6 @@ impl RowsIterator { raw, pluck, timeout_guard: Mutex::new(timeout_guard), - execution_guard: Mutex::new(Some(execution_guard)), } } @@ -1588,8 +1501,6 @@ impl RowsIterator { self.stmt.reset(); let mut timeout_guard = self.timeout_guard.lock().unwrap(); timeout_guard.take(); - let mut execution_guard = self.execution_guard.lock().unwrap(); - execution_guard.take(); } } diff --git a/src/query_timeout.rs b/src/query_timeout.rs index 01888ca..2b9cd55 100644 --- a/src/query_timeout.rs +++ b/src/query_timeout.rs @@ -5,11 +5,33 @@ use std::{ time::{Duration, Instant}, }; +/// Something a timeout can interrupt when its deadline expires. Implemented for +/// both libSQL connections and statements so a single wheel can interrupt at +/// whichever granularity the operation registered with: statement-level for the +/// `Statement` operations (so concurrent operations on the same connection are +/// left untouched) and connection-level for connection-wide operations such as +/// `prepare` and `execute_batch`. +pub trait Interruptible: Send + Sync { + fn interrupt(&self); +} + +impl Interruptible for libsql::Connection { + fn interrupt(&self) { + let _ = libsql::Connection::interrupt(self); + } +} + +impl Interruptible for libsql::Statement { + fn interrupt(&self) { + let _ = libsql::Statement::interrupt(self); + } +} + /// A process-wide timer wheel: a single background thread that interrupts -/// connections when their currently registered operation exceeds its deadline. +/// operations when they exceed their deadline. /// /// All connections share one manager — and therefore one thread. Each -/// registered timeout carries a weak reference to the connection it should +/// registered timeout carries a weak reference to the target it should /// interrupt, so a single wheel serves any number of connections. pub struct QueryTimeoutManager { inner: Arc, @@ -27,7 +49,7 @@ struct State { /// interrupted if it is still present here; a completed operation removes /// itself via its guard's `Drop`, leaving a stale heap entry that is /// discarded lazily when its deadline is reached. - active: HashMap>, + active: HashMap>, next_id: u64, shutdown: bool, } @@ -102,12 +124,12 @@ impl QueryTimeoutManager { self.inner.cv.notify_one(); } - /// Register a timeout for an operation about to run on `conn`. The returned - /// guard cancels the timeout when dropped (i.e. when the operation + /// Register a timeout for an operation about to run against `target`. The + /// returned guard cancels the timeout when dropped (i.e. when the operation /// completes). - pub fn register( + pub fn register( &self, - conn: &Arc, + target: &Arc, timeout: Duration, ) -> QueryTimeoutGuard { let mut state = self.inner.state.lock().unwrap(); @@ -125,7 +147,9 @@ impl QueryTimeoutManager { .peek() .map_or(true, |Reverse(existing)| entry.deadline < existing.deadline); - state.active.insert(id, Arc::downgrade(conn)); + let weak = Arc::downgrade(target); + let target: Weak = weak; + state.active.insert(id, target); state.heap.push(Reverse(entry)); drop(state); @@ -164,9 +188,9 @@ impl QueryTimeoutManager { // Interrupt only if the operation is still in flight; a completed // operation will have removed itself from `active` already. - if let Some(conn) = state.active.remove(&entry.id) { - if let Some(conn) = conn.upgrade() { - let _ = conn.interrupt(); + if let Some(target) = state.active.remove(&entry.id) { + if let Some(target) = target.upgrade() { + target.interrupt(); } } } @@ -263,6 +287,93 @@ mod tests { assert!(result.is_err(), "query should have been interrupted"); } + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + #[ntest::timeout(10000)] + async fn deadline_expires_interrupts_statement() { + let conn = test_conn().await; + let stmt = Arc::new( + conn.prepare( + "WITH RECURSIVE r(n) AS (SELECT 1 UNION ALL SELECT n+1 FROM r) SELECT * FROM r", + ) + .await + .unwrap(), + ); + let mgr = QueryTimeoutManager::new(); + + let _guard = mgr.register(&stmt, Duration::from_millis(200)); + + // Drain the (effectively infinite) result set; the statement-level + // interrupt must abort it once the deadline expires. + let result = async { + let mut rows = stmt.query(()).await?; + while rows.next().await?.is_some() {} + Ok::<_, libsql::Error>(()) + } + .await; + + assert!(result.is_err(), "statement should have been interrupted"); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + #[ntest::timeout(10000)] + async fn deadline_expires_interrupts_single_step_aggregate() { + // A query whose entire work happens inside the first `step()` (an + // aggregate over a large recursive CTE). libSQL runs that first step + // synchronously inside `query()`, so this exercises whether a + // statement-level interrupt aborts an in-progress step. + let conn = test_conn().await; + let stmt = Arc::new( + conn.prepare( + "WITH RECURSIVE numbers(value) AS ( + SELECT 1 UNION ALL SELECT value + 1 FROM numbers WHERE value < 1000000000 + ) SELECT sum(value) FROM numbers", + ) + .await + .unwrap(), + ); + let mgr = QueryTimeoutManager::new(); + + let _guard = mgr.register(&stmt, Duration::from_millis(200)); + + let result = async { + let mut rows = stmt.query(()).await?; + while rows.next().await?.is_some() {} + Ok::<_, libsql::Error>(()) + } + .await; + + assert!(result.is_err(), "statement should have been interrupted"); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + #[ntest::timeout(10000)] + async fn statement_interrupt_flag_is_sticky_until_reset() { + // `libsql_stmt_interrupt` sets a per-statement flag that `sqlite3_step` + // checks at entry and does NOT clear; only `reset()` clears it. So a + // statement interrupted by a timeout stays poisoned until reset — every + // operation must reset before stepping it again. + let conn = test_conn().await; + let stmt = Arc::new(conn.prepare("SELECT 1").await.unwrap()); + + stmt.interrupt().unwrap(); + + // The flag survives into the next execution. + let mut rows = stmt.query(()).await.unwrap(); + assert!( + rows.next().await.is_err(), + "sticky interrupt flag should fail the next step" + ); + drop(rows); + + // ...and reset clears it, making the statement reusable. + stmt.reset(); + let mut rows = stmt.query(()).await.unwrap(); + assert!( + rows.next().await.unwrap().is_some(), + "statement should be reusable after reset" + ); + } + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[ntest::timeout(10000)] async fn guard_dropped_before_deadline_cancels_timeout() {