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.
 

integration-tests/tests/async.test.js

@@ -393,6 +393,35 @@ test.serial("Timeout option", async (t) => {
   fs.unlinkSync(path);
 });
 
+test.serial("Query timeout option interrupts long-running query", async (t) => {
+  const queryTimeout = 100;
+  const path = genDatabaseFilename();
+  const [db, errorType] = await connect(path, { queryTimeout });
+  const stmt = await db.prepare(
+    "WITH RECURSIVE infinite_loop(n) AS (SELECT 1 UNION ALL SELECT n + 1 FROM infinite_loop) SELECT * FROM infinite_loop;"
+  );
+
+  await t.throwsAsync(async () => {
+    await stmt.all();
+  }, {
+    instanceOf: errorType,
+    message: "interrupted",
+    code: "SQLITE_INTERRUPT",
+  });
+
+  db.close();
+  fs.unlinkSync(path);
+});
+
+test.serial("Query timeout option allows short-running query", async (t) => {
+  const path = genDatabaseFilename();
+  const [db] = await connect(path, { queryTimeout: 100 });
+  const stmt = await db.prepare("SELECT 1 AS value");
+  t.deepEqual(await stmt.get(), { value: 1 });
+  db.close();
+  fs.unlinkSync(path);
+});
+
 test.serial("Concurrent writes over same connection", async (t) => {
   const db = t.context.db;
   await db.exec(`

integration-tests/tests/sync.test.js

@@ -457,6 +457,44 @@ test.serial("Timeout option", async (t) => {
   fs.unlinkSync(path);
 });
 
+test.serial("Query timeout option interrupts long-running query", async (t) => {
+  if (t.context.provider === "sqlite") {
+    t.assert(true);
+    return;
+  }
+
+  const path = genDatabaseFilename();
+  const [db, errorType] = await connect(path, { queryTimeout: 100 });
+  const stmt = db.prepare(
+    "WITH RECURSIVE infinite_loop(n) AS (SELECT 1 UNION ALL SELECT n + 1 FROM infinite_loop) SELECT * FROM infinite_loop;"
+  );
+
+  t.throws(() => {
+    stmt.all();
+  }, {
+    instanceOf: errorType,
+    message: "interrupted",
+    code: "SQLITE_INTERRUPT",
+  });
+
+  db.close();
+  fs.unlinkSync(path);
+});
+
+test.serial("Query timeout option allows short-running query", async (t) => {
+  if (t.context.provider === "sqlite") {
+    t.assert(true);
+    return;
+  }
+
+  const path = genDatabaseFilename();
+  const [db] = await connect(path, { queryTimeout: 100 });
+  const stmt = db.prepare("SELECT 1 AS value");
+  t.deepEqual(stmt.get(), { value: 1 });
+  db.close();
+  fs.unlinkSync(path);
+});
+
 test.serial("Statement.reader [SELECT is true]", async (t) => {
   const db = t.context.db;
   const stmt = db.prepare("SELECT * FROM users WHERE id = ?");

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,
         }
     }