Sort submissions by metadata count by jerbaroo · Pull Request #164 · channable/opsqueue

jerbaroo · 2026-08-10T13:35:44Z

This PR makes changes to how the chunk selection query of the PreferDistinct strategy. In the Files changed tab, you can see the impact on the benchmark SVG .

Key changes:

Includes a test of fairness in the Python test suite.
We no longer do index walks over all chunks, instead we sort submissions by metadata.
Already-reserved check is from SQLite via FFI, commit from @ReinierMaas
Metadata counts are not passed to SQLite via the generated query, but rather checked from SQLite via FFI, commit from @ReinierMaas

Separate chunk/submission selection

Support random submission selection

Add assert_streaming_chunks

No explicit destructor call

Call the metadata count FFI function

Avoid duplicate MATERIALIZE

Copilot

Pull request overview

Moves PreferDistinct selection from chunk scans to metadata-ranked submission selection with indexed chunk lookup.

Changes:

Ranks submissions using reservation metadata exposed through SQLite FFI.
Adds submission random-order indexing and simplifies metastate tracking.
Adds fairness and reservation-filtering tests.

Reviewed changes

Copilot reviewed 12 out of 15 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
`workspace-hack/Cargo.toml`	Updates generated dependency features.
`opsqueue/src/consumer/strategy.rs`	Reworks strategy query generation.
`opsqueue/src/consumer/dispatcher/reserver.rs`	Exposes reservation lookup.
`opsqueue/src/consumer/dispatcher/mod.rs`	Registers SQLite FFI callbacks.
`opsqueue/src/consumer/dispatcher/metastate.rs`	Simplifies counts and adds JSON export.
`opsqueue/src/common/submission.rs`	Reuses one SQL parameter.
`opsqueue/src/common/chunk.rs`	Adds signed chunk-index conversion.
`opsqueue/migrations/20260803133844_add_random_order_index_to_submissions.up.sql`	Adds submission random ordering.
`opsqueue/migrations/20260803133844_add_random_order_index_to_submissions.down.sql`	Reverts submission random ordering.
`opsqueue/Cargo.toml`	Adds SQLite FFI dependency.
`libs/opsqueue_python/tests/test_roundtrip.py`	Tests `PreferDistinct` fairness.
`Cargo.toml`	Defines the workspace SQLite dependency.
`Cargo.lock`	Updates resolved dependencies.

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

Don't call destructor twice

Fix comment in migration file

ReinierMaas · 2026-08-11T07:53:07Z

There are some CI failures but the changes in the last fixup don't explain the errors on CI... Are they introduced with the rebase?

Re-generate .db file

jerbaroo · 2026-08-11T09:12:47Z

There are some CI failures but the changes in the last fixup don't explain the errors on CI... Are they introduced with the rebase?

Somehow I had incorrectly generated the DB migrations. Re-generated opsqueue/opsqueue_example_database_schema.db in d644aa5

ReinierMaas

I still have to review: opsqueue/src/consumer/strategy.rs

ReinierMaas · 2026-08-11T13:13:49Z

+    # metastate tracks them as reserved, increasing the busy count for that
+    # company.
+    for _ in range(len(company_ids) * chunks_per_company):
+        [chunk] = consumer_client.reserve_chunks(strategy=strategy)


I was just thinking, does our new strategy selection retain randomness if a client reserves multiple chunks in one go? Do we care for that or would documenting the multi-chunk reservation behaviour be enough?

ReinierMaas · 2026-08-11T13:41:10Z

+        Some(field) => field.to_json(),
+        None => "{}".to_string(),
+    };
+    let len = i32::try_from(json.len()).unwrap_or(i32::MAX);


We should instead report to SQLite that the requested data is too big:

https://docs.rs/rust-sqlite/latest/sqlite3/ffi/fn.sqlite3_result_error_toobig.html

Note: I have included this in the example update to the FFI-code below.

ReinierMaas · 2026-08-11T13:42:49Z

+    #[must_use]
+    pub fn to_json(&self) -> String {
+        use std::fmt::Write as _;
+        let mut out = String::from("{");
+        for entry in &self.vals_to_counts {
+            if out.len() > 1 {
+                out.push(',');
+            }
+            let _ = write!(out, "\"{}\":{}", entry.key(), entry.value());
+        }
+        out.push('}');
+        out
    }


Let's not implement our own JSON serialiser, these keys are user controlled and we need to implement correct escaping:

