diff --git a/Cargo.lock b/Cargo.lock index e2bc16b..7af1e64 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -147,7 +147,7 @@ checksum = "ae36dc4177970ef04fde5178d3e2429882def40e57a451f919c098f72baa6cec" dependencies = [ "proc-macro2", "quote", - "syn 3.0.2", + "syn 3.0.3", ] [[package]] @@ -451,9 +451,9 @@ dependencies = [ [[package]] name = "clap" -version = "4.6.3" +version = "4.6.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0fb99565819980999fb7b4a1796046a5c949e6d4ff132cf5fadf5a641e20d776" +checksum = "d91e0c145792ef73a6ad36d27c75ac09f1832222a3c209689d90f534685ee5b7" dependencies = [ "clap_builder", "clap_derive", @@ -473,14 +473,14 @@ dependencies = [ [[package]] name = "clap_derive" -version = "4.6.3" +version = "4.6.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "32f2392eae7f16557a3d727ef3a12e57b2b2ca6f98566a5f4fb41ffe305df077" +checksum = "d012d2b9d65aca7f18f4d9878a045bc17899bba951561ba5ec3c2ba1eed9a061" dependencies = [ "heck", "proc-macro2", "quote", - "syn 2.0.119", + "syn 3.0.3", ] [[package]] @@ -1631,9 +1631,9 @@ dependencies = [ [[package]] name = "libc" -version = "0.2.188" +version = "0.2.189" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "22053b6a34f84abc97f9129e61334f40174659a1b9bd18c970b83db6a9a6348b" +checksum = "3eaf3ede3fee6db1a4c2ee091bf8a8b4dccdc6d17f656fb07896ee72867612f2" [[package]] name = "libsqlite3-sys" @@ -2183,7 +2183,7 @@ dependencies = [ [[package]] name = "opsqueue" -version = "0.36.0" +version = "0.35.5" dependencies = [ "anyhow", "arc-swap", @@ -2203,6 +2203,7 @@ dependencies = [ "humantime", "insta", "itertools 0.15.0", + "libsqlite3-sys", "moka", "moro-local", "object_store", @@ -2238,7 +2239,7 @@ dependencies = [ [[package]] name = "opsqueue_python" -version = "0.36.0" +version = "0.35.5" dependencies = [ "anyhow", "chrono", @@ -3158,7 +3159,7 @@ checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" dependencies = [ "proc-macro2", "quote", - "syn 3.0.2", + "syn 3.0.3", ] [[package]] @@ -3563,9 +3564,9 @@ dependencies = [ [[package]] name = "syn" -version = "3.0.2" +version = "3.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a207d6d6a2b7fc470b80443726053f18a2481b7e1eee970597051596567987a3" +checksum = "53e9bae58849f64dfa4f5d5ae372c8341f7305f82a3868709269343628b659a3" dependencies = [ "proc-macro2", "quote", @@ -3634,7 +3635,7 @@ checksum = "43cbfe0cf76104d42a574802844187e84a305e531ed54455f11fbde0f10541cd" dependencies = [ "proc-macro2", "quote", - "syn 3.0.2", + "syn 3.0.3", ] [[package]] @@ -3750,9 +3751,9 @@ dependencies = [ [[package]] name = "tokio-stream" -version = "0.1.18" +version = "0.1.19" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "32da49809aab5c3bc678af03902d4ccddea2a87d028d86392a4b1560c6906c70" +checksum = "a3d06f0b082ba57c26b79407372e57cf2a1e28124f78e9479fe80322cf53420b" dependencies = [ "futures-core", "pin-project-lite", diff --git a/Cargo.toml b/Cargo.toml index 17b6eae..7bf754d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,7 +1,7 @@ [workspace] -# Resolver 3 is default in the 2024 Edition, but the workspace doesn't know that -resolver = "3" +# Resolver 2 is default in the 2021 Edition, but the workspace doesn't know that +resolver = "2" members = [ "opsqueue/", @@ -9,7 +9,7 @@ members = [ ] [workspace.package] -version = "0.36.0" +version = "0.35.5" [workspace.lints.clippy] cargo = { level = "warn", priority = -1 } @@ -27,10 +27,6 @@ restriction = { level = "allow", priority = -1 } style = { level = "warn", priority = -1 } suspicious = { level = "warn", priority = -1 } -# Prevent duplicate versions of crates from being used in the workspace. If an easy fix is not -# possible, you can add the crate to the `allowed-duplicate-crates` list in `clippy.toml`. -multiple_crate_versions = "deny" - [profile.release] # Full fat LTO makes a big difference in performance, at the cost of a large # compile time hit. You should only use release builds when performance matters. diff --git a/default.nix b/default.nix index 2da1a73..a46cd37 100644 --- a/default.nix +++ b/default.nix @@ -63,6 +63,8 @@ let # For the shell, libpython needs to be in the search path. pythonEnv ]; + # Fixes: libstdc++.so.6: cannot open shared object file: No such file or directory + LD_LIBRARY_PATH = pkgs.lib.makeLibraryPath [ pkgs.stdenv.cc.cc.lib ]; }; }; in diff --git a/justfile b/justfile index c73f5c6..f9f9a7a 100644 --- a/justfile +++ b/justfile @@ -53,7 +53,7 @@ test-integration *TEST_ARGS: build-bin build-python cd libs/opsqueue_python source "./.setup_local_venv.sh" - timeout {{pytest_timeout_seconds}} pytest --color=yes {{TEST_ARGS}} + pytest --timeout {{pytest_timeout_seconds}} --color=yes {{TEST_ARGS}} # Python integration test suite, using artefacts built through Nix. Args are forwarded to pytest [group('nix')] @@ -68,7 +68,19 @@ nix-test-integration *TEST_ARGS: nix-build export OPSQUEUE_BIN="${nix_build_bin_dir}/bin/opsqueue" export RUST_LOG="opsqueue=debug" - timeout {{pytest_timeout_seconds}} pytest --color=yes {{TEST_ARGS}} + pytest --timeout {{pytest_timeout_seconds}} --color=yes {{TEST_ARGS}} + +# Benchmark query cost and render a graph. +[group('bench')] +bench-chunks-select: + #!/usr/bin/env bash + set -euo pipefail + cargo bench --bench chunks_select + + cd libs/opsqueue_python + source "./.setup_local_venv.sh" + cd - + python opsqueue/benches/plot_chunks_select.py # Run all linters, fast and slow [group('lint')] diff --git a/libs/opsqueue_python/pyproject.toml b/libs/opsqueue_python/pyproject.toml index c30ad39..74c2512 100644 --- a/libs/opsqueue_python/pyproject.toml +++ b/libs/opsqueue_python/pyproject.toml @@ -41,9 +41,10 @@ include = ["**/py.typed"] # because otherwise `pytest` won't be able to see the locally build python package! [project.optional-dependencies] test = [ + "matplotlib==3.11.1", + "multiprocess==0.70.19", "pytest==9.0.3", "pytest-xdist==3.8.0", - "multiprocess==0.70.19", "pytest-timeout==2.4.0", ] diff --git a/libs/opsqueue_python/tests/test_roundtrip.py b/libs/opsqueue_python/tests/test_roundtrip.py index 23d9a14..2e887ac 100644 --- a/libs/opsqueue_python/tests/test_roundtrip.py +++ b/libs/opsqueue_python/tests/test_roundtrip.py @@ -667,3 +667,32 @@ def test_lookup_too_many_submission_ids_by_strategic_metadata() -> None: ) assert exc.type is TooManyMatchingSubmissionsError assert exc.value.max_submissions == max_ + + +def test_prefer_distinct_strategy_fairness(opsqueue: OpsqueueProcess) -> None: + """Test the PreferDistinct strategy fairly interleaves chunks from different + submissions based on their strategic metadata. + + """ + url = "file:///tmp/opsqueue/test_prefer_distinct_fairness" + producer_client = ProducerClient(f"localhost:{opsqueue.port}", url) + consumer_client = ConsumerClient(f"localhost:{opsqueue.port}", url) + company_ids = [1, 2, 3] + chunks_per_company = 4 + company_id_per_submission = {} + for company_id in company_ids: + sub_id = producer_client.insert_submission( + [company_id] * chunks_per_company, + chunk_size=1, + strategic_metadata={"company_id": company_id}, + ) + company_id_per_submission[sub_id] = company_id + strategy = strategy_from_description(("PreferDistinct", "company_id", "Oldest")) + reserved_company_order = [] + # Fetch 1 chunk at a time. Because we don't complete chunks, opsqueue's + # 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) + reserved_company_order.append(company_id_per_submission[chunk.submission_id]) + assert reserved_company_order == [1, 2, 3] * chunks_per_company diff --git a/opsqueue/Cargo.toml b/opsqueue/Cargo.toml index 2d69b21..e5610a2 100644 --- a/opsqueue/Cargo.toml +++ b/opsqueue/Cargo.toml @@ -33,6 +33,7 @@ ux = "0.1.6" anyhow = "1.0.102" # Database: sqlx = { version = "0.9.0", 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" @@ -83,13 +84,14 @@ sqlformat = "0.5.0" workspace = true [dev-dependencies] +assert_matches = { version = "1.5.0" } criterion = {version = "0.8", features = ["async_tokio"]} insta = { version = "1.47.2" } -assert_matches = { version = "1.5.0" } -# [[bench]] -# name = "chunks_select" -# harness = false +[[bench]] +name = "chunks_select" +harness = false +required-features = ["server-logic"] # [[bench]] # name = "submissions_insert" @@ -114,4 +116,4 @@ client-logic = [ "dep:tokio-tungstenite", "dep:arc-swap", ] -default = ["server-logic"] +default = ["server-logic", "dep:libsqlite3-sys"] diff --git a/opsqueue/benches/chunks_select.rs b/opsqueue/benches/chunks_select.rs new file mode 100644 index 0000000..857a096 --- /dev/null +++ b/opsqueue/benches/chunks_select.rs @@ -0,0 +1,283 @@ +/// FOR EACH shape: +/// FOR EACH strategy: +/// FOR EACH amount of chunks: +/// +/// // Setup +/// Create fresh temporary SQLite database +/// Insert chunks, assigning them to companies/submissions based on `shape` +/// +/// // MEASURE +/// FOR _ in (SAMPLES + WARMUP): +/// Start Timer +/// Execute SQL Query -> Fetch EXACTLY ONE chunk +/// Stop Timer +/// IF (not warmup): save duration +/// +/// // REPORT & CLEANUP +/// Calculate stats +/// Write result to terminal and CSV +/// Delete temporary database files +use opsqueue::common::StrategicMetadataMap; +use opsqueue::common::chunk::{ChunkId, ChunkSize}; +use opsqueue::common::submission::db::insert_submission_from_chunks; +use opsqueue::consumer::dispatcher::Dispatcher; +use opsqueue::consumer::strategy::Strategy; +use opsqueue::db::{self}; +use std::io::Write; +use std::num::NonZero; +use std::path::PathBuf; +use std::time::{Duration, Instant}; + +// Human readable name for an OpsQueue strategy. +type StrategyName = &'static str; + +// Shape of the data we are inserting. +#[derive(Debug, Eq, Ord, PartialEq, PartialOrd)] +enum Shape { + FewSubmissionsManyChunks, + ManySubmissionsFewChunks, + Realistic, +} + +const CHUNKS_PER_COMPANY: u64 = CHUNKS_PER_SUBMISSION * SUBMISSIONS_PER_COMPANY; +const CHUNKS_PER_SUBMISSION: u64 = 1024; +// Increase in the total amount of chunks for the `Shape::Realistic` strategy. +const CHUNKS_STEP: usize = 500_000; +// Amount of companies in Channable 23 July 2026. +const COMPANIES: u64 = 93221; +// We use this as an approximation of the amount of submissions. +const EXPORTS: u64 = 303_262; +const MAX_CHUNKS: u64 = 10_000_001; +// Unused: amount of projects in Channable 23 July 2026 +// const PROJECTS: u64 = 166774; +// The amount of samples to collect per iteration of the +// shape/strategy/amount_of_chunks main loop. Note that each "sample" requires a +// number of reservations to be performed. See `bench_strategy` for more info. +const SAMPLES_PER_DATUM: usize = 10; +const SUBMISSIONS_PER_COMPANY: u64 = EXPORTS / COMPANIES; +// The amount of reservations to collect a single sample. +// Note that the time between these is NOT INDEPENDENT: +// the earlier reservations affect the later ones. +const RESERVATIONS_IN_SAMPLE: usize = 50; +// The amount of reservations to perform before we begin collecting a sample. +const RESERVATIONS_IN_WARMUP: usize = 5; + +#[derive(Debug)] +pub struct BenchStats { + pub median: f64, // Median of medians + pub p10: f64, // Lower bound of medians + pub p90: f64, // Upper bound of medians +} + +#[allow(clippy::missing_panics_doc)] +impl BenchStats { + #[must_use] + pub fn new(runs: Vec>) -> Self { + // Calculate the median for each individual run + let mut run_medians: Vec = runs + .into_iter() + .filter_map(|mut run_samples| { + if run_samples.is_empty() { + None + } else { + run_samples.sort_by(|a, b| a.partial_cmp(b).unwrap()); + Some(run_samples[run_samples.len() / 2]) + } + }) + .collect(); + if run_medians.is_empty() { + return BenchStats { + p10: 0., + median: 0., + p90: 0., + }; + } + // Sort the N medians to find our distribution. + run_medians.sort_by(|a, b| a.partial_cmp(b).unwrap()); + let len = run_medians.len(); + BenchStats { + p10: run_medians[len * 10 / 100], + median: run_medians[len / 2], + p90: run_medians[len * 90 / 100], + } + } +} + +/// Maximum amount of chunks for the bench run. +/// +/// This is based on the shape of the data, as that can dramatically affect +/// run-time, and for some shapes we can push the total chunks higher without +/// waiting too long. +fn max_chunks_by_shape(shape: &Shape) -> Vec { + let mut vector: Vec = vec![ + 100, 500, 1_000, 2_000, 4_000, 6_000, 8_000, 10_000, 15_000, 20_000, 30_000, + ]; + if shape == &Shape::Realistic { + let max = MAX_CHUNKS.min(COMPANIES * CHUNKS_PER_COMPANY); + vector.extend((50_000..=max).step_by(CHUNKS_STEP)); + } + vector +} + +/// All the strategies we are benchmarking. +fn strategies() -> [(StrategyName, Strategy); 2] { + [ + ("Random", Strategy::Random), + ( + "PreferDistinct(company_id, Oldest)", + Strategy::PreferDistinct { + meta_key: "company_id".to_string(), + underlying: Box::new(Strategy::Oldest), + }, + ), + ] +} + +/// (total companies, `submissions_per_company`, chunks per submission) for a shape and total chunks. +fn layout(shape: &Shape, total_chunks: u64) -> (u64, u64, u64) { + match shape { + Shape::ManySubmissionsFewChunks => (total_chunks, 1, 1), + Shape::FewSubmissionsManyChunks => { + let submissions = 10.min(total_chunks.max(1)); + (submissions, 1, total_chunks.div_ceil(submissions)) + } + Shape::Realistic => { + let companies = total_chunks.div_ceil(CHUNKS_PER_COMPANY); + (companies, SUBMISSIONS_PER_COMPANY, CHUNKS_PER_SUBMISSION) + } + } +} + +/// Seeds a fresh DB with chunks according to a given `Shape`. +async fn seed(shape: &Shape, total_chunks: u64) -> (db::DBPools, Dispatcher, PathBuf) { + let (total_companies, submissions_per_company, chunks_per_submission) = + layout(shape, total_chunks); + println!( + " Companies: {total_companies:<30}\n Submissions per company: {submissions_per_company}\n Chunks per submission: {chunks_per_submission}" + ); + let db_path = std::env::temp_dir().join("opsqueue_bench.sqlite"); + // Remove the DB before the run. + let _ = std::fs::remove_file(&db_path); + let db_pools = db::open_and_setup(db_path.to_str().unwrap(), NonZero::new(16).unwrap()).await; + let mut conn = db_pools.writer_conn().await.unwrap(); + let dispatcher = Dispatcher::new(Duration::from_mins(1_000_000)); + for company_id in 0..total_companies { + for _ in 0..submissions_per_company { + let mut metadata = StrategicMetadataMap::default(); + metadata.insert("company_id".to_string(), i64::try_from(company_id).unwrap()); + let chunks = vec![Some(b"x".to_vec()); usize::try_from(chunks_per_submission).unwrap()]; + insert_submission_from_chunks( + None, + chunks, + None, + metadata, + ChunkSize::default(), + &mut conn, + ) + .await + .unwrap(); + } + } + (db_pools, dispatcher, db_path) +} + +/// Runs the selection query and fetches just the first chunk. +async fn fetch_and_reserve_chunk( + db_pools: &db::DBPools, + strategy: Strategy, + dispatcher: &Dispatcher, +) -> ChunkId { + use tokio::sync::mpsc::unbounded_channel; + let (notifier, _) = unbounded_channel(); + let reserved = dispatcher + .fetch_and_reserve_chunks(db_pools.reader_pool(), strategy, 1, ¬ifier) + .await + .unwrap(); + if let [(chunk, _)] = reserved.as_slice() { + ChunkId::from((chunk.submission_id, chunk.chunk_index)) + } else { + panic!( + "Yowza! Expected exactly 1 chunk, but got {}", + reserved.len() + ); + } +} + +/// Executes reservations for a given DB state and returns `BenchStats`. +async fn bench_strategy( + db_pools: &db::DBPools, + strategy: &Strategy, + dispatcher: &Dispatcher, +) -> BenchStats { + let mut all_reservation_durations: Vec> = Vec::with_capacity(SAMPLES_PER_DATUM); + // For each of the "samples" we want to collect. + for _ in 0..SAMPLES_PER_DATUM { + let mut sample_reservation_durations = Vec::with_capacity(RESERVATIONS_IN_SAMPLE); + let mut chunks_reserved = Vec::new(); + // For each sample, reserve N chunks (some are warmups). + for i in 0..(RESERVATIONS_IN_WARMUP + RESERVATIONS_IN_SAMPLE) { + let start = Instant::now(); + let chunk_id = fetch_and_reserve_chunk(db_pools, strategy.clone(), dispatcher).await; + if i >= RESERVATIONS_IN_WARMUP { + sample_reservation_durations.push(start.elapsed().as_secs_f64() * 1e6); + } + chunks_reserved.push(chunk_id); + } + all_reservation_durations.push(sample_reservation_durations); + // Reset the queue state for the next run + let mut conn = db_pools.writer_conn().await.unwrap(); + for chunk_id in chunks_reserved { + dispatcher + .finish_reservation(&mut conn, chunk_id, false) + .await; + } + } + BenchStats::new(all_reservation_durations) +} + +fn main() { + let runtime = tokio::runtime::Runtime::new().unwrap(); + let csv_path = PathBuf::from("chunks_select_bench.csv"); + if let Some(parent) = csv_path.parent() { + let _ = std::fs::create_dir_all(parent); + } + let mut csv = std::fs::File::create(&csv_path) + .unwrap_or_else(|e| panic!("Failed to create CSV at {}: {}", csv_path.display(), e)); + writeln!(csv, "shape,strategy,backlog_size,p10_us,median_us,p90_us").unwrap(); + println!( + "{:<30} {:<35} {:<12} {:<10} {:<15}", + "SHAPE", "STRATEGY", "SIZE", "MEDIAN", "P10 / P90" + ); + println!("{}", "-".repeat(105)); + for shape in &[ + Shape::FewSubmissionsManyChunks, + Shape::ManySubmissionsFewChunks, + Shape::Realistic, + ] { + for (strategy_label, strategy) in strategies() { + for size in max_chunks_by_shape(shape) { + let (db_pools, dispatcher, path) = runtime.block_on(seed(shape, size)); + let stats = runtime.block_on(bench_strategy(&db_pools, &strategy, &dispatcher)); + let bounds = format!("{:.1} / {:.1}", stats.p10, stats.p90); + println!( + "{:<30} {:<35} {:<12} {:<10.1} {:<15}", + format!("{shape:?}"), + strategy_label, + size, + stats.median, + bounds + ); + writeln!( + csv, + "{shape:?},\"{strategy_label}\",{size},{:.1},{:.1},{:.1}", + stats.p10, stats.median, stats.p90 + ) + .unwrap(); + drop(db_pools); + let _ = std::fs::remove_file(&path); + let _ = std::fs::remove_file(path.with_extension("sqlite-wal")); + let _ = std::fs::remove_file(path.with_extension("sqlite-shm")); + } + } + } +} diff --git a/opsqueue/benches/chunks_select_bench.png b/opsqueue/benches/chunks_select_bench.png new file mode 100644 index 0000000..5d9380f Binary files /dev/null and b/opsqueue/benches/chunks_select_bench.png differ diff --git a/opsqueue/benches/plot_chunks_select.py b/opsqueue/benches/plot_chunks_select.py new file mode 100644 index 0000000..2d7346e --- /dev/null +++ b/opsqueue/benches/plot_chunks_select.py @@ -0,0 +1,65 @@ +#!/usr/bin/env python3 +"""Plot the chunks_select benchmark: one subplot per backlog shape, a line per strategy. + +Reads target/chunks_select_bench.csv (written by the bench) and writes a PNG next to +this script so it can be committed alongside the code it measures. + +""" + +import csv +import matplotlib.pyplot as plt # noqa: E402 +from collections import defaultdict + +CSV_PATH = "opsqueue/chunks_select_bench.csv" +OUT_PATH = "chunks_select_bench.svg" + + +def main() -> None: + # (shape, strategy) -> {backlog_size: (p10_us, median_us, p90_us)} + series = defaultdict(lambda: defaultdict(dict)) + shapes: list[str] = [] + max_x = 0 + + print(f"Reading CSV: {CSV_PATH}") + with open(CSV_PATH, newline="") as f: + for row in csv.DictReader(f): + shape = row["shape"] + size = int(row["backlog_size"]) + max_x = max(max_x, size) + series[shape][row["strategy"]][size] = ( + float(row["p10_us"]), + float(row["median_us"]), + float(row["p90_us"]), + ) + if shape not in shapes: + shapes.append(shape) + + fig, axes = plt.subplots( + 1, len(shapes), figsize=(7.5 * len(shapes), 5.5), squeeze=False + ) + fig.suptitle("Reservation duration (no network in path)") + + for ax, shape in zip(axes.flatten(), shapes): + for strategy, points in sorted(series[shape].items()): + xs = sorted(points) + p10s, medians, p90s = zip(*[points[x] for x in xs]) + (line,) = ax.plot(xs, medians, marker="o", label=strategy) + ax.fill_between(xs, p10s, p90s, color=line.get_color(), alpha=0.2) + + ax.set_xscale("log") + ax.set_yscale("log") + ax.set_xlim(right=max_x) + ax.set_ylim(top=1_000_000) + ax.set_title(shape) + ax.set_xlabel("Backlog size (total chunks)") + ax.set_ylabel("Latency to select first chunk (us) [Median ± P10/P90]") + ax.grid(True, which="both", linestyle=":", alpha=0.5) + ax.legend(fontsize=9) + + fig.tight_layout() + fig.savefig(OUT_PATH, dpi=150, bbox_inches="tight") + print(f"Wrote {OUT_PATH}") + + +if __name__ == "__main__": + main() diff --git a/opsqueue/migrations/20260721133844_add_random_order_index_to_submissions.down.sql b/opsqueue/migrations/20260721133844_add_random_order_index_to_submissions.down.sql new file mode 100644 index 0000000..fb04450 --- /dev/null +++ b/opsqueue/migrations/20260721133844_add_random_order_index_to_submissions.down.sql @@ -0,0 +1,2 @@ +DROP INDEX IF EXISTS random_submissions_order; +ALTER TABLE submissions DROP COLUMN IF EXISTS random_order; diff --git a/opsqueue/migrations/20260721133844_add_random_order_index_to_submissions.up.sql b/opsqueue/migrations/20260721133844_add_random_order_index_to_submissions.up.sql new file mode 100644 index 0000000..3797c3f --- /dev/null +++ b/opsqueue/migrations/20260721133844_add_random_order_index_to_submissions.up.sql @@ -0,0 +1,9 @@ +-- Uses the same formula as '20250803174028_better_random_order_formula.down.sql'. +ALTER TABLE submissions ADD COLUMN random_order INTEGER NOT NULL GENERATED ALWAYS AS ( + (((id + (id >> 22)) % 65536) * 40503) % 65536 +) VIRTUAL; + +CREATE INDEX random_submissions_order ON submissions ( + random_order + , id +); diff --git a/opsqueue/src/common/chunk.rs b/opsqueue/src/common/chunk.rs index c27c17e..9010b92 100644 --- a/opsqueue/src/common/chunk.rs +++ b/opsqueue/src/common/chunk.rs @@ -130,6 +130,18 @@ impl From for i64 { } } +impl TryFrom for ChunkIndex { + type Error = crate::common::errors::TryFromIntError; + + fn try_from(value: i64) -> Result { + if value < 0 { + return Err(crate::common::errors::TryFromIntError(())); + } + + Ok(Self(u63::new(value.cast_unsigned()))) + } +} + impl TryFrom for ChunkIndex { type Error = crate::common::errors::TryFromIntError; fn try_from(value: u64) -> Result { diff --git a/opsqueue/src/common/submission.rs b/opsqueue/src/common/submission.rs index 9cc469a..365252c 100644 --- a/opsqueue/src/common/submission.rs +++ b/opsqueue/src/common/submission.rs @@ -579,12 +579,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()) diff --git a/opsqueue/src/consumer/dispatcher/metastate.rs b/opsqueue/src/consumer/dispatcher/metastate.rs index f96534c..7085b83 100644 --- a/opsqueue/src/consumer/dispatcher/metastate.rs +++ b/opsqueue/src/consumer/dispatcher/metastate.rs @@ -1,4 +1,3 @@ -use crossbeam_skiplist::SkipSet; use dashmap::{DashMap, Entry}; use rustc_hash::FxBuildHasher; use tracing; @@ -70,7 +69,6 @@ pub type MetaStateVal = i64; #[derive(Debug, Default)] pub struct MetaStateField { vals_to_counts: DashMap, - counts_to_vals: SkipSet<(usize, MetaStateVal)>, } impl MetaStateField { @@ -80,55 +78,31 @@ 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!() - } - 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); - } - } + if let Entry::Occupied(entry) = self + .vals_to_counts + .entry(val) + .and_modify(|count| *count -= 1) + && *entry.get() == 0 + { + entry.remove(); } } + #[must_use] pub fn is_empty(&self) -> bool { self.vals_to_counts.is_empty() } - pub fn too_high_counts(&self, max: usize) -> impl Iterator + '_ { - tracing::debug!("metastate: {self:?}"); - self.counts_to_vals - .range((max, 0)..) - .map(|entry| entry.value().1) + #[must_use] + pub fn get(&self, val: &MetaStateVal) -> Option { + self.vals_to_counts.get(val).map(|count| *count) } } @@ -158,14 +132,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()); diff --git a/opsqueue/src/consumer/dispatcher/mod.rs b/opsqueue/src/consumer/dispatcher/mod.rs index 8ba9b03..a4853b7 100644 --- a/opsqueue/src/consumer/dispatcher/mod.rs +++ b/opsqueue/src/consumer/dispatcher/mod.rs @@ -3,15 +3,17 @@ pub mod reserver; use crate::{ common::{ - chunk::{Chunk, ChunkId}, - submission::Submission, + chunk::{Chunk, ChunkId, ChunkIndex}, + submission::{Submission, SubmissionId}, }, db::{Connection, Pool, ReaderPool, magic::Bool}, }; use futures::stream::{StreamExt as _, TryStreamExt as _}; +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; @@ -21,6 +23,124 @@ use std::sync::Arc; use super::strategy; use crate::common::StrategicMetadataMap; +pub unsafe extern "C" fn sqlite_reserved_chunk_lookup( + context: *mut ffi::sqlite3_context, + n_args: i32, + args: *mut *mut ffi::sqlite3_value, +) { + unsafe { + 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; + 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) { + unsafe { + if ptr.is_null() { + return; + } + let _boxed: Box> = 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, +) { + unsafe { + 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; + 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) { + unsafe { + if ptr.is_null() { + return; + } + let _boxed: Box> = Box::from_raw(ptr.cast()); + } +} + #[derive(Debug, Clone)] pub struct Dispatcher { reserver: Reserver, @@ -74,30 +194,95 @@ impl Dispatcher { stale_chunks_notifier: &UnboundedSender, ) -> Result, sqlx::Error> { let mut conn = pool.reader_conn().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 - .try_filter_map(|chunk| { - futures::future::ready(Ok(self.reserve_chunk(chunk, stale_chunks_notifier))) - }) + .try_filter_map(|chunk| self.reserve_chunk(chunk, stale_chunks_notifier)) .and_then(|chunk| self.join_chunk_with_submission_info(chunk, pool)) .take(limit) .try_collect() .await } - fn reserve_chunk( + pub 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 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. + let user_data = Box::new(self.reserver.clone()); + let user_data = Box::into_raw(user_data).cast::(); + + let rc = unsafe { + ffi::sqlite3_create_function_v2( + sqlite, + reserved_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}" + ))); + } + + // 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::(); + + 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(()) + } + + async fn reserve_chunk( &self, chunk: Chunk, stale_chunks_notifier: &UnboundedSender, - ) -> Option { + ) -> Result, E> { let chunk_id = ChunkId::from((chunk.submission_id, chunk.chunk_index)); - self.reserver + let val = self + .reserver .try_reserve(chunk_id, chunk_id, stale_chunks_notifier) - .map(|_| chunk) + .map(|_| chunk); + Ok(val) } async fn join_chunk_with_submission_info( @@ -152,3 +337,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(migrator = "crate::MIGRATOR")] + 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_mins(1)); + let (stale_chunks_notifier, mut _stale_chunks_receiver) = unbounded_channel::(); + + 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 + })); + } +} diff --git a/opsqueue/src/consumer/dispatcher/reserver.rs b/opsqueue/src/consumer/dispatcher/reserver.rs index 0699857..91b93e6 100644 --- a/opsqueue/src/consumer/dispatcher/reserver.rs +++ b/opsqueue/src/consumer/dispatcher/reserver.rs @@ -83,6 +83,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. /// diff --git a/opsqueue/src/consumer/strategy.rs b/opsqueue/src/consumer/strategy.rs index 89a7ff8..8156ae2 100644 --- a/opsqueue/src/consumer/strategy.rs +++ b/opsqueue/src/consumer/strategy.rs @@ -1,15 +1,17 @@ +#[cfg(feature = "server-logic")] +use std::string::ToString; + #[cfg(feature = "server-logic")] use futures::stream::BoxStream; + use serde::{Deserialize, Serialize}; + #[cfg(feature = "server-logic")] use sqlx::{QueryBuilder, Sqlite}; #[cfg(feature = "server-logic")] use crate::common::chunk::Chunk; -#[cfg(feature = "server-logic")] -use super::dispatcher::metastate::MetaState; - #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub enum Strategy { Oldest, @@ -21,75 +23,186 @@ pub enum Strategy { }, } +pub struct MetaKeysIter<'a> { + strategy: &'a Strategy, +} + +impl<'a> MetaKeysIter<'a> { + #[must_use] + pub fn take(self) -> &'a Strategy { + self.strategy + } +} + +impl<'a> Iterator for MetaKeysIter<'a> { + type Item = &'a str; + + fn next(&mut self) -> Option { + match self.strategy { + Strategy::Oldest | Strategy::Newest | Strategy::Random => None, + Strategy::PreferDistinct { + meta_key, + underlying, + } => { + self.strategy = underlying.as_ref(); + Some(meta_key.as_str()) + } + } + } +} + +impl Strategy { + #[must_use] + pub fn iter(&self) -> MetaKeysIter<'_> { + MetaKeysIter { strategy: self } + } +} + +impl<'a> IntoIterator for &'a Strategy { + type Item = &'a str; + type IntoIter = MetaKeysIter<'a>; + + fn into_iter(self) -> Self::IntoIter { + self.iter() + } +} + #[cfg(feature = "server-logic")] impl Strategy { pub fn build_query<'a>( &'a self, qb: &'a mut QueryBuilder, - metastate: &MetaState, ) -> &'a mut QueryBuilder { - let qb = self.build_query_snippet(qb, metastate); + let qb = self.build_query_snippet_returning_chunks(qb); tracing::trace!("sql: {:?}", qb.sql()); qb } - fn build_query_snippet<'a>( + fn build_query_snippet_returning_chunks<'a>( &'a self, qb: &'a mut QueryBuilder, - metastate: &MetaState, ) -> &'a mut QueryBuilder { use Strategy::{Newest, Oldest, PreferDistinct, Random}; match self { - Oldest => qb.push("SELECT * FROM chunks ORDER BY submission_id ASC"), - Newest => qb.push("SELECT * FROM chunks ORDER BY submission_id DESC"), - Random => { - let random_offset: u16 = rand::random(); - qb.push("SELECT * FROM chunks WHERE random_order >= ") - .push_bind(random_offset) - .push(" UNION ALL SELECT * FROM chunks WHERE random_order < ") - .push_bind(random_offset) + Oldest => qb + .push("SELECT * FROM chunks") + .push(" WHERE opsqueue_is_reserved(submission_id, chunk_index) = 0") + .push(" ORDER BY submission_id ASC"), + Newest => qb + .push("SELECT * FROM chunks") + .push(" WHERE opsqueue_is_reserved(submission_id, chunk_index) = 0") + .push(" ORDER BY submission_id DESC"), + Random => Self::push_random_order_query( + qb, + "*", + "chunks", + Some("opsqueue_is_reserved(submission_id, chunk_index) = 0"), + ), + PreferDistinct { .. } => { + qb.push("WITH underlying_submissions AS MATERIALIZED ("); + let qb = self.build_query_snippet_returning_submission_ids(qb); + qb.push(") "); + // In SQLite, CROSS JOIN ON/WHERE does NOT produce N + // x M rows, it acts as an INNER JOIN forcing the query + // planner to use '
as the outer loop, preserving its + // sort order. + // c.f. https://sqlite.org/optoverview.html#manual_control_of_query_plans_using_cross_join + qb.push( + " SELECT chunks.* + FROM underlying_submissions + CROSS JOIN chunks + WHERE opsqueue_is_reserved(chunks.submission_id, chunk_index) = 0 + AND chunks.submission_id = underlying_submissions.submission_id", + ) } + } + } - PreferDistinct { - meta_key, - underlying, - } => { - let qb = qb.push(format_args!("WITH inner_{meta_key} AS NOT MATERIALIZED (")); - let qb = underlying.build_query_snippet(qb, metastate); - qb.push(format_args!( - r"), - taken_{meta_key} AS ( - SELECT * FROM submissions_metadata - WHERE - submissions_metadata.metadata_key = ", - )); - qb.push_bind(meta_key); - qb.push( - r" AND submissions_metadata.metadata_value IN (SELECT value FROM json_each(", - ); - match metastate.get(meta_key) { - None => { - tracing::trace!("No metastatefield for key: {meta_key}"); - } - Some(field) => { - let taken_values: Vec<_> = field.too_high_counts(1).collect(); - let taken_values_string = - serde_json::to_string(&taken_values).expect("Always valid JSON"); - tracing::trace!( - "Taken values that are left out of PreferDistinct: {taken_values_string:?}" - ); - qb.push_bind(taken_values_string); + /// Append a query snippet resulting in an ordered "`submission_id`" column. + fn build_query_snippet_returning_submission_ids<'a>( + &'a self, + qb: &'a mut QueryBuilder, + ) -> &'a mut QueryBuilder { + use Strategy::{Newest, Oldest, PreferDistinct, Random}; + match self { + Oldest => qb.push("SELECT id AS submission_id FROM submissions ORDER BY id ASC"), + Newest => qb.push("SELECT id AS submission_id FROM submissions ORDER BY id DESC"), + Random => Self::push_random_order_query(qb, "id as submission_id", "submissions", None), + PreferDistinct { .. } => { + let mut metaiter = self.iter(); + let prefer_distinct_metakeys = metaiter + .by_ref() + .map(ToString::to_string) + .collect::>(); + let underlying = metaiter.take(); + // Unique submission IDs from the underlying strategy. + let qb = qb.push("WITH inner AS NOT MATERIALIZED ("); + let qb = underlying.build_query_snippet_returning_submission_ids(qb); + qb.push(")"); + // Count of in-flight chunks per submission, for each meta key. + // Notice the GROUP BY: this is to avoid duplicate FFI calls for the same values. + for (i, meta_key) in prefer_distinct_metakeys.iter().enumerate() { + qb.push(format_args!( + ", counts_{i} AS ( + SELECT submission_id, ffi_meta_counts.count + FROM submissions_metadata sm + JOIN ( SELECT metadata_value, opsqueue_metadata_count (" + )); + qb.push_bind(meta_key.clone()); + qb.push( + ", metadata_value) AS count + FROM submissions_metadata + WHERE metadata_key = ", + ); + qb.push_bind(meta_key.clone()); + qb.push( + " GROUP BY metadata_value + ) ffi_meta_counts ON sm.metadata_value = ffi_meta_counts.metadata_value + WHERE sm.metadata_key = ", + ); + qb.push_bind(meta_key.clone()); + qb.push(")"); + } + // Submissions ranked by in-flight chunks. + qb.push(" SELECT inner.submission_id FROM inner"); + for i in 0..prefer_distinct_metakeys.len() { + qb.push(format_args!( + " LEFT JOIN counts_{i} ON inner.submission_id = counts_{i}.submission_id", + )); + } + qb.push(" ORDER BY "); + for i in 0..prefer_distinct_metakeys.len() { + if i > 0 { + qb.push(", "); } + qb.push(format_args!("counts_{i}.count ASC NULLS FIRST")); } - qb.push(format_args!(")) - ) - SELECT * FROM inner_{meta_key} WHERE NOT EXISTS (SELECT 1 FROM taken_{meta_key} WHERE inner_{meta_key}.submission_id = taken_{meta_key}.submission_id) - UNION ALL - SELECT * FROM inner_{meta_key} WHERE EXISTS (SELECT 1 FROM taken_{meta_key} WHERE inner_{meta_key}.submission_id = taken_{meta_key}.submission_id) - ")) + qb } } } + + fn push_random_order_query<'a>( + qb: &'a mut QueryBuilder, + columns: &str, + table_name: &str, + condition: Option<&str>, + ) -> &'a mut QueryBuilder { + let random_offset: u16 = rand::random(); + let push_select = |qb: &mut QueryBuilder, operator: &str| { + qb.push(format!( + "SELECT {columns} FROM {table_name} WHERE random_order {operator} " + )) + .push_bind(random_offset); + if let Some(condition_) = condition { + qb.push(format!(" AND {condition_}")); + } + }; + push_select(qb, ">="); + qb.push(" UNION ALL "); + push_select(qb, "<"); + qb + } } #[cfg(feature = "server-logic")] @@ -98,15 +211,76 @@ pub type ChunkStream<'a> = BoxStream<'a, Result>; #[cfg(test)] #[cfg(feature = "server-logic")] pub mod test { - use crate::common::StrategicMetadataMap; - use crate::common::chunk::ChunkSize; - - use super::*; use itertools::Itertools; + use libsqlite3_sys as ffi; use sqlformat::{FormatOptions, QueryParams, format}; use sqlx::Row; use sqlx::{QueryBuilder, Sqlite, SqliteConnection}; + use super::*; + use crate::common::StrategicMetadataMap; + use crate::common::chunk::ChunkSize; + + unsafe extern "C" fn sqlite_reserved_chunk_lookup_noop( + context: *mut ffi::sqlite3_context, + _n_args: i32, + _args: *mut *mut ffi::sqlite3_value, + ) { + unsafe { + ffi::sqlite3_result_int(context, 0); + } + } + + unsafe extern "C" fn sqlite_metadata_count_lookup_noop( + context: *mut ffi::sqlite3_context, + _n_args: i32, + _args: *mut *mut ffi::sqlite3_value, + ) { + unsafe { + ffi::sqlite3_result_null(context); + } + } + + async fn register_lookup_noops(conn: &mut SqliteConnection) { + let mut handle = conn.lock_handle().await.unwrap(); + let sqlite = handle.as_raw_handle().as_ptr(); + let function_name = b"opsqueue_is_reserved\0"; + let rc = unsafe { + ffi::sqlite3_create_function_v2( + sqlite, + function_name.as_ptr().cast(), + 2, + ffi::SQLITE_UTF8, + std::ptr::null_mut(), + Some(sqlite_reserved_chunk_lookup_noop), + None, + None, + None, + ) + }; + assert_eq!(rc, ffi::SQLITE_OK, "register opsqueue_is_reserved failed"); + + let function_name = b"opsqueue_metadata_count\0"; + let rc = unsafe { + ffi::sqlite3_create_function_v2( + sqlite, + function_name.as_ptr().cast(), + 2, + ffi::SQLITE_UTF8, + std::ptr::null_mut(), + Some(sqlite_metadata_count_lookup_noop), + None, + None, + None, + ) + }; + assert_eq!( + rc, + ffi::SQLITE_OK, + "register opsqueue_metadata_count failed" + ); + } + async fn explain(qb: &mut sqlx::QueryBuilder, conn: &mut SqliteConnection) -> String { let formatted_query = format( qb.sql().as_str(), @@ -143,13 +317,31 @@ pub mod test { ); } + /// Weaker invariant than `assert_streaming_query`, used by `PreferDistinct` + /// which deliberately materializes and sorts the set of in-flight + /// submissions to rank them by metadata value. That bounded sort is + /// considered acceptable; what must not happen is a full scan or sort of + /// the potentially very large `chunks`. + fn assert_streaming_chunks(qb: &sqlx::QueryBuilder, explained: &str) { + let query_binding = qb.sql(); + let query = query_binding.as_str(); + assert!( + !explained.contains("SCAN chunks"), + "Query should never scan the whole `chunks` backlog, but it did.\n\nQuery: {query}\n\nPlan: \n\n{explained}" + ); + assert!( + explained.contains("SEARCH chunks"), + "Query should reach `chunks` via an index seek, but it did not.\n\nQuery: {query}\n\nPlan: \n\n{explained}" + ); + } + #[sqlx::test(migrator = "crate::MIGRATOR")] pub async fn test_query_plan_oldest(db: sqlx::SqlitePool) { let mut conn = db.acquire().await.unwrap(); + register_lookup_noops(&mut conn).await; let mut qb = QueryBuilder::new(""); - let metastate = MetaState::default(); - let qb = Strategy::Oldest.build_query(&mut qb, &metastate); + let qb = Strategy::Oldest.build_query(&mut qb); let options = FormatOptions::default(); let formatted_query = format(qb.sql().as_str(), &QueryParams::None, &options); insta::assert_snapshot!(formatted_query, @" @@ -157,6 +349,8 @@ pub mod test { * FROM chunks + WHERE + opsqueue_is_reserved(submission_id, chunk_index) = 0 ORDER BY submission_id ASC "); @@ -169,10 +363,10 @@ pub mod test { #[sqlx::test(migrator = "crate::MIGRATOR")] pub async fn test_query_plan_newest(db: sqlx::SqlitePool) { let mut conn = db.acquire().await.unwrap(); + register_lookup_noops(&mut conn).await; let mut qb = QueryBuilder::new(""); - let metastate = MetaState::default(); - let qb = Strategy::Newest.build_query(&mut qb, &metastate); + let qb = Strategy::Newest.build_query(&mut qb); let options = FormatOptions::default(); let formatted_query = format(qb.sql().as_str(), &QueryParams::None, &options); insta::assert_snapshot!(formatted_query, @" @@ -180,6 +374,8 @@ pub mod test { * FROM chunks + WHERE + opsqueue_is_reserved(submission_id, chunk_index) = 0 ORDER BY submission_id DESC "); @@ -192,10 +388,10 @@ pub mod test { #[sqlx::test(migrator = "crate::MIGRATOR")] pub async fn test_query_plan_random(db: sqlx::SqlitePool) { let mut conn = db.acquire().await.unwrap(); - let metastate = MetaState::default(); + register_lookup_noops(&mut conn).await; let mut qb = QueryBuilder::new(""); - let qb = Strategy::Random.build_query(&mut qb, &metastate); + let qb = Strategy::Random.build_query(&mut qb); let formatted_query = format( qb.sql().as_str(), @@ -209,6 +405,7 @@ pub mod test { chunks WHERE random_order >= ? + AND opsqueue_is_reserved(submission_id, chunk_index) = 0 UNION ALL SELECT * @@ -216,16 +413,17 @@ pub mod test { chunks WHERE random_order < ? + AND opsqueue_is_reserved(submission_id, chunk_index) = 0 "); let explained = explain(qb, &mut conn).await; assert_streaming_query(qb, &explained); - insta::assert_snapshot!(explained, @r" + insta::assert_snapshot!(explained, @" 1, 0, COMPOUND QUERY 2, 1, LEFT-MOST SUBQUERY 5, 2, SEARCH chunks USING INDEX random_chunks_order (random_order>?) - 22, 1, UNION ALL - 25, 22, SEARCH chunks USING INDEX random_chunks_order (random_order= ? - UNION ALL - SELECT - * - FROM - chunks - WHERE - random_order < ? - ), - taken_company_id AS ( - SELECT - * - FROM - submissions_metadata - WHERE - submissions_metadata.metadata_key = ? - AND submissions_metadata.metadata_value IN ( - SELECT - value - FROM - json_each() - ) - ) - SELECT - * - FROM - inner_company_id - WHERE - NOT EXISTS ( + underlying_submissions AS MATERIALIZED ( + WITH + inner AS NOT MATERIALIZED ( SELECT - 1 + id as submission_id FROM - taken_company_id + submissions WHERE - inner_company_id.submission_id = taken_company_id.submission_id - ) - UNION ALL - SELECT - * - FROM - inner_company_id - WHERE - EXISTS ( + random_order >= ? + UNION ALL SELECT - 1 + id as submission_id FROM - taken_company_id + submissions + WHERE + random_order < ? + ), + counts_0 AS ( + SELECT + submission_id, + ffi_meta_counts.count + FROM + submissions_metadata sm + JOIN ( + SELECT + metadata_value, + opsqueue_metadata_count (?, metadata_value) AS count + FROM + submissions_metadata + WHERE + metadata_key = ? + GROUP BY + metadata_value + ) ffi_meta_counts ON sm.metadata_value = ffi_meta_counts.metadata_value WHERE - inner_company_id.submission_id = taken_company_id.submission_id + sm.metadata_key = ? ) + SELECT + inner.submission_id + FROM + inner + LEFT JOIN counts_0 ON inner.submission_id = counts_0.submission_id + ORDER BY + counts_0.count ASC NULLS FIRST + ) + SELECT + chunks.* + FROM + underlying_submissions + CROSS JOIN chunks + WHERE + opsqueue_is_reserved(chunks.submission_id, chunk_index) = 0 + AND chunks.submission_id = underlying_submissions.submission_id "); let explained = explain(qb, &mut conn).await; - assert_streaming_query(qb, &explained); + assert_streaming_chunks(qb, &explained); insta::assert_snapshot!(explained, @" - 1, 0, COMPOUND QUERY - 2, 1, LEFT-MOST SUBQUERY - 3, 2, COMPOUND QUERY - 4, 3, LEFT-MOST SUBQUERY - 7, 4, SEARCH chunks USING INDEX random_chunks_order (random_order>?) - 16, 4, CORRELATED SCALAR SUBQUERY 5 - 20, 16, SEARCH submissions_metadata USING PRIMARY KEY (submission_id=? AND metadata_key=?) - 30, 16, LIST SUBQUERY 3 - 32, 30, SCAN json_each VIRTUAL TABLE INDEX 0: - 58, 3, UNION ALL - 61, 58, SEARCH chunks USING INDEX random_chunks_order (random_order?) - 127, 115, CORRELATED SCALAR SUBQUERY 7 - 131, 127, SEARCH submissions_metadata USING PRIMARY KEY (submission_id=? AND metadata_key=?) - 141, 127, LIST SUBQUERY 3 - 143, 141, SCAN json_each VIRTUAL TABLE INDEX 0: - 169, 114, UNION ALL - 172, 169, SEARCH chunks USING INDEX random_chunks_order (random_order?) + 19, 6, UNION ALL + 22, 19, SEARCH submissions USING INDEX random_submissions_order (random_order= ? - UNION ALL + submissions + ORDER BY + id ASC + ), + counts_0 AS ( SELECT - * + submission_id, + ffi_meta_counts.count FROM - chunks + submissions_metadata sm + JOIN ( + SELECT + metadata_value, + opsqueue_metadata_count (?, metadata_value) AS count + FROM + submissions_metadata + WHERE + metadata_key = ? + GROUP BY + metadata_value + ) ffi_meta_counts ON sm.metadata_value = ffi_meta_counts.metadata_value WHERE - random_order < ? + sm.metadata_key = ? ), - taken_priority AS ( + counts_1 AS ( SELECT - * + submission_id, + ffi_meta_counts.count FROM - submissions_metadata - WHERE - submissions_metadata.metadata_key = ? - AND submissions_metadata.metadata_value IN ( + submissions_metadata sm + JOIN ( SELECT - value + metadata_value, + opsqueue_metadata_count (?, metadata_value) AS count FROM - json_each() - ) + submissions_metadata + WHERE + metadata_key = ? + GROUP BY + metadata_value + ) ffi_meta_counts ON sm.metadata_value = ffi_meta_counts.metadata_value + WHERE + sm.metadata_key = ? ) SELECT - * - FROM - inner_priority - WHERE - NOT EXISTS ( - SELECT - 1 - FROM - taken_priority - WHERE - inner_priority.submission_id = taken_priority.submission_id - ) - UNION ALL - SELECT - * + inner.submission_id FROM - inner_priority - WHERE - EXISTS ( - SELECT - 1 - FROM - taken_priority - WHERE - inner_priority.submission_id = taken_priority.submission_id - ) - ), - taken_company_id AS ( - SELECT - * - FROM - submissions_metadata - WHERE - submissions_metadata.metadata_key = ? - AND submissions_metadata.metadata_value IN ( - SELECT - value - FROM - json_each() - ) + inner + LEFT JOIN counts_0 ON inner.submission_id = counts_0.submission_id + LEFT JOIN counts_1 ON inner.submission_id = counts_1.submission_id + ORDER BY + counts_0.count ASC NULLS FIRST, + counts_1.count ASC NULLS FIRST ) SELECT - * - FROM - inner_company_id - WHERE - NOT EXISTS ( - SELECT - 1 - FROM - taken_company_id - WHERE - inner_company_id.submission_id = taken_company_id.submission_id - ) - UNION ALL - SELECT - * + chunks.* FROM - inner_company_id + underlying_submissions + CROSS JOIN chunks WHERE - EXISTS ( - SELECT - 1 - FROM - taken_company_id - WHERE - inner_company_id.submission_id = taken_company_id.submission_id - ) + opsqueue_is_reserved(chunks.submission_id, chunk_index) = 0 + AND chunks.submission_id = underlying_submissions.submission_id "); let explained = explain(qb, &mut conn).await; - assert_streaming_query(qb, &explained); + assert_streaming_chunks(qb, &explained); insta::assert_snapshot!(explained, @" - 1, 0, COMPOUND QUERY - 2, 1, LEFT-MOST SUBQUERY - 3, 2, COMPOUND QUERY - 4, 3, LEFT-MOST SUBQUERY - 5, 4, COMPOUND QUERY - 6, 5, LEFT-MOST SUBQUERY - 9, 6, SEARCH chunks USING INDEX random_chunks_order (random_order>?) - 18, 6, CORRELATED SCALAR SUBQUERY 5 - 22, 18, SEARCH submissions_metadata USING PRIMARY KEY (submission_id=? AND metadata_key=?) - 32, 18, LIST SUBQUERY 3 - 34, 32, SCAN json_each VIRTUAL TABLE INDEX 0: - 52, 6, CORRELATED SCALAR SUBQUERY 11 - 56, 52, SEARCH submissions_metadata USING PRIMARY KEY (submission_id=? AND metadata_key=?) - 66, 52, LIST SUBQUERY 9 - 68, 66, SCAN json_each VIRTUAL TABLE INDEX 0: - 94, 5, UNION ALL - 97, 94, SEARCH chunks USING INDEX random_chunks_order (random_order?) - 197, 185, CORRELATED SCALAR SUBQUERY 7 - 201, 197, SEARCH submissions_metadata USING PRIMARY KEY (submission_id=? AND metadata_key=?) - 211, 197, LIST SUBQUERY 3 - 213, 211, SCAN json_each VIRTUAL TABLE INDEX 0: - 231, 185, CORRELATED SCALAR SUBQUERY 11 - 235, 231, SEARCH submissions_metadata USING PRIMARY KEY (submission_id=? AND metadata_key=?) - 245, 231, LIST SUBQUERY 9 - 247, 245, SCAN json_each VIRTUAL TABLE INDEX 0: - 273, 184, UNION ALL - 276, 273, SEARCH chunks USING INDEX random_chunks_order (random_order?) - 378, 366, CORRELATED SCALAR SUBQUERY 5 - 382, 378, SEARCH submissions_metadata USING PRIMARY KEY (submission_id=? AND metadata_key=?) - 392, 378, LIST SUBQUERY 3 - 394, 392, SCAN json_each VIRTUAL TABLE INDEX 0: - 412, 366, CORRELATED SCALAR SUBQUERY 13 - 416, 412, SEARCH submissions_metadata USING PRIMARY KEY (submission_id=? AND metadata_key=?) - 426, 412, LIST SUBQUERY 9 - 428, 426, SCAN json_each VIRTUAL TABLE INDEX 0: - 454, 365, UNION ALL - 457, 454, SEARCH chunks USING INDEX random_chunks_order (random_order?) - 557, 545, CORRELATED SCALAR SUBQUERY 7 - 561, 557, SEARCH submissions_metadata USING PRIMARY KEY (submission_id=? AND metadata_key=?) - 571, 557, LIST SUBQUERY 3 - 573, 571, SCAN json_each VIRTUAL TABLE INDEX 0: - 591, 545, CORRELATED SCALAR SUBQUERY 13 - 595, 591, SEARCH submissions_metadata USING PRIMARY KEY (submission_id=? AND metadata_key=?) - 605, 591, LIST SUBQUERY 9 - 607, 605, SCAN json_each VIRTUAL TABLE INDEX 0: - 633, 544, UNION ALL - 636, 633, SEARCH chunks USING INDEX random_chunks_order (random_order = Strategy::Random - .build_query(&mut query_builder, &MetaState::default()) + .build_query(&mut query_builder) .build_query_as() .fetch(conn.get_inner()) .try_collect() @@ -784,7 +867,7 @@ pub mod test { let mut query_builder = QueryBuilder::default(); let vals2: Vec = Strategy::Random - .build_query(&mut query_builder, &MetaState::default()) + .build_query(&mut query_builder) .build_query_as() .fetch(conn.get_inner()) .try_collect()