Pass reserver_ids to SQLite to pre-filter the queried chunks

opsqueue/Cargo.toml

@@ -33,6 +33,7 @@ ux = "0.1.6"
 anyhow = "1.0.102"
 # Database:
 sqlx = { version = "0.8.2", features = ["sqlite", "runtime-tokio", "chrono"], optional = true }
+libsqlite3-sys = { version = "0.30.1", optional = true }
 # Serialization:
 serde = { version = "1.0.203", features = ["derive"] }
 serde_json = "1.0.149"
@@ -96,6 +97,7 @@ assert_matches = { version = "1.5.0" }
 # Dependencies only in use by the server-logic:
 server-logic = [
     "dep:sqlx",
+    "dep:libsqlite3-sys",
     "dep:opentelemetry-otlp",
     "dep:opentelemetry-semantic-conventions",
     "dep:moka",

opsqueue/src/common/chunk.rs

@@ -124,6 +124,18 @@ impl From<ChunkIndex> for i64 {
     }
 }
 
+impl TryFrom<i64> for ChunkIndex {
+    type Error = crate::common::errors::TryFromIntError;
+
+    fn try_from(value: i64) -> Result<Self, Self::Error> {
+        if value < 0 {
+            return Err(crate::common::errors::TryFromIntError(()));
+        }
+
+        Ok(Self(u63::new(value as u64)))
+    }
+}
+
 impl TryFrom<u64> for ChunkIndex {
     type Error = crate::common::errors::TryFromIntError;
     fn try_from(value: u64) -> Result<Self, Self::Error> {

opsqueue/src/common/submission.rs

@@ -532,12 +532,10 @@ pub mod db {
             r#"
             SELECT id AS "id: SubmissionId" FROM submissions WHERE prefix = $1
             UNION ALL
-            SELECT id AS "id: SubmissionId" FROM submissions_completed WHERE prefix = $2
+            SELECT id AS "id: SubmissionId" FROM submissions_completed WHERE prefix = $1
             UNION ALL
-            SELECT id AS "id: SubmissionId" FROM submissions_failed WHERE prefix = $3
+            SELECT id AS "id: SubmissionId" FROM submissions_failed WHERE prefix = $1
             "#,
-            prefix,
-            prefix,
             prefix
         )
         .fetch_optional(conn.get_inner())

opsqueue/src/consumer/dispatcher/mod.rs

@@ -3,12 +3,13 @@ pub mod reserver;
 
 use crate::{
     common::{
-        chunk::{Chunk, ChunkId},
-        submission::Submission,
+        chunk::{Chunk, ChunkId, ChunkIndex},
+        submission::{Submission, SubmissionId},
     },
     db::{magic::Bool, Connection, Pool, ReaderPool},
 };
 use futures::stream::{StreamExt as _, TryStreamExt as _};
+use libsqlite3_sys as ffi;
 use metastate::MetaState;
 use reserver::Reserver;
 use sqlx::QueryBuilder;
@@ -21,6 +22,63 @@ use std::sync::Arc;
 use super::strategy;
 use crate::common::StrategicMetadataMap;
 
+unsafe extern "C" fn sqlite_reserved_chunk_lookup(
+    context: *mut ffi::sqlite3_context,
+    n_args: i32,
+    args: *mut *mut ffi::sqlite3_value,
+) {
+    if n_args != 2 {
+        tracing::error!(
+            n_args,
+            "opsqueue_is_reserved called with unexpected argument count"
+        );
+        // Fail open: this callback is an optimization only.
+        ffi::sqlite3_result_int(context, 0);
+        return;
+    }
+
+    let user_data = ffi::sqlite3_user_data(context) as *const Reserver<ChunkId, ChunkId>;
+    if user_data.is_null() {
+        tracing::error!("opsqueue_is_reserved called without registered reserver user_data");
+        // Fail open: this callback is an optimization only.
+        ffi::sqlite3_result_int(context, 0);
+        return;
+    }
+
+    let submission_id_raw = ffi::sqlite3_value_int64(*args.add(0));
+    let chunk_index_raw = ffi::sqlite3_value_int64(*args.add(1));
+
+    let Ok(submission_id) = SubmissionId::try_from(submission_id_raw) else {
+        tracing::error!(
+            submission_id_raw,
+            "opsqueue_is_reserved got invalid submission_id"
+        );
+        // Fail open: this callback is an optimization only.
+        ffi::sqlite3_result_int(context, 0);
+        return;
+    };
+    let Ok(chunk_index) = ChunkIndex::try_from(chunk_index_raw) else {
+        tracing::error!(
+            chunk_index_raw,
+            "opsqueue_is_reserved got invalid chunk_index"
+        );
+        // Fail open: this callback is an optimization only.
+        ffi::sqlite3_result_int(context, 0);
+        return;
+    };
+
+    let chunk_id = ChunkId::from((submission_id, chunk_index));
+    let is_reserved = (*user_data).is_reserved(&chunk_id);
+    ffi::sqlite3_result_int(context, i32::from(is_reserved));
+}
+
+unsafe extern "C" fn sqlite_reserved_chunk_lookup_destructor(ptr: *mut std::ffi::c_void) {
+    if ptr.is_null() {
+        return;
+    }
+    let _boxed: Box<Reserver<ChunkId, ChunkId>> = Box::from_raw(ptr.cast());
+}
+
 #[derive(Debug, Clone)]
 pub struct Dispatcher {
     reserver: Reserver<ChunkId, ChunkId>,
@@ -66,6 +124,8 @@ impl Dispatcher {
         stale_chunks_notifier: &UnboundedSender<ChunkId>,
     ) -> Result<Vec<(Chunk, Submission)>, sqlx::Error> {
         let mut conn = pool.reader_conn().await?;
+        self.register_reserved_chunk_lookup(conn.get_inner())
+            .await?;
         let mut query_builder = QueryBuilder::new("");
         let stream = strategy
             .build_query(&mut query_builder, &self.metastate)
@@ -79,6 +139,43 @@ impl Dispatcher {
             .await
     }
 
+    async fn register_reserved_chunk_lookup(
+        &self,
+        conn: &mut sqlx::SqliteConnection,
+    ) -> Result<(), sqlx::Error> {
+        let mut handle = conn.lock_handle().await?;
+        let sqlite = handle.as_raw_handle().as_ptr();
+        let function_name = b"opsqueue_is_reserved\0";
+
+        // Register the current reserver state on this connection.
+        // Re-registering replaces any previous callback on this handle.
+        let user_data = Box::new(self.reserver.clone());
+        let user_data = Box::into_raw(user_data).cast::<std::ffi::c_void>();
+
+        let rc = unsafe {
+            ffi::sqlite3_create_function_v2(
+                sqlite,
+                function_name.as_ptr().cast(),
+                2,
+                ffi::SQLITE_UTF8,
+                user_data,
+                Some(sqlite_reserved_chunk_lookup),
+                None,
+                None,
+                Some(sqlite_reserved_chunk_lookup_destructor),
+            )
+        };
+
+        if rc != ffi::SQLITE_OK {
+            unsafe { sqlite_reserved_chunk_lookup_destructor(user_data) };
+            return Err(sqlx::Error::Protocol(format!(
+                "sqlite3_create_function_v2 failed with rc={rc}"
+            )));
+        }
+
+        Ok(())
+    }
+
     async fn reserve_chunk<E>(
         &self,
         chunk: Chunk,
@@ -144,3 +241,58 @@ impl Dispatcher {
             .run_pending_tasks_periodically(cancellation_token);
     }
 }
+
+#[cfg(test)]
+#[cfg(feature = "server-logic")]
+mod test {
+    use super::*;
+    use crate::common::chunk::ChunkId;
+    use crate::common::chunk::ChunkSize;
+    use crate::db::DBPools;
+    use tokio::sync::mpsc::unbounded_channel;
+    use ux::u63;
+
+    #[sqlx::test]
+    async fn fetch_and_reserve_chunks_excludes_already_reserved(db: sqlx::SqlitePool) {
+        let pools = DBPools::from_test_pool(&db);
+        let dispatcher = Dispatcher::new(Duration::from_secs(60));
+        let (stale_chunks_notifier, mut _stale_chunks_receiver) = unbounded_channel::<ChunkId>();
+
+        let mut writer_conn = pools.writer_conn().await.unwrap();
+        let submission_id = crate::common::submission::db::insert_submission_from_chunks(
+            None,
+            vec![Some("a".into()), Some("b".into()), Some("c".into())],
+            None,
+            StrategicMetadataMap::default(),
+            ChunkSize::default(),
+            &mut writer_conn,
+        )
+        .await
+        .unwrap();
+
+        let pre_reserved_chunk = ChunkId::from((submission_id, u63::new(0).into()));
+        dispatcher
+            .reserver()
+            .try_reserve(
+                pre_reserved_chunk,
+                pre_reserved_chunk,
+                &stale_chunks_notifier,
+            )
+            .expect("precondition: pre-reserving chunk should succeed");
+
+        let reserved = dispatcher
+            .fetch_and_reserve_chunks(
+                pools.reader_pool(),
+                strategy::Strategy::Oldest,
+                10,
+                &stale_chunks_notifier,
+            )
+            .await
+            .unwrap();
+
+        assert_eq!(reserved.len(), 2);
+        assert!(reserved.iter().all(|(chunk, _submission)| {
+            ChunkId::from((chunk.submission_id, chunk.chunk_index)) != pre_reserved_chunk
+        }));
+    }
+}

opsqueue/src/consumer/dispatcher/reserver.rs

@@ -82,6 +82,11 @@ where
         }
     }
 
+    /// Returns whether a key currently has an active reservation.
+    pub fn is_reserved(&self, key: &K) -> bool {
+        self.reservations.contains_key(key)
+    }
+
     /// Removes a particular key-val from the reserver.
     /// Afterwards, it is possible to reserve it again.
     ///