Add query timeout option to interrupt long-running queries

A single background tokio task with a min-heap manages all query
deadlines efficiently. When a query starts, a TimeoutGuard is
acquired; if the deadline expires before the guard is dropped,
the connection is interrupted via sqlite3_interrupt().

Author: penberg

docs/api.md

@@ -22,6 +22,7 @@ You can use the `options` parameter to specify various options. Options supporte
 - `syncPeriod`: synchronize the database periodically every `syncPeriod` seconds.
 - `authToken`: authentication token for the provider URL (optional).
 - `timeout`: number of milliseconds to wait on locked database before returning `SQLITE_BUSY` error
+- `queryTimeout`: maximum number of milliseconds a query is allowed to run before being interrupted with `SQLITE_INTERRUPT` error
 
 The function returns a `Database` object.
 

src/lib.rs

@@ -21,6 +21,7 @@
 #![allow(deprecated)]
 
 mod auth;
+mod query_timeout;
 
 use napi::{
     bindgen_prelude::{Array, FromNapiValue, ToNapiValue},
@@ -200,6 +201,8 @@ pub struct Options {
     pub encryptionKey: Option<String>,
     // Encryption key for remote encryption at rest.
     pub remoteEncryptionKey: Option<String>,
+    // Maximum time in milliseconds that a query is allowed to run.
+    pub queryTimeout: Option<f64>,
 }
 
 /// Access mode.
@@ -224,6 +227,10 @@ pub struct Database {
     default_safe_integers: AtomicBool,
     // Whether to use memory-only mode.
     memory: bool,
+    // Maximum time in milliseconds that a query is allowed to run.
+    query_timeout: Option<Duration>,
+    // Shared timeout manager for efficient query timeout handling.
+    timeout_manager: Arc<query_timeout::QueryTimeoutManager>,
 }
 
 impl Drop for Database {
@@ -320,11 +327,19 @@ pub async fn connect(path: String, opts: Option<Options>) -> Result<Database> {
         conn.busy_timeout(Duration::from_millis(timeout as u64))
             .map_err(Error::from)?
     }
+    let query_timeout = opts
+        .as_ref()
+        .and_then(|o| o.queryTimeout)
+        .filter(|&t| t > 0.0)
+        .map(|t| Duration::from_millis(t as u64));
+    let timeout_manager = Arc::new(query_timeout::QueryTimeoutManager::new());
     Ok(Database {
         db,
         conn: Some(Arc::new(conn)),
         default_safe_integers,
         memory,
+        query_timeout,
+        timeout_manager,
     })
 }
 
@@ -387,7 +402,7 @@ impl Database {
             pluck: false.into(),
             timing: false.into(),
         };
-        Ok(Statement::new(conn, stmt, mode))
+        Ok(Statement::new(conn, stmt, mode, self.query_timeout, self.timeout_manager.clone()))
     }
 
     /// Sets the authorizer for the database.
@@ -515,6 +530,7 @@ impl Database {
                 ));
             }
         };
+        let _guard = self.query_timeout.map(|t| self.timeout_manager.register(&conn, t));
         conn.execute_batch(&sql).await.map_err(Error::from)?;
         Ok(())
     }
@@ -620,6 +636,10 @@ pub struct Statement {
     column_names: Vec<std::ffi::CString>,
     // The access mode.
     mode: AccessMode,
+    // Maximum time in milliseconds that a query is allowed to run.
+    query_timeout: Option<Duration>,
+    // Shared timeout manager.
+    timeout_manager: Arc<query_timeout::QueryTimeoutManager>,
 }
 
 #[napi]
@@ -635,6 +655,8 @@ impl Statement {
         conn: Arc<libsql::Connection>,
         stmt: libsql::Statement,
         mode: AccessMode,
+        query_timeout: Option<Duration>,
+        timeout_manager: Arc<query_timeout::QueryTimeoutManager>,
     ) -> Self {
         let column_names: Vec<std::ffi::CString> = stmt
             .columns()
@@ -647,6 +669,8 @@ impl Statement {
             stmt,
             column_names,
             mode,
+            query_timeout,
+            timeout_manager,
         }
     }
 
