SQL: Order by ascending metadata counts, null first

opsqueue/migrations/20260717143000_submission_metadata_pk_by_key.down.sql

@@ -0,0 +1,25 @@
+CREATE TABLE submissions_metadata_old
+(
+    submission_id INTEGER NOT NULL,
+    metadata_key TEXT NOT NULL,
+    metadata_value INTEGER NOT NULL,
+
+    PRIMARY KEY (submission_id, metadata_key, metadata_value)
+) WITHOUT ROWID, STRICT;
+
+INSERT INTO submissions_metadata_old (submission_id, metadata_key, metadata_value)
+SELECT
+      submission_id
+    , metadata_key
+    , metadata_value
+FROM submissions_metadata;
+
+DROP INDEX lookup_submission_by_metadata;
+DROP TABLE submissions_metadata;
+ALTER TABLE submissions_metadata_old RENAME TO submissions_metadata;
+
+CREATE INDEX lookup_submission_by_metadata ON submissions_metadata (
+      metadata_key
+    , metadata_value
+    , submission_id
+);

opsqueue/migrations/20260717143000_submission_metadata_pk_by_key.up.sql

@@ -0,0 +1,29 @@
+CREATE TABLE submissions_metadata_new
+(
+    submission_id INTEGER NOT NULL,
+    metadata_key TEXT NOT NULL,
+    metadata_value INTEGER NOT NULL,
+
+    -- Align with StrategicMetadataMap: one value per key per submission.
+    PRIMARY KEY (submission_id, metadata_key)
+) WITHOUT ROWID, STRICT;
+
+-- If legacy duplicates exist for the same (submission_id, metadata_key),
+-- keep one deterministic value.
+INSERT INTO submissions_metadata_new (submission_id, metadata_key, metadata_value)
+SELECT
+      submission_id
+    , metadata_key
+    , MAX(metadata_value) AS metadata_value
+FROM submissions_metadata
+GROUP BY submission_id, metadata_key;
+
+DROP INDEX lookup_submission_by_metadata;
+DROP TABLE submissions_metadata;
+ALTER TABLE submissions_metadata_new RENAME TO submissions_metadata;
+
+CREATE INDEX lookup_submission_by_metadata ON submissions_metadata (
+      metadata_key
+    , metadata_value
+    , submission_id
+);

opsqueue/src/consumer/dispatcher/metastate.rs

@@ -1,4 +1,3 @@
-use crossbeam_skiplist::SkipSet;
 use dashmap::{DashMap, Entry};
 use rustc_hash::FxBuildHasher;
 use tracing;