Suggested change

#[must_use]

pub fn to_json(&self) -> String {

use std::fmt::Write as _;

let mut out = String::from("{");

for entry in &self.vals_to_counts {

if out.len() > 1 {

out.push(',');

}

let _ = write!(out, "\"{}\":{}", entry.key(), entry.value());

}

out.push('}');

out

}

#[must_use]

pub fn to_json(&self) -> String {

serde_json::to_string(&self.vals_to_counts)

.unwrap_or_else(|_| String::from("{}"))

}

diff --git a/Cargo.lock b/Cargo.lock index fe470f8..719bb4a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -647,6 +647,7 @@ dependencies = [ "lock_api", "once_cell", "parking_lot_core", + "serde", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 44ce405..27727cc 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -28,7 +28,7 @@ ciborium = { version = "0.2.2", features = ["std"] } clap = { version = "4.5.60", features = ["derive", "std"] } criterion = { version = "0.8", features = ["async_tokio"] } crossbeam-skiplist = { version = "0.1.3", features = ["std"] } -dashmap = { version = "6.1.0" } +dashmap = { version = "6.1.0", features = ["serde"] } either = { version = "1.13.0" } futures = { version = "0.3.32" } http = { version = "1.4.0" }

ReinierMaas · 2026-08-11T14:15:14Z

+    let json = match unsafe { &*user_data }.get(metadata_key) {
+        Some(field) => field.to_json(),
+        None => "{}".to_string(),
+    };
+    let len = i32::try_from(json.len()).unwrap_or(i32::MAX);
+    unsafe {
+        // SQLITE_TRANSIENT tells SQLite to copy the bytes before we drop them.
+        ffi::sqlite3_result_text(context, json.as_ptr().cast(), len, ffi::SQLITE_TRANSIENT());
+    };


We will be performing this query often and also for submissions without any connected metadata. I have changed it a bit to be lighter on allocations, no allocations for static, no copy for JSON (CString resizes the underlying String allocation iff the null-byte wouldn't fit) and report the too big error correctly instead of truncated JSON:

Suggested change

let json = match unsafe { &*user_data }.get(metadata_key) {

Some(field) => field.to_json(),

None => "{}".to_string(),

};

let len = i32::try_from(json.len()).unwrap_or(i32::MAX);

unsafe {

// SQLITE_TRANSIENT tells SQLite to copy the bytes before we drop them.

ffi::sqlite3_result_text(context, json.as_ptr().cast(), len, ffi::SQLITE_TRANSIENT());

};

match unsafe { &*user_data }.get(metadata_key) {

Some(field) => {

let json = field.to_json();

match i32::try_from(json.len()) {

Ok(len) => {

unsafe {

// No SQLite allocation by passing a destructor!

ffi::sqlite3_result_text(

context,

CString::new(json).expect("JSON doesn't contain null-bytes").into_raw().cast(),

len,

Some(sqlite_cstring_destructor),

);

}

},

Err(_) => unsafe {

// Report to SQLite that the generated response is too large...

ffi::sqlite3_result_error_toobig(context) },

}

}

None => {

unsafe {

// No Rust allocation, and SQLITE_STATIC means no SQLite allocation!

ffi::sqlite3_result_text(

context,

c"{}".as_ptr().cast(),

2,

ffi::SQLITE_STATIC(),

);

}

}

}

diff --git a/opsqueue/src/consumer/dispatcher/mod.rs b/opsqueue/src/consumer/dispatcher/mod.rs index 728566b..63bd606 100644 --- a/opsqueue/src/consumer/dispatcher/mod.rs +++ b/opsqueue/src/consumer/dispatcher/mod.rs @@ -13,7 +13,7 @@ use libsqlite3_sys as ffi; use metastate::MetaState; use reserver::Reserver; use sqlx::QueryBuilder; -use std::ffi::CStr; +use std::ffi::{CStr, CString}; use std::time::{Duration, Instant}; use tokio::sync::mpsc::UnboundedSender; use tokio_util::sync::CancellationToken; @@ -130,6 +130,14 @@ unsafe extern "C" fn sqlite_metadata_count_lookup( } } +unsafe extern "C" fn sqlite_cstring_destructor(ptr: *mut std::ffi::c_void) { + if !ptr.is_null() { + unsafe { + drop(CString::from_raw(ptr.cast())); + } + } +} + /// Returns the whole value -> count map for one metadata key as a JSON object, /// so the caller can obtain every count in a single FFI call. unsafe extern "C" fn sqlite_metadata_counts_lookup(

ReinierMaas

LGTM! I left some more changes but I think that covers all.

ReinierMaas · 2026-08-12T10:55:15Z

+}
+
+impl<'a> MetaKeysIter<'a> {
+    /// The first non-[`Strategy::PreferDistinct`] strategy in the chain.


After the iterator has been exhaustedly polled otherwise the next strategy is returned.

ReinierMaas · 2026-08-12T10:57:50Z

-        metastate: &MetaState,
    ) -> &'a mut QueryBuilder<Sqlite> {
        use Strategy::{Newest, Oldest, PreferDistinct, Random};