@@ -663,8 +687,10 @@ impl Statement {
         let start = std::time::Instant::now();
         let stmt = self.stmt.clone();
         let conn = self.conn.clone();
+        let guard = self.start_timeout_guard();
 
         let future = async move {
+            let _guard = guard;
             stmt.run(params).await.map_err(Error::from)?;
             let changes = if conn.total_changes() == total_changes_before {
                 0
@@ -707,7 +733,9 @@ impl Statement {
         };
 
         let stmt_fut = stmt.clone();
+        let guard = self.start_timeout_guard();
         let future = async move {
+            let _guard = guard;
             let mut rows = stmt_fut.query(params).await.map_err(Error::from)?;
             let row = rows.next().await.map_err(Error::from)?;
             let duration: Option<f64> = start.map(|start| start.elapsed().as_secs_f64());
@@ -771,6 +799,7 @@ impl Statement {
         stmt.reset();
         let params = map_params(&stmt, params).unwrap();
         let stmt = self.stmt.clone();
+        let guard = self.start_timeout_guard();
         let future = async move {
             let rows = stmt.query(params).await.map_err(Error::from)?;
             Ok::<_, napi::Error>(rows)
@@ -783,6 +812,7 @@ impl Statement {
                 safe_ints,
                 raw,
                 pluck,
+                guard,
             ))
         })
     }
@@ -864,6 +894,13 @@ impl Statement {
         self.stmt.interrupt().map_err(Error::from)?;
         Ok(())
     }
+
+}
+
+impl Statement {
+    fn start_timeout_guard(&self) -> Option<query_timeout::TimeoutGuard> {
+        self.query_timeout.map(|t| self.timeout_manager.register(&self.conn, t))
+    }
 }
 
 /// Gets first row from statement in blocking mode.
@@ -885,6 +922,7 @@ pub fn statement_get_sync(
     };
 
     let rt = runtime()?;
+    let _guard = stmt.start_timeout_guard();
     rt.block_on(async move {
         let params = map_params(&stmt.stmt, params)?;
         let mut rows = stmt.stmt.query(params).await.map_err(Error::from)?;
@@ -909,6 +947,7 @@ pub fn statement_get_sync(
 pub fn statement_run_sync(stmt: &Statement, params: Option<napi::JsUnknown>) -> Result<RunResult> {
     stmt.stmt.reset();
     let rt = runtime()?;
+    let _guard = stmt.start_timeout_guard();
     rt.block_on(async move {
         let params = map_params(&stmt.stmt, params)?;
         let total_changes_before = stmt.conn.total_changes();
@@ -940,11 +979,12 @@ pub fn statement_iterate_sync(
     let safe_ints = stmt.mode.safe_ints.load(Ordering::SeqCst);
     let raw = stmt.mode.raw.load(Ordering::SeqCst);
     let pluck = stmt.mode.pluck.load(Ordering::SeqCst);
-    let stmt = stmt.stmt.clone();
+    let guard = stmt.start_timeout_guard();
+    let inner_stmt = stmt.stmt.clone();
     let (rows, column_names) = rt.block_on(async move {
-        stmt.reset();
-        let params = map_params(&stmt, params)?;
-        let rows = stmt.query(params).await.map_err(Error::from)?;
+        inner_stmt.reset();
+        let params = map_params(&inner_stmt, params)?;
+        let rows = inner_stmt.query(params).await.map_err(Error::from)?;
         let mut column_names = Vec::new();
         for i in 0..rows.column_count() {
             column_names
@@ -958,6 +998,7 @@ pub fn statement_iterate_sync(
         safe_ints,
         raw,
         pluck,
+        guard,
     ))
 }
 
@@ -1104,6 +1145,7 @@ pub struct RowsIterator {
     safe_ints: bool,
     raw: bool,
     pluck: bool,
+    _timeout_guard: Option<query_timeout::TimeoutGuard>,
 }
 
 #[napi]
@@ -1114,13 +1156,15 @@ impl RowsIterator {
         safe_ints: bool,
         raw: bool,
         pluck: bool,
+        timeout_guard: Option<query_timeout::TimeoutGuard>,
     ) -> Self {
         Self {
             rows,
             column_names,
             safe_ints,
             raw,
             pluck,
+            _timeout_guard: timeout_guard,
         }
     }
 

src/query_timeout.rs

@@ -0,0 +1,154 @@
+use std::collections::BinaryHeap;
+use std::cmp::Reverse;
+use std::sync::{Arc, Mutex};
+use std::time::Duration;
+use tokio::sync::Notify;
+use tokio::time::Instant;
+
+/// A single-background-task timer wheel that interrupts connections when their
+/// query deadline expires. Registering a query returns a [`TimeoutGuard`] —
+/// dropping the guard cancels the timeout.
+pub struct QueryTimeoutManager {
+    inner: Arc<Inner>,
+}
+
+struct Inner {
+    entries: Mutex<Entries>,
+    /// Wakes the background task when the earliest deadline changes.
+    notify: Notify,
+}
+
+struct Entries {
+    heap: BinaryHeap<Reverse<Entry>>,
+    next_id: u64,
+}
+
+#[derive(Clone)]
+struct Entry {
+    id: u64,
+    deadline: Instant,
+    conn: Arc<libsql::Connection>,
+    /// Cleared when the guard is dropped (query finished in time).
+    active: Arc<AtomicBool>,
+}
+
+use std::sync::atomic::AtomicBool;
+use std::sync::atomic::Ordering;
+
+impl PartialEq for Entry {
+    fn eq(&self, other: &Self) -> bool {
+        self.deadline == other.deadline && self.id == other.id
+    }
+}
+impl Eq for Entry {}
+impl PartialOrd for Entry {
+    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+        Some(self.cmp(other))
+    }
+}
+impl Ord for Entry {
+    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
+        self.deadline.cmp(&other.deadline).then(self.id.cmp(&other.id))
+    }
+}
+
+impl QueryTimeoutManager {
+    pub fn new() -> Self {
+        let inner = Arc::new(Inner {
+            entries: Mutex::new(Entries {
+                heap: BinaryHeap::new(),
+                next_id: 0,
+            }),
+            notify: Notify::new(),
+        });
+        let bg = inner.clone();
+        tokio::spawn(async move {
+            Self::background_task(bg).await;
+        });
+        Self { inner }
+    }
+
+    /// Register a query. The returned guard must be held for the duration of
+    /// the query — dropping it cancels the timeout.
+    pub fn register(&self, conn: &Arc<libsql::Connection>, timeout: Duration) -> TimeoutGuard {
+        let active = Arc::new(AtomicBool::new(true));
+        let mut entries = self.inner.entries.lock().unwrap();
+        let id = entries.next_id;
+        entries.next_id += 1;
+        let entry = Entry {
+            id,
+            deadline: Instant::now() + timeout,
+            conn: conn.clone(),
+            active: active.clone(),
+        };
+        let is_new_earliest = entries
+            .heap
+            .peek()
+            .map_or(true, |Reverse(e)| entry.deadline < e.deadline);
+        entries.heap.push(Reverse(entry));
+        drop(entries);
+        if is_new_earliest {
+            self.inner.notify.notify_one();
+        }
+        TimeoutGuard { active }
+    }
+
+    async fn background_task(inner: Arc<Inner>) {
+        loop {
+            // Find the next deadline, skipping cancelled entries.
+            let next = {
+                let mut entries = inner.entries.lock().unwrap();
+                loop {
+                    match entries.heap.peek() {
+                        Some(Reverse(e)) if !e.active.load(Ordering::Relaxed) => {
+                            entries.heap.pop();
+                        }
+                        Some(Reverse(e)) => break Some(e.clone()),
+                        None => break None,
+                    }
+                }
+            };
+
+            match next {
+                Some(entry) => {
+                    tokio::select! {
+                        _ = tokio::time::sleep_until(entry.deadline) => {
+                            // Deadline reached — interrupt if still active.
+                            if entry.active.load(Ordering::Relaxed) {
+                                let _ = entry.conn.interrupt();
+                            }
+                            // Remove this entry.
+                            let mut entries = inner.entries.lock().unwrap();
+                            // Pop entries that are done (expired or cancelled).
+                            while let Some(Reverse(e)) = entries.heap.peek() {
+                                if !e.active.load(Ordering::Relaxed) || e.id == entry.id {
+                                    entries.heap.pop();
+                                } else {
+                                    break;
+                                }
+                            }
+                        }
+                        _ = inner.notify.notified() => {
+                            // A new earlier deadline was added; re-check.
+                        }
+                    }
+                }
+                None => {
+                    // Nothing to do — wait until a new entry is registered.
+                    inner.notify.notified().await;
+                }
+            }
+        }
+    }
+}
+
+/// Dropping this guard cancels the associated query timeout.
+pub struct TimeoutGuard {
+    active: Arc<AtomicBool>,
+}
+
+impl Drop for TimeoutGuard {
+    fn drop(&mut self) {
+        self.active.store(false, Ordering::Relaxed);
+    }
+}