@@ -67,7 +66,6 @@ pub type MetaStateVal = i64;
 #[derive(Debug, Default)]
 pub struct MetaStateField {
     vals_to_counts: DashMap<MetaStateVal, usize, FxBuildHasher>,
-    counts_to_vals: SkipSet<(usize, MetaStateVal)>,
 }
 
 impl MetaStateField {
@@ -76,55 +74,30 @@ impl MetaStateField {
     }
 
     fn increment(&self, val: &MetaStateVal) {
-        match self.vals_to_counts.entry(*val) {
-            Entry::Vacant(entry) => {
-                self.counts_to_vals.insert((1, *entry.key()));
-                entry.insert(1);
-            }
-            Entry::Occupied(mut entry) => {
-                // The entry is now locked, so we can also safely update the relevant element of the SkipSet
-                let count = entry.get();
-                let mut set_entry = (*count, *entry.key());
-                self.counts_to_vals.remove(&set_entry);
-                set_entry.0 += 1;
-                self.counts_to_vals.insert(set_entry);
-                *entry.get_mut() += 1;
-            }
-        }
+        self.vals_to_counts
+            .entry(*val)
+            .and_modify(|count| *count += 1)
+            .or_insert(1);
     }
 
     fn decrement(&self, val: &MetaStateVal) {
-        match self.vals_to_counts.entry(*val) {
-            Entry::Vacant(_entry) => {
-                unreachable!()
+        if let Entry::Occupied(entry) = self
+            .vals_to_counts
+            .entry(*val)
+            .and_modify(|count| *count -= 1)
+        {
+            if *entry.get() == 0 {
+                entry.remove();
             }
-            Entry::Occupied(mut entry) => {
-                // The entry is now locked, so we can also safely update the relevant element of the SkipSet
-                let count = entry.get();
-                let mut set_entry = (*count, *entry.key());
-                if *count == 1 {
-                    *entry.get_mut() -= 1;
-                    self.counts_to_vals.remove(&set_entry);
-                    entry.remove();
-                } else {
-                    *entry.get_mut() -= 1;
-                    self.counts_to_vals.remove(&set_entry);
-                    set_entry.0 -= 1;
-                    self.counts_to_vals.insert(set_entry);
-                }
-            }
-        };
+        }
     }
 
     pub fn is_empty(&self) -> bool {
         self.vals_to_counts.is_empty()
     }
 
-    pub fn too_high_counts(&self, max: usize) -> impl Iterator<Item = MetaStateVal> + '_ {
-        tracing::debug!("metastate: {self:?}");
-        self.counts_to_vals
-            .range((max, 0)..)
-            .map(|entry| entry.value().1)
+    pub fn get(&self, val: &MetaStateVal) -> Option<usize> {
+        self.vals_to_counts.get(val).map(|count| *count)
     }
 }
 
@@ -154,14 +127,14 @@ mod tests {
             sut.increment(key, val);
         }
 
-        dbg!(&sut);
-
-        let too_highs: Vec<_> = sut
-            .get(key)
-            .expect("Should exist at this stage")
-            .too_high_counts(group_size)
-            .collect();
-        assert_eq!(too_highs.len(), n_groups);
+        {
+            // We have to release the selected state_field before we can decrement it, otherwise we
+            // would deadlock on the DashMap lock.
+            let state_field = sut.get(key).expect("Should exist at this stage");
+            for group in 0..n_groups {
+                assert_eq!(state_field.get(&(group as i64)), Some(group_size));
+            }
+        }
 
         // Decrement in a different order
         vals.shuffle(&mut rand::rng());

opsqueue/src/consumer/dispatcher/mod.rs

@@ -13,6 +13,7 @@ use libsqlite3_sys as ffi;
 use metastate::MetaState;
 use reserver::Reserver;
 use sqlx::QueryBuilder;
+use std::ffi::CStr;
 use std::time::{Duration, Instant};
 use tokio::sync::mpsc::UnboundedSender;
 use tokio_util::sync::CancellationToken;
@@ -79,6 +80,57 @@ unsafe extern "C" fn sqlite_reserved_chunk_lookup_destructor(ptr: *mut std::ffi:
     let _boxed: Box<Reserver<ChunkId, ChunkId>> = Box::from_raw(ptr.cast());
 }
 
+unsafe extern "C" fn sqlite_metadata_count_lookup(
+    context: *mut ffi::sqlite3_context,
+    n_args: i32,
+    args: *mut *mut ffi::sqlite3_value,
+) {
+    if n_args != 2 {
+        tracing::error!(
+            n_args,
+            "opsqueue_metadata_count called with unexpected argument count"
+        );
+        ffi::sqlite3_result_null(context);
+        return;
+    }
+
+    let user_data = ffi::sqlite3_user_data(context) as *const Arc<MetaState>;
+    if user_data.is_null() {
+        tracing::error!("opsqueue_metadata_count called without registered metastate user_data");
+        ffi::sqlite3_result_null(context);
+        return;
+    }
+
+    let metadata_key_ptr = ffi::sqlite3_value_text(*args.add(0));
+    if metadata_key_ptr.is_null() {
+        ffi::sqlite3_result_null(context);
+        return;
+    }
+    let Ok(metadata_key) = CStr::from_ptr(metadata_key_ptr.cast()).to_str() else {
+        tracing::error!("opsqueue_metadata_count got non-utf8 metadata_key");
+        ffi::sqlite3_result_null(context);
+        return;
+    };
+
+    let metadata_value = ffi::sqlite3_value_int64(*args.add(1));
+
+    if let Some(meta_count) = (*user_data)
+        .get(metadata_key)
+        .and_then(|meta_keys| meta_keys.get(&metadata_value))
+    {
+        ffi::sqlite3_result_int64(context, i64::try_from(meta_count).unwrap_or(i64::MAX));
+    } else {
+        ffi::sqlite3_result_null(context);
+    }
+}
+
+unsafe extern "C" fn sqlite_metadata_count_lookup_destructor(ptr: *mut std::ffi::c_void) {
+    if ptr.is_null() {
+        return;
+    }
+    let _boxed: Box<Arc<MetaState>> = Box::from_raw(ptr.cast());
+}
+
 #[derive(Debug, Clone)]
 pub struct Dispatcher {
     reserver: Reserver<ChunkId, ChunkId>,
@@ -124,11 +176,10 @@ 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?;
+        self.register_lookups(conn.get_inner()).await?;
         let mut query_builder = QueryBuilder::new("");
         let stream = strategy
-            .build_query(&mut query_builder, &self.metastate)
+            .build_query(&mut query_builder)
             .build_query_as()
             .fetch(conn.get_inner());
         stream
@@ -139,13 +190,11 @@ impl Dispatcher {
             .await
     }
 
-    async fn register_reserved_chunk_lookup(
-        &self,
-        conn: &mut sqlx::SqliteConnection,
-    ) -> Result<(), sqlx::Error> {
+    async fn register_lookups(&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";
+        let reserved_function_name = b"opsqueue_is_reserved\0";
+        let metadata_count_function_name = b"opsqueue_metadata_count\0";
 
         // Register the current reserver state on this connection.
         // Re-registering replaces any previous callback on this handle.
@@ -155,7 +204,7 @@ impl Dispatcher {
         let rc = unsafe {
             ffi::sqlite3_create_function_v2(
                 sqlite,
-                function_name.as_ptr().cast(),
+                reserved_function_name.as_ptr().cast(),
                 2,
                 ffi::SQLITE_UTF8,
                 user_data,
@@ -173,6 +222,32 @@ impl Dispatcher {
             )));
         }
 
+        // Register metadata count lookup backed by current metastate.
+        // Re-registering replaces any previous callback on this handle.
+        let user_data = Box::new(self.metastate.clone());
+        let user_data = Box::into_raw(user_data).cast::<std::ffi::c_void>();
+
+        let rc = unsafe {
+            ffi::sqlite3_create_function_v2(
+                sqlite,
+                metadata_count_function_name.as_ptr().cast(),
+                2,
+                ffi::SQLITE_UTF8,
+                user_data,
+                Some(sqlite_metadata_count_lookup),
+                None,
+                None,
+                Some(sqlite_metadata_count_lookup_destructor),
+            )
+        };
+
+        if rc != ffi::SQLITE_OK {
+            unsafe { sqlite_metadata_count_lookup_destructor(user_data) };
+            return Err(sqlx::Error::Protocol(format!(
+                "sqlite3_create_function_v2 failed with rc={rc}"
+            )));
+        }
+
         Ok(())
     }