+        let ffi_is_reserved = "opsqueue_is_reserved(chunks.submission_id, chunks.chunk_index) = 0";


This is compared to 0 zero, i.e. false but SQLite doesn't have booleans:

Suggested change

let ffi_is_reserved = "opsqueue_is_reserved(chunks.submission_id, chunks.chunk_index) = 0";

let ffi_is_not_reserved = "opsqueue_is_reserved(chunks.submission_id, chunks.chunk_index) = 0";

ReinierMaas · 2026-08-12T11:13:12Z

+                        "))
+                            ffi_counts


Is this AS ffi_counts, without the explicit AS? I think it is easier to follow if written as:

Suggested change

"))

ffi_counts

")) AS ffi_counts

ReinierMaas · 2026-08-12T11:22:15Z

+                .push("SELECT * FROM chunks")
+                .push(format!(" WHERE {ffi_is_reserved}"))
+                .push(" ORDER BY submission_id DESC"),
+            Random => Self::push_random_order_query(qb, "*", "chunks", Some(ffi_is_reserved)),


Shall we use single push and add the whole query in one go, like at line 100?

ReinierMaas · 2026-08-12T11:23:12Z

+                // In-flight chunk count per submission, per meta key.
+                //
+                // The FFI call returns all counts as JSON in a single call.
+                // The CROSS JOIN ON ensures the json_each is the outer loop,
+                // and only performed once.
+                for (i, meta_key) in meta_keys.iter().enumerate() {
+                    qb.push(format!(
+                        ", counts_{i} AS (
+                            SELECT sm.submission_id, ffi_counts.value AS count
+                            FROM json_each(opsqueue_metadata_counts("
+                    ));
+                    qb.push_bind(*meta_key);


Good to use enumerate here to prevent SQL injection!

ReinierMaas · 2026-08-12T11:28:53Z

+            "register opsqueue_metadata_count failed"
+        );
+
+        let function_name = b"opsqueue_metadata_counts\0";


Suggested change

let function_name = b"opsqueue_metadata_counts\0";

let function_name = c"opsqueue_metadata_counts";

ReinierMaas · 2026-08-12T11:29:50Z

+        let reserved_function_name = b"opsqueue_is_reserved\0";
+        let metadata_count_function_name = b"opsqueue_metadata_count\0";


Suggested change

let reserved_function_name = b"opsqueue_is_reserved\0";

let metadata_count_function_name = b"opsqueue_metadata_count\0";

let reserved_function_name = c"opsqueue_is_reserved";

let metadata_count_function_name = c"opsqueue_metadata_count";

ReinierMaas · 2026-08-12T11:30:46Z

+        }
+
+        // Register the bulk metadata counts lookup backed by current metastate.
+        let counts_function_name = b"opsqueue_metadata_counts\0";


Should we also move this name to the top of the function like the other names we are recording here?

Suggested change

let counts_function_name = b"opsqueue_metadata_counts\0";

let counts_function_name = c"opsqueue_metadata_counts";

ReinierMaas · 2026-08-12T11:31:24Z

+                // we also keep the row number from the underlying query, this
+                // is used as a tie-breaker if metadata counts are equal.
+                let qb = qb.push("WITH inner AS NOT MATERIALIZED (");
+                qb.push("SELECT submission_id, ROW_NUMBER() OVER () as underlying_row FROM ( ");


These pushes can be combined.

ReinierMaas · 2026-08-12T11:37:35Z

+    /// Append a query snippet to select from the `random_order` column on the
+    /// given table using the "cutting the deck" technique.
+    fn push_random_order_query<'a>(
+        qb: &'a mut QueryBuilder<Sqlite>,
+        columns: &str,
+        table_name: &str,
+        condition: Option<&str>,
+    ) -> &'a mut QueryBuilder<Sqlite> {


There is potential for SQL injection here we only pass static strings currently we can reduce the change of using this wrong, i.e. temporary strings such as from a request, by requiring the columns, tables and condition to be of type &'static str. We should also document this in the function documentation. SQLx doesn't expose an escape identifier...

ReinierMaas and others added 13 commits July 13, 2026 17:02

SQL: Filter already reserved chunks from Chunk selection

db1de29

Add failing PreferDistinct fairness test

a6664b3

PreferDistinct sorts submissions by metadata

ada8709

fixup! PreferDistinct sorts submissions by metadata

b40b106

Separate chunk/submission selection

fixup! PreferDistinct sorts submissions by metadata

31cf874

Support random submission selection

fixup! PreferDistinct sorts submissions by metadata

7c1d688

Add assert_streaming_chunks

SQL: Prevent duplicated binds

c903fb1

fixup! SQL: Filter already reserved chunks from Chunk selection

7dfef50

No explicit destructor call

SQL: Add metadata count FFI function

b48dd48

fixup! SQL: Add metadata count FFI function

573b2d6

Call the metadata count FFI function

Flatten PreferDistinct into a single query level

1d2f36a

fixup! Flatten PreferDistinct into a single query level

02fc44d

Avoid duplicate MATERIALIZE

SQL: Deduplicate FFI calls per distinct metadata value

98f022d

jerbaroo mentioned this pull request Aug 10, 2026

PreferDistinct: faster and fairer #148

Closed

jerbaroo requested a balanced review from Copilot August 10, 2026 13:49

Copilot started reviewing on behalf of jerbaroo August 10, 2026 13:49 View session

Copilot AI reviewed Aug 10, 2026

View reviewed changes

Comment thread opsqueue/src/consumer/strategy.rs

Comment thread opsqueue/src/consumer/dispatcher/mod.rs

Comment thread opsqueue/migrations/20260803133844_add_random_order_index_to_submissions.up.sql Outdated

jerbaroo force-pushed the jerbaroo-faster-prefer-distinct branch from 735f40b to 525fbdf Compare August 10, 2026 14:44

jerbaroo added 3 commits August 10, 2026 17:03

SQL: Single FFI call

efe45e2

Break ties of PreferDistinct submissions with equal counts

3d9f2c3

fixup! SQL: Add metadata count FFI function

4887d3d

Don't call destructor twice

jerbaroo force-pushed the jerbaroo-faster-prefer-distinct branch from 525fbdf to 4887d3d Compare August 10, 2026 15:04

fixup! fixup! PreferDistinct sorts submissions by metadata

7903f8b

Fix comment in migration file

jerbaroo marked this pull request as ready for review August 10, 2026 15:14

jerbaroo requested a review from ReinierMaas August 10, 2026 15:20

jerbaroo self-assigned this Aug 11, 2026

fixup! fixup! fixup! PreferDistinct sorts submissions by metadata

d644aa5

Re-generate .db file

ReinierMaas requested changes Aug 11, 2026

View reviewed changes

ReinierMaas self-requested a review August 11, 2026 15:01

ReinierMaas approved these changes Aug 12, 2026

View reviewed changes

	let ffi_is_reserved = "opsqueue_is_reserved(chunks.submission_id, chunks.chunk_index) = 0";
	let ffi_is_not_reserved = "opsqueue_is_reserved(chunks.submission_id, chunks.chunk_index) = 0";

	let function_name = b"opsqueue_metadata_counts\0";
	let function_name = c"opsqueue_metadata_counts";

		let reserved_function_name = b"opsqueue_is_reserved\0";
		let metadata_count_function_name = b"opsqueue_metadata_count\0";

	let counts_function_name = b"opsqueue_metadata_counts\0";
	let counts_function_name = c"opsqueue_metadata_counts";

Uh oh!

Conversation

jerbaroo commented Aug 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ReinierMaas commented Aug 11, 2026

Uh oh!

jerbaroo commented Aug 11, 2026

Uh oh!

ReinierMaas left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ReinierMaas Aug 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ReinierMaas left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

jerbaroo commented Aug 10, 2026 •

edited

Loading

ReinierMaas Aug 11, 2026 •

edited

Loading